From 2b32252817d029c7b4568dc55444627fbf0a4aa0 Mon Sep 17 00:00:00 2001 From: Tanvir Farhad Date: Sun, 5 Jul 2026 12:56:20 +0100 Subject: [PATCH] fix(security): body size limit, Gemini key in header, AI rate limiting, playbook path traversal guard Closes #149. - SEC-002: cap request bodies at 2MB with a JSON 413 handler - SEC-003: send the Gemini API key via x-goog-api-key header instead of a URL query param, so it never leaks into logs - SEC-004: add a lightweight in-memory rate limiter (20 req/min/IP) on all AI endpoints - SEC-005: validate rule_id and confirm the resolved script path stays inside playbooks/cli/ before reading it - SEC-001 was already fixed on dev via #144; added a regression test to lock it in --- api/app.py | 9 +++++ api/rate_limit.py | 43 ++++++++++++++++++++++ api/routes/ai.py | 8 ++++ api/routes/findings.py | 24 ++++++++---- api/services/ai_provider.py | 2 +- tests/test_ai_provider.py | 24 ++++++++++++ tests/test_app_security.py | 33 +++++++++++++++++ tests/test_findings_playbook.py | 60 ++++++++++++++++++++++++++++++ tests/test_rate_limit.py | 65 +++++++++++++++++++++++++++++++++ 9 files changed, 260 insertions(+), 8 deletions(-) create mode 100644 api/rate_limit.py create mode 100644 tests/test_ai_provider.py create mode 100644 tests/test_app_security.py create mode 100644 tests/test_findings_playbook.py create mode 100644 tests/test_rate_limit.py diff --git a/api/app.py b/api/app.py index 9fda9cd..3eef72d 100644 --- a/api/app.py +++ b/api/app.py @@ -21,6 +21,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))\"" @@ -96,6 +100,7 @@ def create_app() -> Flask: # Configuration & Security # # ------------------------------------------------------------------ # app.config["JWT_SECRET"] = _resolve_jwt_secret() + app.config["MAX_CONTENT_LENGTH"] = _MAX_CONTENT_LENGTH # ------------------------------------------------------------------ # # CORS # @@ -242,6 +247,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/rate_limit.py b/api/rate_limit.py new file mode 100644 index 0000000..cd56288 --- /dev/null +++ b/api/rate_limit.py @@ -0,0 +1,43 @@ +"""Minimal in-memory sliding-window rate limiter for Flask routes. + +Not a substitute for a distributed limiter (e.g. Redis-backed) behind a +load balancer with multiple worker processes, but sufficient to blunt basic +abuse of endpoints that proxy to metered third-party APIs. +""" + +import threading +import time +from functools import wraps + +from flask import current_app, jsonify, request + +_lock = threading.Lock() +_hits: dict[str, list[float]] = {} + + +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. + """ + + def decorator(fn): + @wraps(fn) + def wrapped(*args, **kwargs): + if current_app.testing: + return fn(*args, **kwargs) + + key = f"{request.remote_addr}:{request.path}" + now = time.monotonic() + with _lock: + recent = [t for t in _hits.get(key, ()) if now - t < window_seconds] + if len(recent) >= max_requests: + return jsonify({"error": "Rate limit exceeded, try again later"}), 429 + recent.append(now) + _hits[key] = recent + + return fn(*args, **kwargs) + + return wrapped + + return decorator diff --git a/api/routes/ai.py b/api/routes/ai.py index 8ebf196..3e8259a 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, @@ -159,6 +162,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: @@ -209,6 +213,7 @@ def insights(): @ai_bp.post("/api/ai/summary") +@rate_limit(_AI_RATE_LIMIT) def ai_summary(): body, error = _read_request() if error: @@ -247,6 +252,7 @@ def ai_summary(): @ai_bp.post("/api/ai/prioritise") +@rate_limit(_AI_RATE_LIMIT) def ai_prioritise(): body, error = _read_request() if error: @@ -292,6 +298,7 @@ def ai_prioritise(): @ai_bp.post("/api/ai/ask") +@rate_limit(_AI_RATE_LIMIT) def ai_ask(): body, error = _read_request() if error: @@ -333,6 +340,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 054891b..6af9199 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__) @@ -77,12 +81,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 70d2380..278eb79 100644 --- a/api/services/ai_provider.py +++ b/api/services/ai_provider.py @@ -95,7 +95,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..b7bde7c --- /dev/null +++ b/tests/test_ai_provider.py @@ -0,0 +1,24 @@ +"""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..571db57 --- /dev/null +++ b/tests/test_app_security.py @@ -0,0 +1,33 @@ +"""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 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..7cceeb7 --- /dev/null +++ b/tests/test_rate_limit.py @@ -0,0 +1,65 @@ +"""Tests for api.rate_limit.rate_limit (SEC-004).""" + +import pytest +from flask import Flask, jsonify + +from api.rate_limit import _hits, rate_limit + + +@pytest.fixture(autouse=True) +def _reset_rate_limit_state(): + _hits.clear() + yield + _hits.clear() + + +def _make_app(limit=3, window=60): + app = Flask(__name__) + + @app.get("/ping") + @rate_limit(limit, window) + def ping(): + return jsonify({"ok": True}) + + return app + + +def test_requests_under_limit_succeed(): + 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(): + 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(): + app = Flask(__name__) + + @app.get("/a") + @rate_limit(1) + def a(): + return jsonify({"ok": True}) + + @app.get("/b") + @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