Security hardening: body size limit, Gemini key header, AI rate limiting, playbook path traversal

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)

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

api/routes/ai.py

@@ -5,13 +5,16 @@
 
 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
 
 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:

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()

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,
         )

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"

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

tests/test_findings_playbook.py

@@ -0,0 +1,60 @@
+"""Tests for GET /api/findings/<id>/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"] == []

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