File size: 2,492 Bytes
4b9d59b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """
Core security utilities for the Atom platform.
This package provides:
- Rate limiting middleware
- Security headers middleware
- Role-based access control (RBAC)
"""
import logging
import time
from collections import defaultdict
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from .rbac import require_role
logger = logging.getLogger(__name__)
class RateLimitMiddleware(BaseHTTPMiddleware):
def __init__(self, app, requests_per_minute: int = 60):
super().__init__(app)
self.requests_per_minute = requests_per_minute
self.request_counts = defaultdict(list)
async def dispatch(self, request: Request, call_next):
# Allow OPTIONS requests through without rate limiting to avoid CORS preflight issues
if request.method == "OPTIONS":
return await call_next(request)
# Identify client by IP
client_ip = request.client.host
# Clean old requests
current_time = time.time()
self.request_counts[client_ip] = [
t for t in self.request_counts[client_ip]
if current_time - t < 60
]
# Check limit
if len(self.request_counts[client_ip]) >= self.requests_per_minute:
logger.warning(f"Rate limit exceeded for {client_ip}")
return Response("Rate limit exceeded", status_code=429)
# Record request
self.request_counts[client_ip].append(current_time)
# Process request
response = await call_next(request)
return response
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
response = await call_next(request)
# Add 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' 'unsafe-inline' 'unsafe-eval' cdn.jsdelivr.net fonts.googleapis.com fonts.gstatic.com; connect-src 'self' http://localhost:* http://127.0.0.1:* ws://localhost:* ws://127.0.0.1:*; img-src 'self' data: https:;"
return response
__all__ = ["require_role", "RateLimitMiddleware", "SecurityHeadersMiddleware"]
|