Security Hardening¶

# Application TLS configuration
import ssl
from fastapi import FastAPI
from fastapi.middleware.httpsredirect import HTTPSRedirectMiddleware

app = FastAPI()

# Force HTTPS redirect
app.add_middleware(HTTPSRedirectMiddleware)

# TLS configuration
ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
ssl_context.load_cert_chain("cert.pem", "key.pem")
ssl_context.set_ciphers("ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES128-GCM-SHA256")
ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2

# Input validation
from pydantic import BaseModel, validator
import re

class SearchRequest(BaseModel):
    query: str
    collection_id: str
    limit: int = 10

    @validator('query')
    def validate_query(cls, v):
        if not v or len(v.strip()) == 0:
            raise ValueError('Query cannot be empty')
        if len(v) > 1000:
            raise ValueError('Query too long')
        # Check for SQL injection patterns
        if re.search(r'[;\'"]', v):
            raise ValueError('Invalid characters in query')
        return v.strip()

    @validator('collection_id')
    def validate_collection_id(cls, v):
        if not re.match(r'^[a-zA-Z0-9-_]+$', v):
            raise ValueError('Invalid collection ID format')
        return v

    @validator('limit')
    def validate_limit(cls, v):
        if v < 1 or v > 100:
            raise ValueError('Limit must be between 1 and 100')
        return v

# SQL injection prevention
import psycopg2
from psycopg2 import sql

def safe_query(cursor, query_template, params):
    """Execute query with parameterized statements"""
    try:
        cursor.execute(query_template, params)
        return cursor.fetchall()
    except psycopg2.Error as e:
        logger.error(f"Database error: {e}")
        raise HTTPException(status_code=500, detail="Database error")

# Example usage
def search_documents(collection_id: str, query: str, limit: int):
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            # Use parameterized query
            query_template = """
                SELECT id, title, content, created_at
                FROM documents
                WHERE collection_id = %s
                AND content ILIKE %s
                ORDER BY created_at DESC
                LIMIT %s
            """
            params = (collection_id, f"%{query}%", limit)
            return safe_query(cursor, query_template, params)

# JWT authentication
import jwt
from datetime import datetime, timedelta
from fastapi import HTTPException, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

security = HTTPBearer()

class JWTAuth:
    def __init__(self, secret_key: str, algorithm: str = "HS256"):
        self.secret_key = secret_key
        self.algorithm = algorithm

    def create_token(self, user_id: str, expires_delta: timedelta = None):
        """Create JWT token"""
        if expires_delta:
            expire = datetime.utcnow() + expires_delta
        else:
            expire = datetime.utcnow() + timedelta(hours=24)

        payload = {
            "user_id": user_id,
            "exp": expire,
            "iat": datetime.utcnow()
        }

        return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)

    def verify_token(self, token: str):
        """Verify JWT token"""
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
            return payload
        except jwt.ExpiredSignatureError:
            raise HTTPException(status_code=401, detail="Token expired")
        except jwt.JWTError:
            raise HTTPException(status_code=401, detail="Invalid token")

# Authentication dependency
async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
    jwt_auth = JWTAuth(os.getenv("JWT_SECRET"))
    payload = jwt_auth.verify_token(credentials.credentials)
    return payload["user_id"]

# Role-based access control
from enum import Enum
from functools import wraps

class Role(Enum):
    ADMIN = "admin"
    USER = "user"
    READONLY = "readonly"

def require_role(required_role: Role):
    """Decorator to require specific role"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            user_id = kwargs.get("current_user")
            user_role = get_user_role(user_id)

            if not has_permission(user_role, required_role):
                raise HTTPException(status_code=403, detail="Insufficient permissions")

            return await func(*args, **kwargs)
        return wrapper
    return decorator

def has_permission(user_role: Role, required_role: Role) -> bool:
    """Check if user has required permission"""
    role_hierarchy = {
        Role.ADMIN: [Role.ADMIN, Role.USER, Role.READONLY],
        Role.USER: [Role.USER, Role.READONLY],
        Role.READONLY: [Role.READONLY]
    }

    return required_role in role_hierarchy.get(user_role, [])

# Usage example
@app.post("/api/collections")
@require_role(Role.ADMIN)
async def create_collection(
    collection: CollectionCreate,
    current_user: str = Depends(get_current_user)
):
    # Only admins can create collections
    pass

# Rate limiting
from fastapi import FastAPI, Request
from fastapi.middleware.base import BaseHTTPMiddleware
import time
from collections import defaultdict

class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, calls: int = 100, period: int = 60):
        super().__init__(app)
        self.calls = calls
        self.period = period
        self.clients = defaultdict(list)

    async def dispatch(self, request: Request, call_next):
        client_ip = request.client.host
        now = time.time()

        # Clean old requests
        self.clients[client_ip] = [
            req_time for req_time in self.clients[client_ip]
            if now - req_time < self.period
        ]

        # Check rate limit
        if len(self.clients[client_ip]) >= self.calls:
            return JSONResponse(
                status_code=429,
                content={"detail": "Rate limit exceeded"}
            )

        # Add current request
        self.clients[client_ip].append(now)

        response = await call_next(request)
        return response

# Apply rate limiting
app.add_middleware(RateLimitMiddleware, calls=100, period=60)

# Security event collection
import logging
from datetime import datetime
from typing import Dict, Any

class SecurityEventLogger:
    def __init__(self):
        self.logger = logging.getLogger("security")
        self.logger.setLevel(logging.INFO)

        # Create security event handler
        handler = logging.StreamHandler()
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)

    def log_auth_failure(self, user_id: str, ip_address: str, reason: str):
        """Log authentication failure"""
        self.logger.warning(
            f"Authentication failure - User: {user_id}, IP: {ip_address}, Reason: {reason}"
        )

    def log_suspicious_activity(self, activity: str, details: Dict[str, Any]):
        """Log suspicious activity"""
        self.logger.warning(
            f"Suspicious activity - {activity}: {details}"
        )

    def log_security_event(self, event_type: str, details: Dict[str, Any]):
        """Log general security event"""
        self.logger.info(
            f"Security event - {event_type}: {details}"
        )

# Global security logger
security_logger = SecurityEventLogger()

# Security metrics
from prometheus_client import Counter, Histogram, Gauge

# Security event counters
auth_failures = Counter('auth_failures_total', 'Total authentication failures', ['user_id', 'reason'])
suspicious_activities = Counter('suspicious_activities_total', 'Total suspicious activities', ['activity_type'])
security_events = Counter('security_events_total', 'Total security events', ['event_type'])

# Security response time
security_response_time = Histogram('security_response_time_seconds', 'Security response time')

# Active security threats
active_threats = Gauge('active_threats', 'Number of active security threats')

# Example usage
def log_auth_failure(user_id: str, reason: str):
    auth_failures.labels(user_id=user_id, reason=reason).inc()
    security_logger.log_auth_failure(user_id, get_client_ip(), reason)

def log_suspicious_activity(activity: str, details: Dict[str, Any]):
    suspicious_activities.labels(activity_type=activity).inc()
    security_logger.log_suspicious_activity(activity, details)

# Incident response automation
import asyncio
from datetime import datetime
from typing import List, Dict

class SecurityIncidentResponse:
    def __init__(self):
        self.active_incidents = {}
        self.response_team = ["devops@company.com", "security@company.com"]

    async def handle_security_alert(self, alert: Dict[str, Any]):
        """Handle security alert"""
        incident_id = f"SEC-{datetime.now().strftime('%Y%m%d-%H%M%S')}"

        # Create incident
        incident = {
            "id": incident_id,
            "type": alert["type"],
            "severity": alert["severity"],
            "timestamp": datetime.now(),
            "status": "open",
            "details": alert["details"]
        }

        self.active_incidents[incident_id] = incident

        # Notify response team
        await self.notify_response_team(incident)

        # Take automated actions
        await self.take_automated_actions(incident)

        return incident_id

    async def notify_response_team(self, incident: Dict[str, Any]):
        """Notify security response team"""
        # Send email notification
        await self.send_email_notification(incident)

        # Send Slack notification
        await self.send_slack_notification(incident)

    async def take_automated_actions(self, incident: Dict[str, Any]):
        """Take automated security actions"""
        if incident["severity"] == "critical":
            # Block suspicious IP
            await self.block_suspicious_ip(incident["details"]["ip_address"])

            # Increase monitoring
            await self.increase_monitoring(incident["details"]["user_id"])

            # Generate security report
            await self.generate_security_report(incident)

# Global incident response
incident_response = SecurityIncidentResponse()

# Audit logging
import json
from datetime import datetime
from typing import Dict, Any

class AuditLogger:
    def __init__(self):
        self.logger = logging.getLogger("audit")
        self.logger.setLevel(logging.INFO)

        # Create audit handler
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(message)s')
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)

    def log_user_action(self, user_id: str, action: str, resource: str, details: Dict[str, Any]):
        """Log user action"""
        audit_event = {
            "timestamp": datetime.utcnow().isoformat(),
            "event_type": "user_action",
            "user_id": user_id,
            "action": action,
            "resource": resource,
            "details": details,
            "ip_address": get_client_ip(),
            "user_agent": get_user_agent()
        }

        self.logger.info(json.dumps(audit_event))

    def log_system_event(self, event_type: str, details: Dict[str, Any]):
        """Log system event"""
        audit_event = {
            "timestamp": datetime.utcnow().isoformat(),
            "event_type": "system_event",
            "system_event_type": event_type,
            "details": details
        }

        self.logger.info(json.dumps(audit_event))

    def log_security_event(self, event_type: str, details: Dict[str, Any]):
        """Log security event"""
        audit_event = {
            "timestamp": datetime.utcnow().isoformat(),
            "event_type": "security_event",
            "security_event_type": event_type,
            "details": details
        }

        self.logger.info(json.dumps(audit_event))

# Global audit logger
audit_logger = AuditLogger()

# Compliance reporting
from datetime import datetime, timedelta
from typing import List, Dict

class ComplianceReporter:
    def __init__(self):
        self.audit_logger = AuditLogger()

    def generate_compliance_report(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """Generate compliance report"""
        report = {
            "report_period": {
                "start": start_date.isoformat(),
                "end": end_date.isoformat()
            },
            "user_actions": self.get_user_actions(start_date, end_date),
            "system_events": self.get_system_events(start_date, end_date),
            "security_events": self.get_security_events(start_date, end_date),
            "compliance_summary": self.get_compliance_summary(start_date, end_date)
        }

        return report

    def get_user_actions(self, start_date: datetime, end_date: datetime) -> List[Dict[str, Any]]:
        """Get user actions for compliance report"""
        # Query audit logs for user actions
        # This would typically query a database or log aggregation system
        pass

    def get_system_events(self, start_date: datetime, end_date: datetime) -> List[Dict[str, Any]]:
        """Get system events for compliance report"""
        # Query audit logs for system events
        pass

    def get_security_events(self, start_date: datetime, end_date: datetime) -> List[Dict[str, Any]]:
        """Get security events for compliance report"""
        # Query audit logs for security events
        pass

    def get_compliance_summary(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """Get compliance summary"""
        return {
            "total_events": 0,
            "security_incidents": 0,
            "compliance_score": 100,
            "recommendations": []
        }

# Global compliance reporter
compliance_reporter = ComplianceReporter()

# Data classification
from enum import Enum
from typing import Dict, Any

class DataClassification(Enum):
    PUBLIC = "public"
    INTERNAL = "internal"
    CONFIDENTIAL = "confidential"
    RESTRICTED = "restricted"

class DataClassifier:
    def __init__(self):
        self.classification_rules = {
            "email": DataClassification.CONFIDENTIAL,
            "phone": DataClassification.CONFIDENTIAL,
            "ssn": DataClassification.RESTRICTED,
            "credit_card": DataClassification.RESTRICTED,
            "api_key": DataClassification.RESTRICTED,
            "password": DataClassification.RESTRICTED
        }

    def classify_data(self, data: Dict[str, Any]) -> Dict[str, DataClassification]:
        """Classify data based on content"""
        classifications = {}

        for key, value in data.items():
            classification = DataClassification.INTERNAL  # Default

            for pattern, data_class in self.classification_rules.items():
                if pattern.lower() in key.lower():
                    classification = data_class
                    break

            classifications[key] = classification

        return classifications

    def apply_data_protection(self, data: Dict[str, Any], classifications: Dict[str, DataClassification]) -> Dict[str, Any]:
        """Apply data protection based on classification"""
        protected_data = {}

        for key, value in data.items():
            classification = classifications.get(key, DataClassification.INTERNAL)

            if classification == DataClassification.RESTRICTED:
                # Mask or remove restricted data
                protected_data[key] = "***REDACTED***"
            elif classification == DataClassification.CONFIDENTIAL:
                # Partially mask confidential data
                if isinstance(value, str) and len(value) > 4:
                    protected_data[key] = value[:2] + "***" + value[-2:]
                else:
                    protected_data[key] = "***MASKED***"
            else:
                protected_data[key] = value

        return protected_data

# Global data classifier
data_classifier = DataClassifier()