""" Security middleware and utilities for PolicyEye. - API key authentication - Rate limiting per IP - Input sanitization and prompt injection protection - Request size limits - Sensitive data masking in logs """ import time import re import hashlib import logging import secrets import os from collections import defaultdict from fastapi import Request, HTTPException, Depends from fastapi.security import APIKeyHeader from starlette.middleware.base import BaseHTTPMiddleware from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials import jwt from core.config import OPENROUTER_API_KEY, SUPABASE_JWT_SECRET logger = logging.getLogger(__name__) # Initialize a global Supabase client for RS256 token verification _supabase_client = None def get_supabase_client(): global _supabase_client if _supabase_client is None: from supabase import create_client supabase_url = os.environ.get("SUPABASE_URL") supabase_key = os.environ.get("SUPABASE_KEY") if supabase_url and supabase_key: _supabase_client = create_client(supabase_url, supabase_key) return _supabase_client # --- API Key Management --- # Generate a default API key on first run (stored in memory for prototype) # In production, this would come from a secrets manager _VALID_API_KEYS: set[str] = set() _MASTER_API_KEY: str | None = None def generate_api_key() -> str: """Generate a new API key.""" key = f"ss_{secrets.token_urlsafe(32)}" _VALID_API_KEYS.add(key) return key def get_or_create_master_key() -> str: """Get the master API key, creating one if needed.""" global _MASTER_API_KEY if _MASTER_API_KEY is None: _MASTER_API_KEY = generate_api_key() logger.info(f"[Security] Master API key generated: {_MASTER_API_KEY[:12]}...") return _MASTER_API_KEY def validate_api_key(key: str) -> bool: """Validate an API key.""" return key in _VALID_API_KEYS api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) security_bearer = HTTPBearer(auto_error=False) async def require_api_key(api_key: str = Depends(api_key_header)): """Legacy dependency that requires a valid API key (for development/scripts).""" if api_key is None: raise HTTPException(status_code=401, detail="Missing API key. Include X-API-Key header.") if not validate_api_key(api_key): raise HTTPException(status_code=403, detail="Invalid API key.") return api_key async def verify_jwt_token( credentials: HTTPAuthorizationCredentials = Depends(security_bearer), api_key: str = Depends(api_key_header), ): """ Validates either a Supabase JWT (Authorization: Bearer ) OR a valid API Key. Returns a dict with user context if authenticated via JWT, or string if API key. """ async def _verify_supabase_jwt(token: str) -> dict: if not SUPABASE_JWT_SECRET: raise HTTPException(status_code=500, detail="Server not configured for JWT auth (missing SUPABASE_JWT_SECRET).") try: unverified_header = jwt.get_unverified_header(token) alg = unverified_header.get("alg", "HS256") if alg == "HS256": payload = jwt.decode( token, SUPABASE_JWT_SECRET, algorithms=["HS256"], options={"verify_aud": False}, ) return {"type": "jwt", "sub": payload.get("sub"), "email": payload.get("email")} supabase = get_supabase_client() if not supabase: raise HTTPException(status_code=500, detail="Missing Supabase URL/Key to verify RS256 tokens securely.") from fastapi.concurrency import run_in_threadpool import asyncio user_res = None for attempt in range(3): try: user_res = await run_in_threadpool(supabase.auth.get_user, token) break except Exception as e: if attempt == 2: raise e await asyncio.sleep(0.2) if user_res and getattr(user_res, 'user', None): return {"type": "jwt", "sub": user_res.user.id, "email": user_res.user.email} if user_res and getattr(user_res, 'data', None) and getattr(user_res.data, 'user', None): return {"type": "jwt", "sub": user_res.data.user.id, "email": user_res.data.user.email} raise HTTPException(status_code=401, detail="Invalid token.") except jwt.ExpiredSignatureError: raise HTTPException(status_code=401, detail="Token has expired.") except jwt.InvalidTokenError as e: raise HTTPException(status_code=401, detail=f"Invalid token: {str(e)}") except Exception as e: logger.error(f"Supabase Auth verification failed: {e}") raise HTTPException(status_code=401, detail=f"Token verification failed: {str(e)}") if credentials: token = credentials.credentials try: return await _verify_supabase_jwt(token) except HTTPException as jwt_exc: if api_key and validate_api_key(api_key): return {"type": "api_key", "sub": "api_client"} raise jwt_exc if api_key and validate_api_key(api_key): return {"type": "api_key", "sub": "api_client"} raise HTTPException(status_code=401, detail="Missing authentication token.") async def verify_jwt_token_optional( credentials: HTTPAuthorizationCredentials = Depends(security_bearer), api_key: str = Depends(api_key_header), ): """Same as verify_jwt_token, but returns None instead of raising an error if auth fails.""" try: return await verify_jwt_token(credentials, api_key) except HTTPException as e: logger.error(f"[Security] Optional Auth failed: {e.detail}") return None # --- Rate Limiting --- class RateLimiter: """ In-memory sliding window rate limiter per IP address. Limits: 30 requests/minute, 200 requests/hour. """ def __init__(self, default_per_minute: int = 60, default_per_hour: int = 300): self.default_per_minute = default_per_minute self.default_per_hour = default_per_hour # Track by key: f"{ip}:{path}" self._requests: dict[str, list[float]] = defaultdict(list) # Strict endpoint limits self.endpoint_limits = { "/api/check-eligibility": {"min": 5, "hour": 20}, "/api/chat": {"min": 15, "hour": 60}, "/api/upload-policy": {"min": 10, "hour": 30}, } def _cleanup(self, key: str, now: float): """Remove timestamps older than 1 hour.""" cutoff = now - 3600 self._requests[key] = [t for t in self._requests[key] if t > cutoff] def check(self, ip: str, path: str) -> tuple[bool, str]: """Check if the request is allowed. Returns (allowed, reason).""" now = time.time() # We apply rate limits globally per IP, and per IP+path global_key = ip path_key = f"{ip}:{path}" self._cleanup(global_key, now) self._cleanup(path_key, now) # Global limits check global_timestamps = self._requests[global_key] if len([t for t in global_timestamps if t > now - 60]) >= self.default_per_minute: return False, "Global rate limit exceeded." if len(global_timestamps) >= self.default_per_hour: return False, "Global rate limit exceeded." # Path limits check path_limits = self.endpoint_limits.get(path) if path_limits: path_timestamps = self._requests[path_key] if len([t for t in path_timestamps if t > now - 60]) >= path_limits["min"]: return False, f"Rate limit exceeded for {path}: {path_limits['min']} requests/minute allowed." if len(path_timestamps) >= path_limits["hour"]: return False, f"Rate limit exceeded for {path}: {path_limits['hour']} requests/hour allowed." # Allow self._requests[global_key].append(now) if path_limits: self._requests[path_key].append(now) return True, "OK" rate_limiter = RateLimiter() # --- Rate Limiter Middleware --- class RateLimitMiddleware(BaseHTTPMiddleware): """Middleware to enforce rate limits.""" async def dispatch(self, request: Request, call_next): # Skip rate limiting for health check if request.url.path == "/api/health": return await call_next(request) client_ip = request.client.host if request.client else "unknown" allowed, reason = rate_limiter.check(client_ip, request.url.path) if not allowed: logger.warning(f"[RateLimit] {client_ip} — {reason}") raise HTTPException(status_code=429, detail=reason) return await call_next(request) # --- Input Sanitization --- # Patterns that indicate prompt injection attempts _INJECTION_PATTERNS = [ r"ignore\s+(previous|above|all)\s+(instructions|prompts|rules)", r"you\s+are\s+now\s+a", r"forget\s+(everything|all|your)", r"system\s*:\s*", r"<\s*/?script", r"javascript\s*:", r"data\s*:\s*text/html", r"\{\{\s*.*\s*\}\}", # Template injection r"IGNORE_PREVIOUS", r"OVERRIDE_SYSTEM", r"base64[=:]", # Catch base64 encoded attacks r"hex\s*(decode|encode)", r"(exec|eval|os\.system|subprocess)\(", ] _compiled_patterns = [re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS] def sanitize_text_input(text: str, field_name: str = "input") -> str: """ Sanitize text input to prevent prompt injection and XSS. Raises ValueError if injection attempt is detected. """ if not text: return text # Check for prompt injection for pattern in _compiled_patterns: if pattern.search(text): logger.warning(f"[Security] Prompt injection detected in {field_name}: {text[:100]}") raise ValueError(f"Invalid input detected in {field_name}. Please use standard medical/insurance terminology.") # Strip potentially dangerous characters (but allow medical notation) # Keep: letters, numbers, spaces, periods, commas, hyphens, parentheses, slashes, ₹, % cleaned = re.sub(r'[<>{}|\\`\x00-\x08\x0b\x0c\x0e-\x1f]', '', text) return cleaned.strip() def sanitize_case_input(case_data: dict) -> dict: """Sanitize all string fields in case input data.""" sanitized = {} for key, value in case_data.items(): if isinstance(value, str): sanitized[key] = sanitize_text_input(value, key) elif isinstance(value, list): sanitized[key] = [ sanitize_text_input(v, key) if isinstance(v, str) else v for v in value ] else: sanitized[key] = value return sanitized # --- Sensitive Data Masking --- def mask_patient_data(data: dict) -> dict: """Mask sensitive patient data for logging purposes.""" masked = data.copy() sensitive_fields = ["patient_name", "hospital_name", "policy_start_date"] for field in sensitive_fields: if field in masked and masked[field]: value = str(masked[field]) if len(value) > 3: masked[field] = value[:2] + "***" + value[-1:] else: masked[field] = "***" return masked # --- File Upload Validation --- MAX_PDF_SIZE_MB = 20 MAX_PDF_SIZE_BYTES = MAX_PDF_SIZE_MB * 1024 * 1024 ALLOWED_MIME_TYPES = {"application/pdf"} # PDF magic bytes PDF_MAGIC = b"%PDF" def validate_pdf_upload(content: bytes, filename: str) -> None: """ Validate an uploaded PDF file: 1. Size check 2. Magic bytes check (is it actually a PDF?) 3. Filename sanitization """ # Size check if len(content) > MAX_PDF_SIZE_BYTES: raise ValueError( f"File too large: {len(content) / 1024 / 1024:.1f}MB " f"(max {MAX_PDF_SIZE_MB}MB)" ) # Empty file check if len(content) < 100: raise ValueError("File appears to be empty or corrupted") # Magic bytes — verify it's actually a PDF if not content[:4].startswith(PDF_MAGIC): raise ValueError("File is not a valid PDF document") # Filename sanitization safe_filename = re.sub(r'[^a-zA-Z0-9._-]', '_', filename) if not safe_filename.lower().endswith('.pdf'): raise ValueError("Only .pdf files are accepted") logger.info(f"[Security] PDF validated: {safe_filename} ({len(content)/1024:.0f}KB)")