Fix lint for protected AI routes

Author: MaxGhenis

policyengine_api/routes/simulation_analysis_routes.py

@@ -1,6 +1,9 @@
+import json
+
 from flask import Blueprint, request, Response, stream_with_context
 from werkzeug.exceptions import BadRequest
-from policyengine_api.utils.payload_validators import validate_country
+
+from policyengine_api.security import require_simulation_analysis_api_key
 from policyengine_api.services.simulation_analysis_service import (
     SimulationAnalysisService,
 )
@@ -10,8 +13,6 @@
 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__)
 simulation_analysis_service = SimulationAnalysisService()

policyengine_api/routes/tracer_analysis_routes.py

@@ -1,16 +1,16 @@
+import json
+
 from flask import Blueprint, request, Response, stream_with_context
 from werkzeug.exceptions import BadRequest
+
+from policyengine_api.security import require_simulation_analysis_api_key
 from policyengine_api.utils.payload_validators import (
     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,
 )
-import json
-from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS
-import re
 
 tracer_analysis_bp = Blueprint("tracer_analysis", __name__)
 tracer_analysis_service = TracerAnalysisService()
@@ -30,8 +30,6 @@ def execute_tracer_analysis(country_id):
     household_id = payload.get("household_id")
     policy_id = payload.get("policy_id")
     variable = payload.get("variable")
-    api_version = COUNTRY_PACKAGE_VERSIONS[country_id]
-
     if not isinstance(variable, str):
         raise BadRequest("variable must be a string")
 

policyengine_api/security.py

@@ -6,22 +6,28 @@
 from flask import request
 from werkzeug.exceptions import Unauthorized
 
-_LOCAL_CLIENT_HOSTS = {"127.0.0.1", "::1", "localhost"}
+_ALLOW_UNAUTHENTICATED_AI_ANALYSIS_ENV = (
+    "POLICYENGINE_API_ALLOW_UNAUTHENTICATED_AI_ANALYSIS"
+)
 
 
 def require_simulation_analysis_api_key(view):
-    """Require a shared API key for non-local simulation analysis requests."""
+    """Require a shared API key for simulation analysis requests."""
 
     @wraps(view)
     def wrapped(*args, **kwargs):
-        client_host = request.remote_addr
-        if client_host in _LOCAL_CLIENT_HOSTS:
+        if os.getenv(_ALLOW_UNAUTHENTICATED_AI_ANALYSIS_ENV, "").strip().lower() in {
+            "1",
+            "true",
+            "yes",
+        }:
             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:
+        expected_key = os.getenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "").strip()
+        if not expected_key:
+            raise Unauthorized("Simulation analysis API key is not configured")
+
+        if request.headers.get("X-PolicyEngine-Api-Key") == expected_key:
             return view(*args, **kwargs)
 
         raise Unauthorized("API key required for simulation analysis")

tests/unit/routes/test_ai_route_auth.py

@@ -29,6 +29,19 @@ def test_ai_prompt_rejects_requests_without_api_key(client, monkeypatch):
     assert "API key required" in response.json["message"]
 
 
+def test_ai_prompt_rejects_loopback_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": "127.0.0.1"},
+    )
+
+    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")
 
@@ -65,6 +78,23 @@ def test_tracer_analysis_rejects_requests_without_api_key(client, monkeypatch):
     assert "API key required" in response.json["message"]
 
 
+def test_requests_fail_closed_when_api_key_is_not_configured(client, monkeypatch):
+    monkeypatch.delenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", raising=False)
+
+    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 "not configured" 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")
 
@@ -85,6 +115,4 @@ def test_tracer_analysis_allows_requests_with_api_key(client, monkeypatch):
 
     assert response.status_code == 200
     assert response.json["result"] == "Existing analysis"
-    mock_execute_analysis.assert_called_once_with(
-        "us", 1500, 2, "disposable_income"
-    )
+    mock_execute_analysis.assert_called_once_with("us", 1500, 2, "disposable_income")