utils/security.py

"""Security utilities — rate limiting, validation, sanitization."""

import re
import time
import logging
from collections import defaultdict
from typing import Optional

from config import config

logger = logging.getLogger("synapse.security")


class RateLimiter:
    """Simple in-memory rate limiter."""

    def __init__(self, max_per_minute: int = None):
        self.max_per_minute = max_per_minute or config.security.rate_limit_per_minute
        self._requests: dict[str, list[float]] = defaultdict(list)

    def check(self, client_id: str) -> bool:
        now = time.time()
        window = now - 60
        self._requests[client_id] = [
            t for t in self._requests[client_id] if t > window
        ]
        if len(self._requests[client_id]) >= self.max_per_minute:
            logger.warning(f"Rate limit exceeded for {client_id}")
            return False
        self._requests[client_id].append(now)
        return True

    def get_remaining(self, client_id: str) -> int:
        now = time.time()
        window = now - 60
        recent = [t for t in self._requests[client_id] if t > window]
        return max(0, self.max_per_minute - len(recent))


rate_limiter = RateLimiter()


def sanitize_input(text: str, max_length: int = 10000) -> str:
    if not text:
        return ""
    text = text[:max_length]
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
    return text.strip()


def validate_email(email: str) -> bool:
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
    return bool(re.match(pattern, email))


def validate_username(username: str) -> bool:
    pattern = r'^[a-zA-Z0-9_]{3,30}$'
    return bool(re.match(pattern, username))


def validate_file_extension(filename: str) -> bool:
    from pathlib import Path
    ext = Path(filename).suffix.lower()
    return ext in config.security.allowed_extensions


def validate_file_size(size_bytes: int) -> bool:
    max_bytes = config.security.max_upload_size_mb * 1024 * 1024
    return size_bytes <= max_bytes


def mask_sensitive(text: str) -> str:
    patterns = [
        (r'(GROQ_API_KEY=)[^\s]+', r'\1***'),
        (r'(HF_API_TOKEN=)[^\s]+', r'\1***'),
        (r'(sk-[a-zA-Z0-9]+)', r'***'),
        (r'(gsk_[a-zA-Z0-9]+)', r'***'),
    ]
    for pattern, replacement in patterns:
        text = re.sub(pattern, replacement, text)
    return text


class InputValidator:
    """Validates and sanitizes various input types."""

    @staticmethod
    def validate_chat_message(content: str) -> tuple[bool, str]:
        if not content or not content.strip():
            return False, "Message cannot be empty"
        if len(content) > 50000:
            return False, "Message too long (max 50,000 characters)"
        return True, sanitize_input(content)

    @staticmethod
    def validate_prompt(prompt: str) -> tuple[bool, str]:
        if not prompt or not prompt.strip():
            return False, "Prompt cannot be empty"
        if len(prompt) > 10000:
            return False, "Prompt too long (max 10,000 characters)"
        return True, sanitize_input(prompt)

    @staticmethod
    def validate_image_params(width: int, height: int, steps: int) -> tuple[bool, str]:
        if not (128 <= width <= 2048):
            return False, "Width must be between 128 and 2048"
        if not (128 <= height <= 2048):
            return False, "Height must be between 128 and 2048"
        if not (1 <= steps <= 100):
            return False, "Steps must be between 1 and 100"
        return True, "ok"


validator = InputValidator()