🛡️ Sentinel: Add CORS and TrustedHost middleware for security hardening

src/regression_model_template/controller/kafka_app.py

@@ -12,6 +12,8 @@
 import uvicorn
 import pandas as pd
 from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
 from pydantic import BaseModel
 
 from confluent_kafka import Producer, Consumer, KafkaError, Message
@@ -29,6 +31,8 @@
 DEFAULT_OUTPUT_TOPIC = os.getenv("DEFAULT_OUTPUT_TOPIC", "output_topic")
 DEFAULT_FASTAPI_HOST = os.getenv("DEFAULT_FASTAPI_HOST", "127.0.0.1")
 DEFAULT_FASTAPI_PORT = int(os.getenv("DEFAULT_FASTAPI_PORT", 8100))
+ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",")
+ALLOWED_HOSTS = os.getenv("ALLOWED_HOSTS", "*").split(",")
 LOGGING_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
 
 
@@ -44,6 +48,20 @@
 )
 
 
+# Security Middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=ALLOWED_ORIGINS,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+app.add_middleware(
+    TrustedHostMiddleware,
+    allowed_hosts=ALLOWED_HOSTS,
+)
+
+
 # Data Models
 class PredictionRequest(BaseModel):
     """Request model for prediction."""

tests/controller/test_kafka_app_middleware.py

@@ -0,0 +1,28 @@
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
+from regression_model_template.controller.kafka_app import app
+
+
+def test_middleware_configuration():
+    """Test that security middleware is correctly configured in the app."""
+
+    # Get all middleware classes
+    middleware_classes = [m.cls for m in app.user_middleware]
+
+    # Verify CORSMiddleware is present
+    assert CORSMiddleware in middleware_classes, "CORSMiddleware should be present"
+
+    # Verify TrustedHostMiddleware is present
+    assert TrustedHostMiddleware in middleware_classes, "TrustedHostMiddleware should be present"
+
+    # Inspect CORSMiddleware configuration
+    cors_middleware = next(m for m in app.user_middleware if m.cls == CORSMiddleware)
+    # The 'kwargs' dict contains the kwargs passed to the middleware
+    assert cors_middleware.kwargs["allow_origins"] == ["*"]
+    assert cors_middleware.kwargs["allow_methods"] == ["*"]
+    assert cors_middleware.kwargs["allow_headers"] == ["*"]
+    assert cors_middleware.kwargs["allow_credentials"] is True
+
+    # Inspect TrustedHostMiddleware configuration
+    trusted_host_middleware = next(m for m in app.user_middleware if m.cls == TrustedHostMiddleware)
+    assert trusted_host_middleware.kwargs["allowed_hosts"] == ["*"]