Spaces:
Runtime error
Runtime error
| """ | |
| Security middleware for production environment | |
| """ | |
| import time | |
| import hashlib | |
| from collections import defaultdict | |
| from datetime import datetime, timedelta | |
| from typing import Dict, Optional | |
| from fastapi import Request, HTTPException, status | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.middleware.trustedhost import TrustedHostMiddleware | |
| from starlette.middleware.base import BaseHTTPMiddleware | |
| from config.security import SecurityConfig | |
| class RateLimitMiddleware(BaseHTTPMiddleware): | |
| """Rate limiting middleware to prevent brute force attacks""" | |
| def __init__(self, app, calls: int = 100, period: int = 60): | |
| super().__init__(app) | |
| self.calls = calls | |
| self.period = period | |
| self.clients: Dict[str, list] = defaultdict(list) | |
| self.blocked_clients: Dict[str, float] = {} # IP -> blocked_until timestamp | |
| # Specific limits for sensitive endpoints | |
| self.endpoint_limits = { | |
| '/api/auth/patient/login': (5, 900), # 5 attempts per 15 minutes | |
| '/api/auth/therapist/login': (5, 900), | |
| '/api/auth/patient/signup': (3, 3600), # 3 signups per hour | |
| '/api/auth/therapist/signup': (3, 3600), | |
| } | |
| def get_client_ip(self, request: Request) -> str: | |
| """Get real client IP considering proxies""" | |
| # Check for real IP behind proxy | |
| forwarded = request.headers.get("X-Forwarded-For") | |
| if forwarded: | |
| return forwarded.split(",")[0] | |
| return request.client.host | |
| async def dispatch(self, request: Request, call_next): | |
| client_ip = self.get_client_ip(request) | |
| path = request.url.path | |
| current_time = time.time() | |
| # Check if client is blocked | |
| if client_ip in self.blocked_clients: | |
| if current_time < self.blocked_clients[client_ip]: | |
| remaining = int(self.blocked_clients[client_ip] - current_time) | |
| raise HTTPException( | |
| status_code=status.HTTP_429_TOO_MANY_REQUESTS, | |
| detail=f"Too many requests. Please try again in {remaining} seconds." | |
| ) | |
| else: | |
| # Unblock client | |
| del self.blocked_clients[client_ip] | |
| # Get limits for this endpoint | |
| if path in self.endpoint_limits: | |
| limit_calls, limit_period = self.endpoint_limits[path] | |
| else: | |
| limit_calls, limit_period = self.calls, self.period | |
| # Clean old requests | |
| min_time = current_time - limit_period | |
| self.clients[client_ip] = [ | |
| req_time for req_time in self.clients[client_ip] | |
| if req_time > min_time | |
| ] | |
| # Check rate limit | |
| if len(self.clients[client_ip]) >= limit_calls: | |
| # Block client temporarily | |
| self.blocked_clients[client_ip] = current_time + 300 # Block for 5 minutes | |
| raise HTTPException( | |
| status_code=status.HTTP_429_TOO_MANY_REQUESTS, | |
| detail="Too many requests. Please slow down." | |
| ) | |
| # Record this request | |
| self.clients[client_ip].append(current_time) | |
| # Process request | |
| response = await call_next(request) | |
| return response | |
| class SecurityHeadersMiddleware(BaseHTTPMiddleware): | |
| """Add security headers to all responses""" | |
| async def dispatch(self, request: Request, call_next): | |
| response = await call_next(request) | |
| # Security headers | |
| response.headers["X-Content-Type-Options"] = "nosniff" | |
| response.headers["X-Frame-Options"] = "DENY" | |
| response.headers["X-XSS-Protection"] = "1; mode=block" | |
| response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains" | |
| response.headers["Content-Security-Policy"] = "default-src 'self'" | |
| response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" | |
| response.headers["Permissions-Policy"] = "geolocation=(), microphone=(), camera=()" | |
| return response | |
| class LoginAttemptTracker: | |
| """Track failed login attempts for additional security""" | |
| def __init__(self): | |
| self.attempts: Dict[str, list] = defaultdict(list) | |
| self.locked_accounts: Dict[str, float] = {} | |
| def record_failed_attempt(self, identifier: str) -> None: | |
| """Record a failed login attempt""" | |
| current_time = time.time() | |
| self.attempts[identifier].append(current_time) | |
| # Clean old attempts (older than 15 minutes) | |
| min_time = current_time - SecurityConfig.LOGIN_ATTEMPTS_WINDOW | |
| self.attempts[identifier] = [ | |
| t for t in self.attempts[identifier] if t > min_time | |
| ] | |
| # Lock account if too many attempts | |
| if len(self.attempts[identifier]) >= SecurityConfig.LOGIN_ATTEMPTS_LIMIT: | |
| self.locked_accounts[identifier] = current_time + 900 # Lock for 15 minutes | |
| def is_locked(self, identifier: str) -> tuple[bool, Optional[int]]: | |
| """Check if account is locked""" | |
| if identifier in self.locked_accounts: | |
| current_time = time.time() | |
| if current_time < self.locked_accounts[identifier]: | |
| remaining = int(self.locked_accounts[identifier] - current_time) | |
| return True, remaining | |
| else: | |
| # Unlock account | |
| del self.locked_accounts[identifier] | |
| if identifier in self.attempts: | |
| del self.attempts[identifier] | |
| return False, None | |
| def clear_attempts(self, identifier: str) -> None: | |
| """Clear attempts after successful login""" | |
| if identifier in self.attempts: | |
| del self.attempts[identifier] | |
| if identifier in self.locked_accounts: | |
| del self.locked_accounts[identifier] | |
| # Global instance for login tracking | |
| login_tracker = LoginAttemptTracker() | |
| def setup_security_middleware(app): | |
| """Setup all security middleware for the application""" | |
| # CORS middleware | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=SecurityConfig.ALLOWED_ORIGINS, | |
| allow_credentials=True, | |
| allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"], | |
| allow_headers=["*"], | |
| expose_headers=["*"] | |
| ) | |
| # Trusted host middleware (prevent host header attacks) | |
| if os.getenv('ENVIRONMENT') == 'production': | |
| allowed_hosts = os.getenv('ALLOWED_HOSTS', 'localhost').split(',') | |
| app.add_middleware(TrustedHostMiddleware, allowed_hosts=allowed_hosts) | |
| # Rate limiting | |
| app.add_middleware(RateLimitMiddleware) | |
| # Security headers | |
| app.add_middleware(SecurityHeadersMiddleware) | |
| return app |