PTFE-672: api authentication

runner_manager/main.py

@@ -1,13 +1,39 @@
 import logging
-
+from fastapi import Depends, FastAPI, HTTPException, Response, Security, status
+from fastapi.security import APIKeyHeader, APIKeyQuery
 from fastapi import FastAPI, Response
 
-from runner_manager.dependencies import get_queue
+from runner_manager.dependencies import get_queue, get_settings
 from runner_manager.jobs.startup import startup
 
 log = logging.getLogger(__name__)
 
 app = FastAPI()
+api_key_query = APIKeyQuery(name="api-key", auto_error=False)
+api_key_header = APIKeyHeader(name="x-api-key", auto_error=False)
+
+def get_api_key(
+    api_key_query: str = Security(api_key_query),
+    api_key_header: str = Security(api_key_header),
+) -> str:
+    settings = get_settings()
+    if not settings.api_key:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="API Key not configured in settings",
+        )
+    api_key = api_key_query or api_key_header
+    if not api_key:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="API Key required for this endpoint",
+        )
+    if api_key != settings.api_key.get_secret_value():
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API Key",
+        )
+    return api_key
 
 
 @app.on_event("startup")
@@ -21,3 +47,15 @@ def startup_event():
 @app.get("/_health")
 def health():
     return Response(status_code=200)
+
+
+@app.get("/public")
+def public():
+    """A public endpoint that does not require any authentication."""
+    return "Public Endpoint"
+
+
+@app.get("/private")
+def private(api_key: str = Depends(get_api_key)):
+    """A private endpoint that requires a valid API key to be provided."""
+    return f"Private Endpoint. API Key: {api_key}"

runner_manager/models/settings.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict, Optional
 
 import yaml
-from pydantic import AnyHttpUrl, BaseSettings, RedisDsn
+from pydantic import AnyHttpUrl, BaseSettings, RedisDsn, SecretStr
 
 
 def yaml_config_settings_source(settings: BaseSettings) -> Dict[str, Any]:
@@ -24,6 +24,7 @@ class Settings(BaseSettings):
     name: str = "runner-manager"
     redis_om_url: Optional[RedisDsn] = None
     github_base_url: Optional[AnyHttpUrl] = None
+    api_key: Optional[SecretStr] = None
 
     class Config:
         config: ConfigFile = ConfigFile()

tests/api/auth_tests.py

@@ -0,0 +1,19 @@
+import os
+
+
+def test_public_endpoint(client):
+    response = client.get("/public")
+
+    assert response.status_code == 200
+
+
+def test_private_endpoint_without_api_key(client):
+    response = client.get("/private")
+    assert response.status_code == 401
+
+
+def test_private_endpoint_with_valid_api_key(client):
+    os.environ["API_KEY"] = "secret"
+    headers = {"x-api-key": "secret"}
+    response = client.get("/private", headers=headers)
+    assert response.status_code == 200