backend/security.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import time
from collections import defaultdict, deque
from threading import Lock

from fastapi import HTTPException, Request, status


class LoginRateLimiter:
    def __init__(self, max_attempts: int = 6, window_seconds: int = 10 * 60):
        self.max_attempts = max_attempts
        self.window_seconds = window_seconds
        self._attempts: dict[str, deque[float]] = defaultdict(deque)
        self._lock = Lock()

    def check(self, client_key: str) -> None:
        now = time.time()
        with self._lock:
            attempts = self._attempts[client_key]
            while attempts and attempts[0] < now - self.window_seconds:
                attempts.popleft()
            if len(attempts) >= self.max_attempts:
                raise HTTPException(
                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                    detail="Too many sign-in attempts. Try again in a few minutes.",
                )

    def record_failure(self, client_key: str) -> None:
        with self._lock:
            self._attempts[client_key].append(time.time())

    def clear(self, client_key: str) -> None:
        with self._lock:
            self._attempts.pop(client_key, None)


def passwords_match(provided: str, expected: str) -> bool:
    return hmac.compare_digest(
        provided.encode("utf-8"),
        expected.encode("utf-8"),
    )


def _signing_key(password: str) -> bytes:
    return hashlib.sha256(
        f"vayuchat-session-v1:{password}".encode("utf-8")
    ).digest()


def create_session_token(password: str, ttl_seconds: int) -> str:
    payload = {
        "exp": int(time.time()) + ttl_seconds,
        "nonce": secrets.token_urlsafe(16),
    }
    encoded = base64.urlsafe_b64encode(
        json.dumps(payload, separators=(",", ":")).encode("utf-8")
    ).rstrip(b"=")
    signature = hmac.new(_signing_key(password), encoded, hashlib.sha256).digest()
    return ".".join(
        (
            encoded.decode("ascii"),
            base64.urlsafe_b64encode(signature).rstrip(b"=").decode("ascii"),
        )
    )


def verify_session_token(token: str, password: str) -> bool:
    try:
        encoded_text, signature_text = token.split(".", maxsplit=1)
        encoded = encoded_text.encode("ascii")
        expected = hmac.new(_signing_key(password), encoded, hashlib.sha256).digest()
        padded_signature = signature_text + "=" * (-len(signature_text) % 4)
        supplied = base64.urlsafe_b64decode(padded_signature)
        if not hmac.compare_digest(supplied, expected):
            return False
        padded_payload = encoded_text + "=" * (-len(encoded_text) % 4)
        payload = json.loads(base64.urlsafe_b64decode(padded_payload))
        return int(payload["exp"]) > int(time.time())
    except (ValueError, KeyError, TypeError, json.JSONDecodeError):
        return False


def client_key(request: Request) -> str:
    host = request.client.host if request.client else "unknown"
    user_agent = request.headers.get("user-agent", "")[:160]
    return hashlib.sha256(f"{host}:{user_agent}".encode("utf-8")).hexdigest()


def require_same_origin(request: Request) -> None:
    origin = request.headers.get("origin")
    if not origin:
        return
    forwarded_host = request.headers.get("x-forwarded-host")
    expected_host = forwarded_host or request.headers.get("host", "")
    if not expected_host or origin.split("://", maxsplit=1)[-1] != expected_host:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Cross-origin request rejected.",
        )