Protect AI and tracing routes with API key

Author: MaxGhenis

changelog.d/protect-ai-routes.fixed.md

@@ -0,0 +1 @@
+Protect AI and tracing routes with the shared API key.

policyengine_api/routes/ai_prompt_routes.py

@@ -1,6 +1,7 @@
 from flask import Blueprint, Response, request
 from copy import deepcopy
 from policyengine_api.services.ai_prompt_service import AIPromptService
+from policyengine_api.security import require_simulation_analysis_api_key
 from policyengine_api.utils.payload_validators import validate_country
 from policyengine_api.utils.payload_validators.ai import (
     validate_sim_analysis_payload,
@@ -12,11 +13,12 @@
 ai_prompt_service = AIPromptService()
 
 
-@validate_country
 @ai_prompt_bp.route(
     "/<country_id>/ai-prompts/<string:prompt_name>",
     methods=["POST"],
 )
+@validate_country
+@require_simulation_analysis_api_key
 def generate_ai_prompt(country_id, prompt_name: str) -> Response:
     """
     Get an AI prompt with a given name, filled with the given data.

policyengine_api/routes/simulation_analysis_routes.py

@@ -10,6 +10,7 @@
 from policyengine_api.utils.payload_validators.ai import (
     validate_sim_analysis_payload,
 )
+from policyengine_api.security import require_simulation_analysis_api_key
 import json
 
 simulation_analysis_bp = Blueprint("simulation_analysis", __name__)
@@ -18,6 +19,7 @@
 
 @simulation_analysis_bp.route("/<country_id>/simulation-analysis", methods=["POST"])
 @validate_country
+@require_simulation_analysis_api_key
 def execute_simulation_analysis(country_id):
     print("Got POST request for simulation analysis")
 

policyengine_api/routes/tracer_analysis_routes.py

@@ -4,6 +4,7 @@
     validate_country,
     validate_tracer_analysis_payload,
 )
+from policyengine_api.security import require_simulation_analysis_api_key
 from policyengine_api.services.tracer_analysis_service import (
     TracerAnalysisService,
 )
@@ -17,6 +18,7 @@
 
 @tracer_analysis_bp.route("/<country_id>/tracer-analysis", methods=["POST"])
 @validate_country
+@require_simulation_analysis_api_key
 def execute_tracer_analysis(country_id):
 
     payload = request.json

policyengine_api/security.py

@@ -0,0 +1,29 @@
+"""Security helpers for sensitive API routes."""
+
+import os
+from functools import wraps
+
+from flask import request
+from werkzeug.exceptions import Unauthorized
+
+_LOCAL_CLIENT_HOSTS = {"127.0.0.1", "::1", "localhost"}
+
+
+def require_simulation_analysis_api_key(view):
+    """Require a shared API key for non-local simulation analysis requests."""
+
+    @wraps(view)
+    def wrapped(*args, **kwargs):
+        client_host = request.remote_addr
+        if client_host in _LOCAL_CLIENT_HOSTS:
+            return view(*args, **kwargs)
+
+        expected_key = os.getenv(
+            "POLICYENGINE_API_AI_ANALYSIS_API_KEY", ""
+        ).strip()
+        if expected_key and request.headers.get("X-PolicyEngine-Api-Key") == expected_key:
+            return view(*args, **kwargs)
+
+        raise Unauthorized("API key required for simulation analysis")
+
+    return wrapped

tests/unit/routes/test_ai_route_auth.py

@@ -0,0 +1,90 @@
+import os
+from unittest.mock import patch
+
+import pytest
+
+os.environ.setdefault("FLASK_DEBUG", "1")
+
+from policyengine_api.api import app
+from tests.fixtures.simulation_analysis_prompt_fixtures import valid_input_us
+
+
+@pytest.fixture
+def client():
+    app.config["TESTING"] = True
+    with app.test_client() as test_client:
+        yield test_client
+
+
+def test_ai_prompt_rejects_requests_without_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    response = client.post(
+        "/us/ai-prompts/simulation_analysis",
+        json=valid_input_us,
+        environ_base={"REMOTE_ADDR": "203.0.113.10"},
+    )
+
+    assert response.status_code == 401
+    assert "API key required" in response.json["message"]
+
+
+def test_ai_prompt_allows_requests_with_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    with patch(
+        "policyengine_api.routes.ai_prompt_routes.ai_prompt_service.get_prompt",
+        return_value="Prompt text",
+    ) as mock_get_prompt:
+        response = client.post(
+            "/us/ai-prompts/simulation_analysis",
+            json=valid_input_us,
+            headers={"X-PolicyEngine-Api-Key": "secret-key"},
+            environ_base={"REMOTE_ADDR": "203.0.113.10"},
+        )
+
+    assert response.status_code == 200
+    assert response.json["result"] == "Prompt text"
+    mock_get_prompt.assert_called_once()
+
+
+def test_tracer_analysis_rejects_requests_without_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    response = client.post(
+        "/us/tracer-analysis",
+        json={
+            "household_id": 1500,
+            "policy_id": 2,
+            "variable": "disposable_income",
+        },
+        environ_base={"REMOTE_ADDR": "203.0.113.10"},
+    )
+
+    assert response.status_code == 401
+    assert "API key required" in response.json["message"]
+
+
+def test_tracer_analysis_allows_requests_with_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    with patch(
+        "policyengine_api.routes.tracer_analysis_routes.tracer_analysis_service.execute_analysis",
+        return_value=("Existing analysis", "static"),
+    ) as mock_execute_analysis:
+        response = client.post(
+            "/us/tracer-analysis",
+            json={
+                "household_id": 1500,
+                "policy_id": 2,
+                "variable": "disposable_income",
+            },
+            headers={"X-PolicyEngine-Api-Key": "secret-key"},
+            environ_base={"REMOTE_ADDR": "203.0.113.10"},
+        )
+
+    assert response.status_code == 200
+    assert response.json["result"] == "Existing analysis"
+    mock_execute_analysis.assert_called_once_with(
+        "us", 1500, 2, "disposable_income"
+    )

tests/unit/routes/test_simulation_analysis_auth.py

@@ -0,0 +1,48 @@
+import os
+from unittest.mock import patch
+
+import pytest
+
+os.environ.setdefault("FLASK_DEBUG", "1")
+
+from policyengine_api.api import app
+from tests.to_refactor.fixtures.simulation_analysis_fixtures import test_json
+
+
+@pytest.fixture
+def client():
+    app.config["TESTING"] = True
+    with app.test_client() as test_client:
+        yield test_client
+
+
+def test_simulation_analysis_rejects_requests_without_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    response = client.post(
+        "/us/simulation-analysis",
+        json=test_json,
+        environ_base={"REMOTE_ADDR": "203.0.113.10"},
+    )
+
+    assert response.status_code == 401
+    assert "API key required" in response.json["message"]
+
+
+def test_simulation_analysis_allows_requests_with_api_key(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+
+    with patch(
+        "policyengine_api.routes.simulation_analysis_routes.simulation_analysis_service.execute_analysis",
+        return_value=("Existing analysis", "static"),
+    ) as mock_execute_analysis:
+        response = client.post(
+            "/us/simulation-analysis",
+            json=test_json,
+            headers={"X-PolicyEngine-Api-Key": "secret-key"},
+            environ_base={"REMOTE_ADDR": "203.0.113.10"},
+        )
+
+    assert response.status_code == 200
+    assert response.json["result"] == "Existing analysis"
+    mock_execute_analysis.assert_called_once()