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,17 @@
 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."""
+    """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:
-            return view(*args, **kwargs)
+        expected_key = os.getenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "").strip()
+        if not expected_key:
+            raise Unauthorized("Simulation analysis API key is not configured")
 
-        expected_key = os.getenv(
-            "POLICYENGINE_API_AI_ANALYSIS_API_KEY", ""
-        ).strip()
-        if expected_key and request.headers.get("X-PolicyEngine-Api-Key") == expected_key:
+        if request.headers.get("X-PolicyEngine-Api-Key") == expected_key:
             return view(*args, **kwargs)
 
         raise Unauthorized("API key required for simulation analysis")

tests/conftest.py

@@ -3,7 +3,7 @@
 import time
 from contextlib import contextmanager
 from subprocess import Popen, TimeoutExpired
-import sys
+import os
 import redis
 import pytest
 from policyengine_api.api import app
@@ -33,9 +33,18 @@ def running(process_arguments, seconds_to_wait_after_launch=0):
 def client():
     """run the app for the tests to run against"""
     app.config["TESTING"] = True
+    previous_api_key = os.environ.get("POLICYENGINE_API_AI_ANALYSIS_API_KEY")
+    os.environ["POLICYENGINE_API_AI_ANALYSIS_API_KEY"] = "test-ai-analysis-key"
     with running(["redis-server"], 3):
         redis_client = redis.Redis()
         redis_client.ping()
         with running([sys.executable, "policyengine_api/worker.py"], 3):
             with app.test_client() as test_client:
+                test_client.environ_base["HTTP_X_POLICYENGINE_API_KEY"] = (
+                    "test-ai-analysis-key"
+                )
                 yield test_client
+    if previous_api_key is None:
+        os.environ.pop("POLICYENGINE_API_AI_ANALYSIS_API_KEY", None)
+    else:
+        os.environ["POLICYENGINE_API_AI_ANALYSIS_API_KEY"] = previous_api_key

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,41 @@ 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_env_flag_does_not_reopen_tracer_analysis(client, monkeypatch):
+    monkeypatch.setenv("POLICYENGINE_API_AI_ANALYSIS_API_KEY", "secret-key")
+    monkeypatch.setenv("POLICYENGINE_API_ALLOW_UNAUTHENTICATED_AI_ANALYSIS", "true")
+
+    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")
 
@@ -85,6 +133,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")