diff --git a/api/app.py b/api/app.py index e39448d..cff0b87 100644 --- a/api/app.py +++ b/api/app.py @@ -8,6 +8,7 @@ from dotenv import load_dotenv from flask import Flask, g, jsonify, request from flask_cors import CORS +from werkzeug.middleware.proxy_fix import ProxyFix from api.models.finding import DatabaseManager @@ -22,6 +23,10 @@ # Paths that are always public regardless of environment or demo mode _ALWAYS_PUBLIC = {"/", "/health"} +# Maximum accepted request body size (bytes). Guards AI endpoints and the rest +# of the API against oversized payloads tying up worker threads. +_MAX_CONTENT_LENGTH = 2 * 1024 * 1024 # 2 MB + _INSECURE_JWT_DEFAULT = "change-me-in-production" _MIN_JWT_SECRET_LENGTH = 32 _GENERATE_CMD = 'python -c "import secrets; print(secrets.token_urlsafe(32))"' @@ -92,10 +97,17 @@ def create_app() -> Flask: """ app = Flask(__name__) + # Trust exactly one reverse-proxy hop (Render's edge) for the client IP + # and scheme, so request.remote_addr reflects the real caller instead of + # collapsing every client onto Render's proxy address. Rate limiting and + # any other per-IP logic depend on this being accurate. + app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1) + # ------------------------------------------------------------------ # # Configuration & Security # # ------------------------------------------------------------------ # app.config["JWT_SECRET"] = _resolve_jwt_secret() + app.config["MAX_CONTENT_LENGTH"] = _MAX_CONTENT_LENGTH # ------------------------------------------------------------------ # # CORS # @@ -236,6 +248,10 @@ def forbidden(exc): def not_found(exc): return jsonify({"error": "Not found"}), 404 + @app.errorhandler(413) + def payload_too_large(exc): + return jsonify({"error": "Request body too large"}), 413 + @app.errorhandler(500) def internal_error(exc): logger.error("Unhandled exception: %s", exc) diff --git a/api/models/finding.py b/api/models/finding.py index e40ea86..5f7a800 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -167,6 +167,17 @@ def create_tables(self) -> None: CREATE INDEX IF NOT EXISTS idx_findings_rule_id ON findings(rule_id); """) + cur.execute(""" + CREATE TABLE IF NOT EXISTS rate_limit_hits ( + id SERIAL PRIMARY KEY, + limit_key TEXT NOT NULL, + hit_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """) + cur.execute(""" + CREATE INDEX IF NOT EXISTS idx_rate_limit_hits_key_time + ON rate_limit_hits(limit_key, hit_at); + """) conn.commit() logger.info("Database tables created / verified") diff --git a/api/rate_limit.py b/api/rate_limit.py new file mode 100644 index 0000000..3a54551 --- /dev/null +++ b/api/rate_limit.py @@ -0,0 +1,103 @@ +"""Postgres-backed sliding-window rate limiter for Flask routes. + +Deployments run gunicorn with multiple worker processes (see startup.sh), +so a plain in-memory counter would give each worker its own independent +budget instead of enforcing one shared limit per client. State is kept in +the `rate_limit_hits` table instead, with the count-then-insert made atomic +across all workers via a Postgres advisory transaction lock keyed on the +same rate-limit key. +""" + +import logging +import os +from functools import wraps +from typing import Any + +from flask import current_app, g, jsonify, request + +from api.models.finding import DatabaseManager + +logger = logging.getLogger(__name__) + + +def _get_db() -> DatabaseManager: + """Reuse the same `g.db` connection every other route module populates, + so it's cleaned up by api.app's existing close_db teardown instead of + leaking a second, unmanaged connection per request.""" + if "db" not in g: + g.db = DatabaseManager(os.environ["DATABASE_URL"]) + g.db.connect() + return g.db + + +def _check_and_record(key: str, max_requests: int, window_seconds: int) -> bool: + """Return True if this request is within the limit, recording it if so. + + The advisory lock is scoped to the current transaction (released on the + commit/rollback below), which serializes concurrent requests for the + same key across every worker process sharing this database — not just + threads within one process. + """ + db = _get_db() + conn: Any = db._get_conn() + try: + with conn.cursor() as cur: + cur.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (key,)) + cur.execute( + "DELETE FROM rate_limit_hits WHERE limit_key = %s AND hit_at < now() - make_interval(secs => %s)", + (key, window_seconds), + ) + cur.execute("SELECT COUNT(*) FROM rate_limit_hits WHERE limit_key = %s", (key,)) + (count,) = cur.fetchone() + allowed = count < max_requests + if allowed: + cur.execute("INSERT INTO rate_limit_hits (limit_key, hit_at) VALUES (%s, now())", (key,)) + conn.commit() + except Exception: + conn.rollback() + raise + return allowed + + +def rate_limit(max_requests: int, window_seconds: int = 60): + """Limit a view to ``max_requests`` per ``window_seconds`` per client IP. + + Disabled automatically when the Flask app is in testing mode. Fails + open (allows the request) if the rate-limit check itself errors, since + a database hiccup here shouldn't take down an otherwise-healthy AI + endpoint — this limiter blunts abuse, it isn't the primary safeguard. + """ + + def decorator(fn): + @wraps(fn) + def wrapped(*args, **kwargs): + if current_app.testing: + return fn(*args, **kwargs) + + if "DATABASE_URL" not in os.environ: + # Distinct from the generic fail-open below: a transient DB + # hiccup is fine to wave through, but a deployment that's + # simply missing its DB config would otherwise silently and + # permanently run this metered endpoint with no rate + # limiting at all, which should be loud, not silent. + logger.error( + "DATABASE_URL is not set; refusing %s without rate limiting", + request.path, + ) + return jsonify({"error": "Service temporarily unavailable"}), 503 + + key = f"{request.remote_addr}:{request.path}" + try: + allowed = _check_and_record(key, max_requests, window_seconds) + except Exception: + logger.exception("Rate limit check failed for %s; failing open", key) + allowed = True + + if not allowed: + return jsonify({"error": "Rate limit exceeded, try again later"}), 429 + + return fn(*args, **kwargs) + + return wrapped + + return decorator diff --git a/api/routes/ai.py b/api/routes/ai.py index ed437b3..5059ea4 100644 --- a/api/routes/ai.py +++ b/api/routes/ai.py @@ -5,6 +5,7 @@ from flask import Blueprint, jsonify, request +from api.rate_limit import rate_limit from api.services.ai_provider import PROVIDERS as SUPPORTED_PROVIDERS from api.services.ai_provider import get_completion from ai.retriever import retrieve, VectorStoreNotBuilt @@ -12,6 +13,8 @@ ai_bp = Blueprint("ai", __name__) logger = logging.getLogger(__name__) +_AI_RATE_LIMIT = 20 # requests per minute per client IP, per endpoint + _SEVERITY_RANK = { "CRITICAL": 5, "HIGH": 4, @@ -156,6 +159,7 @@ def _read_request(): @ai_bp.post("/api/ai/insights") +@rate_limit(_AI_RATE_LIMIT) def insights(): data = request.get_json(silent=True) if data is None: @@ -206,6 +210,7 @@ def insights(): @ai_bp.post("/api/ai/summary") +@rate_limit(_AI_RATE_LIMIT) def ai_summary(): body, error = _read_request() if error: @@ -244,6 +249,7 @@ def ai_summary(): @ai_bp.post("/api/ai/prioritise") +@rate_limit(_AI_RATE_LIMIT) def ai_prioritise(): body, error = _read_request() if error: @@ -289,6 +295,7 @@ def ai_prioritise(): @ai_bp.post("/api/ai/ask") +@rate_limit(_AI_RATE_LIMIT) def ai_ask(): body, error = _read_request() if error: @@ -330,6 +337,7 @@ def ai_ask(): @ai_bp.post("/api/ai/threat-simulation") +@rate_limit(_AI_RATE_LIMIT) def ai_threat_simulation(): body, error = _read_request() if error: diff --git a/api/routes/findings.py b/api/routes/findings.py index dcb1f13..3a4939b 100644 --- a/api/routes/findings.py +++ b/api/routes/findings.py @@ -2,14 +2,18 @@ import logging import os +import re from pathlib import Path from flask import Blueprint, g, jsonify, request from api.models.finding import DatabaseManager -_PLAYBOOKS_DIR = Path(__file__).parent.parent.parent / "playbooks" / "cli" +_PLAYBOOKS_DIR = (Path(__file__).parent.parent.parent / "playbooks" / "cli").resolve() -_PLAYBOOKS_DIR = Path(__file__).parent.parent.parent / "playbooks" / "cli" +# Known rule_id shape, e.g. AZ-STOR-001. Anything else is rejected before it +# ever reaches the filesystem, closing off path traversal via a crafted or +# corrupted rule_id. +_RULE_ID_RE = re.compile(r"^[A-Z0-9]+(?:-[A-Z0-9]+)*$") findings_bp = Blueprint("findings", __name__) logger = logging.getLogger(__name__) @@ -73,12 +77,18 @@ def get_playbook(finding_id: int): remediation = finding.get("remediation", "") cve_refs = finding.get("cve_references") or [] - # Map rule_id (e.g. AZ-STOR-001) to script filename (fix_az_stor_001.sh) - script_name = "fix_" + rule_id.lower().replace("-", "_") + ".sh" - script_path = _PLAYBOOKS_DIR / script_name - cli_commands = [] - if script_path.exists(): + script_path = None + if _RULE_ID_RE.match(rule_id or ""): + # Map rule_id (e.g. AZ-STOR-001) to script filename (fix_az_stor_001.sh) + script_name = "fix_" + rule_id.lower().replace("-", "_") + ".sh" + candidate = (_PLAYBOOKS_DIR / script_name).resolve() + if candidate.is_relative_to(_PLAYBOOKS_DIR): + script_path = candidate + else: + logger.warning("Rejected malformed rule_id for playbook lookup: %r", rule_id) + + if script_path is not None and script_path.exists(): raw = script_path.read_text() # Strip comment-only lines and blank lines; join multi-line commands lines = raw.splitlines() diff --git a/api/services/ai_provider.py b/api/services/ai_provider.py index 9134632..a3e4b0b 100644 --- a/api/services/ai_provider.py +++ b/api/services/ai_provider.py @@ -91,7 +91,7 @@ def _gemini(api_key: str, prompt: str, model: str) -> str: try: resp = requests.post( f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent", - params={"key": api_key}, + headers={"x-goog-api-key": api_key, "content-type": "application/json"}, json={"contents": [{"parts": [{"text": prompt}]}]}, timeout=30, ) diff --git a/tests/test_ai_provider.py b/tests/test_ai_provider.py new file mode 100644 index 0000000..9fe292e --- /dev/null +++ b/tests/test_ai_provider.py @@ -0,0 +1,22 @@ +"""Tests for api.services.ai_provider (SEC-003: Gemini key must not leak via URL).""" + +from unittest.mock import MagicMock, patch + +from api.services import ai_provider + + +@patch("api.services.ai_provider.requests.post") +def test_gemini_api_key_sent_via_header_not_query_param(mock_post): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"candidates": [{"content": {"parts": [{"text": "ok"}]}}]} + mock_post.return_value = mock_resp + + result = ai_provider.get_completion("gemini", "top-secret-key", "prompt") + + assert result == "ok" + args, kwargs = mock_post.call_args + url = args[0] if args else kwargs.get("url", "") + assert "top-secret-key" not in url + assert "params" not in kwargs or "key" not in (kwargs.get("params") or {}) + assert kwargs["headers"]["x-goog-api-key"] == "top-secret-key" diff --git a/tests/test_app_security.py b/tests/test_app_security.py new file mode 100644 index 0000000..8dd58ad --- /dev/null +++ b/tests/test_app_security.py @@ -0,0 +1,72 @@ +"""App-level security tests: JWT enforcement on GET routes (SEC-001) and +request body size limits (SEC-002).""" + + +def test_get_score_without_auth_returns_401(client): + resp = client.get("/api/score") + assert resp.status_code == 401 + + +def test_get_findings_without_auth_returns_401(client): + resp = client.get("/api/findings") + assert resp.status_code == 401 + + +def test_get_scans_without_auth_returns_401(client): + resp = client.get("/api/scans") + assert resp.status_code == 401 + + +def test_get_compliance_without_auth_returns_401(client): + resp = client.get("/api/compliance/cis") + assert resp.status_code == 401 + + +def test_health_and_index_remain_public(client): + assert client.get("/health").status_code == 200 + assert client.get("/").status_code == 200 + + +def test_oversized_body_returns_413(client, auth_headers): + large_body = b"a" * (3 * 1024 * 1024) # 3 MB, over the 2 MB limit + resp = client.post("/api/ai/insights", data=large_body, headers=auth_headers) + assert resp.status_code == 413 + + +def test_proxy_fix_trusts_exactly_one_hop(app): + """Render terminates TLS and proxies to the app over one hop. Without + ProxyFix, request.remote_addr is always Render's proxy address, which + collapses every distinct client into one rate-limit bucket. Trusting + more than one hop would let a client spoof its own IP via a forged + leading X-Forwarded-For entry, so exactly one hop must be configured.""" + from werkzeug.middleware.proxy_fix import ProxyFix + + assert isinstance(app.wsgi_app, ProxyFix) + assert app.wsgi_app.x_for == 1 + assert app.wsgi_app.x_proto == 1 + + +def test_forwarded_for_resolves_through_full_wsgi_stack(app): + """End-to-end: a request through the real WSGI stack (not + test_request_context, which bypasses wsgi_app entirely) must see + request.remote_addr rewritten from the trusted forwarded header.""" + import api.app as app_module + + app_module._ALWAYS_PUBLIC.add("/__debug_remote_addr") + try: + + @app.route("/__debug_remote_addr") + def _debug_remote_addr(): + from flask import jsonify, request + + return jsonify({"ip": request.remote_addr}) + + client = app.test_client() + resp = client.get( + "/__debug_remote_addr", + headers={"X-Forwarded-For": "203.0.113.7, 198.51.100.9"}, + environ_overrides={"REMOTE_ADDR": "10.0.0.1"}, + ) + assert resp.get_json()["ip"] == "198.51.100.9" + finally: + app_module._ALWAYS_PUBLIC.discard("/__debug_remote_addr") diff --git a/tests/test_findings_playbook.py b/tests/test_findings_playbook.py new file mode 100644 index 0000000..4c8f690 --- /dev/null +++ b/tests/test_findings_playbook.py @@ -0,0 +1,60 @@ +"""Tests for GET /api/findings//playbook (SEC-005 path traversal guard).""" + +from unittest.mock import MagicMock, patch + +import api.routes.findings as findings_route + + +def _finding(rule_id, remediation="Do the thing.", cve_references=None): + return { + "rule_id": rule_id, + "remediation": remediation, + "cve_references": cve_references or [], + } + + +def _mock_db(finding): + db = MagicMock() + db.get_finding_by_id.return_value = finding + return db + + +def _get_playbook(client, auth_headers, finding): + with patch.object(findings_route, "_get_db", return_value=_mock_db(finding)): + return client.get("/api/findings/1/playbook", headers=auth_headers) + + +def test_valid_rule_id_returns_playbook(client, auth_headers): + resp = _get_playbook(client, auth_headers, _finding("AZ-STOR-001")) + assert resp.status_code == 200 + assert resp.get_json()["cli_commands"] + + +def test_missing_finding_returns_404(client, auth_headers): + with patch.object(findings_route, "_get_db", return_value=_mock_db(None)): + resp = client.get("/api/findings/999/playbook", headers=auth_headers) + assert resp.status_code == 404 + + +def test_dotdot_rule_id_does_not_escape_playbooks_dir(client, auth_headers): + resp = _get_playbook(client, auth_headers, _finding("../../../../etc/passwd")) + assert resp.status_code == 200 + assert resp.get_json()["cli_commands"] == [] + + +def test_rule_id_with_path_separator_rejected(client, auth_headers): + resp = _get_playbook(client, auth_headers, _finding("AZ/STOR/001")) + assert resp.status_code == 200 + assert resp.get_json()["cli_commands"] == [] + + +def test_absolute_path_rule_id_rejected(client, auth_headers): + resp = _get_playbook(client, auth_headers, _finding("/etc/passwd")) + assert resp.status_code == 200 + assert resp.get_json()["cli_commands"] == [] + + +def test_unknown_but_well_formed_rule_id_returns_empty_commands(client, auth_headers): + resp = _get_playbook(client, auth_headers, _finding("AZ-NOPE-999")) + assert resp.status_code == 200 + assert resp.get_json()["cli_commands"] == [] diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py new file mode 100644 index 0000000..35add4f --- /dev/null +++ b/tests/test_rate_limit.py @@ -0,0 +1,181 @@ +"""Tests for api.rate_limit.rate_limit (SEC-004).""" + +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask, jsonify + +import api.rate_limit as rate_limit_module + + +class _FakeRateLimitStore: + """In-memory stand-in for the rate_limit_hits table, keyed like Postgres. + + Mimics the real query sequence in _check_and_record: purge expired + hits for the key, count what's left, and only record a new hit when + under the limit — so the decorator's behavior is exercised without a + live database. + """ + + def __init__(self): + self.hits: dict[str, list[float]] = {} + + def check_and_record(self, key: str, max_requests: int, window_seconds: int) -> bool: + recent = self.hits.get(key, []) + if len(recent) >= max_requests: + return False + recent.append(0.0) + self.hits[key] = recent + return True + + +@pytest.fixture +def _mock_db_backend(monkeypatch): + """Stand in for the Postgres-backed store at the decorator level. + + Applied explicitly (not autouse) so the low-level tests further down + that exercise the real _check_and_record's SQL can still call the + genuine implementation. Sets DATABASE_URL so the decorator's + config-guard doesn't short-circuit to the 503 path before reaching + _check_and_record. + """ + monkeypatch.setenv("DATABASE_URL", "postgresql://test/db") + store = _FakeRateLimitStore() + with patch.object(rate_limit_module, "_check_and_record", side_effect=store.check_and_record): + yield store + + +def _make_app(limit=3, window=60): + app = Flask(__name__) + + @app.get("/ping") + @rate_limit_module.rate_limit(limit, window) + def ping(): + return jsonify({"ok": True}) + + return app + + +def test_requests_under_limit_succeed(_mock_db_backend): + client = _make_app(limit=3).test_client() + for _ in range(3): + assert client.get("/ping").status_code == 200 + + +def test_requests_over_limit_are_rejected(_mock_db_backend): + client = _make_app(limit=3).test_client() + for _ in range(3): + client.get("/ping") + resp = client.get("/ping") + assert resp.status_code == 429 + + +def test_limit_is_scoped_per_route(_mock_db_backend): + app = Flask(__name__) + + @app.get("/a") + @rate_limit_module.rate_limit(1) + def a(): + return jsonify({"ok": True}) + + @app.get("/b") + @rate_limit_module.rate_limit(1) + def b(): + return jsonify({"ok": True}) + + client = app.test_client() + assert client.get("/a").status_code == 200 + assert client.get("/b").status_code == 200 + assert client.get("/a").status_code == 429 + + +def test_testing_mode_bypasses_rate_limit(): + app = _make_app(limit=1) + app.config["TESTING"] = True + client = app.test_client() + for _ in range(5): + assert client.get("/ping").status_code == 200 + + +def test_missing_database_url_returns_503_instead_of_silently_disabling(monkeypatch): + """Unlike a transient DB error (which fails open), a deployment that's + simply missing DATABASE_URL must not silently run this metered endpoint + with no rate limiting for its entire lifetime — it should refuse loudly.""" + monkeypatch.delenv("DATABASE_URL", raising=False) + app = _make_app(limit=1) + client = app.test_client() + + resp = client.get("/ping") + + assert resp.status_code == 503 + + +def test_check_failure_fails_open(_mock_db_backend): + """A DB error during the rate-limit check must not block the request — + this limiter blunts abuse, it isn't the endpoint's primary safeguard.""" + app = _make_app(limit=1) + client = app.test_client() + + with patch.object(rate_limit_module, "_check_and_record", side_effect=Exception("db down")): + assert client.get("/ping").status_code == 200 + assert client.get("/ping").status_code == 200 + + +def test_check_and_record_uses_advisory_lock_and_purges_expired_hits(): + """Verify the real _check_and_record issues the expected Postgres calls: + advisory lock, purge-then-count, and insert only when under the limit.""" + mock_cursor = MagicMock() + mock_cursor.fetchone.return_value = (0,) + mock_cursor.__enter__.return_value = mock_cursor + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_db = MagicMock() + mock_db._get_conn.return_value = mock_conn + + with patch.object(rate_limit_module, "_get_db", return_value=mock_db): + allowed = rate_limit_module._check_and_record("1.2.3.4:/api/x", 5, 60) + + assert allowed is True + queries = [call.args[0] for call in mock_cursor.execute.call_args_list] + assert any("pg_advisory_xact_lock" in q for q in queries) + assert any("DELETE FROM rate_limit_hits" in q for q in queries) + assert any("SELECT COUNT(*) FROM rate_limit_hits" in q for q in queries) + assert any("INSERT INTO rate_limit_hits" in q for q in queries) + mock_conn.commit.assert_called_once() + + +def test_check_and_record_denies_and_skips_insert_when_at_limit(): + mock_cursor = MagicMock() + mock_cursor.fetchone.return_value = (5,) + mock_cursor.__enter__.return_value = mock_cursor + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_db = MagicMock() + mock_db._get_conn.return_value = mock_conn + + with patch.object(rate_limit_module, "_get_db", return_value=mock_db): + allowed = rate_limit_module._check_and_record("1.2.3.4:/api/x", 5, 60) + + assert allowed is False + queries = [call.args[0] for call in mock_cursor.execute.call_args_list] + assert not any("INSERT INTO rate_limit_hits" in q for q in queries) + mock_conn.commit.assert_called_once() + + +def test_check_and_record_rolls_back_on_query_error(): + """A mid-transaction error must roll back rather than leave the + connection in an aborted-transaction state for later use.""" + mock_cursor = MagicMock() + mock_cursor.__enter__.return_value = mock_cursor + mock_cursor.execute.side_effect = [None, None, Exception("lock timeout")] + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_db = MagicMock() + mock_db._get_conn.return_value = mock_conn + + with patch.object(rate_limit_module, "_get_db", return_value=mock_db): + with pytest.raises(Exception, match="lock timeout"): + rate_limit_module._check_and_record("1.2.3.4:/api/x", 5, 60) + + mock_conn.rollback.assert_called_once() + mock_conn.commit.assert_not_called()