geointel/backend/app/services/auth_service.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Literal, cast
from uuid import UUID

from app.core.config import Settings
from app.core.public_demo import PUBLIC_DEMO_PROJECT_ID


@dataclass(frozen=True)
class AuthPrincipal:
    username: str
    expires_at: int
    session_id: str = field(default_factory=lambda: secrets.token_urlsafe(12))
    role: Literal["operator", "guest"] = "operator"
    project_id: UUID | None = None


class AuthService:
    HASH_NAME = "pbkdf2_sha256"
    HASH_ITERATIONS = 600_000
    MAX_FAILURES = 5
    FAILURE_WINDOW_SECONDS = 300
    _failures: dict[str, deque[float]] = {}
    _failure_lock = threading.Lock()
    _guest_requests: dict[str, deque[float]] = {}
    _guest_request_lock = threading.Lock()
    _active_guest_compute = 0
    _guest_compute_lock = threading.Lock()

    @staticmethod
    def _b64_encode(value: bytes) -> str:
        return base64.urlsafe_b64encode(value).decode("ascii").rstrip("=")

    @staticmethod
    def _b64_decode(value: str) -> bytes:
        return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))

    @classmethod
    def hash_password(
        cls,
        password: str,
        *,
        salt: bytes | None = None,
        iterations: int | None = None,
    ) -> str:
        resolved_salt = salt or secrets.token_bytes(18)
        resolved_iterations = iterations or cls.HASH_ITERATIONS
        digest = hashlib.pbkdf2_hmac(
            "sha256",
            password.encode("utf-8"),
            resolved_salt,
            resolved_iterations,
        )
        return "$".join(
            (
                cls.HASH_NAME,
                str(resolved_iterations),
                cls._b64_encode(resolved_salt),
                cls._b64_encode(digest),
            )
        )

    @classmethod
    def verify_password(cls, password: str, encoded: str) -> bool:
        try:
            algorithm, iterations_raw, salt_raw, expected_raw = encoded.split("$", 3)
            if algorithm != cls.HASH_NAME:
                return False
            iterations = int(iterations_raw)
            if iterations < 100_000 or iterations > 2_000_000:
                return False
            salt = cls._b64_decode(salt_raw)
            expected = cls._b64_decode(expected_raw)
            actual = hashlib.pbkdf2_hmac(
                "sha256",
                password.encode("utf-8"),
                salt,
                iterations,
            )
            return hmac.compare_digest(actual, expected)
        except (TypeError, ValueError):
            return False

    @classmethod
    def credentials_match(cls, username: str, password: str, settings: Settings) -> bool:
        expected_username = settings.auth_username or ""
        expected_password_hash = settings.auth_password_hash or ""
        username_matches = hmac.compare_digest(
            username.encode("utf-8"),
            expected_username.encode("utf-8"),
        )
        password_matches = cls.verify_password(password, expected_password_hash)
        return username_matches and password_matches

    @classmethod
    def create_session_token(
        cls,
        username: str,
        settings: Settings,
        *,
        role: Literal["operator", "guest"] = "operator",
        project_id: UUID | None = None,
        ttl_seconds: int | None = None,
        now: int | None = None,
    ) -> str:
        issued_at = int(time.time() if now is None else now)
        if role == "guest" and project_id is None:
            raise ValueError("Guest sessions must be scoped to a demo project")
        resolved_ttl = ttl_seconds if ttl_seconds is not None else (
            settings.guest_session_ttl_seconds if role == "guest" else settings.auth_session_ttl_seconds
        )
        payload = {
            "exp": issued_at + resolved_ttl,
            "iat": issued_at,
            "jti": secrets.token_urlsafe(12),
            "role": role,
            "sub": username,
            "v": 2,
        }
        if project_id is not None:
            payload["project_id"] = str(project_id)
        encoded_payload = cls._b64_encode(
            json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
        )
        signature = hmac.new(
            (settings.auth_session_secret or "").encode("utf-8"),
            encoded_payload.encode("ascii"),
            hashlib.sha256,
        ).digest()
        return f"{encoded_payload}.{cls._b64_encode(signature)}"

    @classmethod
    def verify_session_token(
        cls,
        token: str | None,
        settings: Settings,
        *,
        now: int | None = None,
    ) -> AuthPrincipal | None:
        if not token:
            return None
        try:
            encoded_payload, encoded_signature = token.split(".", 1)
            expected_signature = hmac.new(
                (settings.auth_session_secret or "").encode("utf-8"),
                encoded_payload.encode("ascii"),
                hashlib.sha256,
            ).digest()
            supplied_signature = cls._b64_decode(encoded_signature)
            if not hmac.compare_digest(expected_signature, supplied_signature):
                return None
            payload = json.loads(cls._b64_decode(encoded_payload))
            username = str(payload.get("sub") or "")
            expires_at = int(payload.get("exp") or 0)
            issued_at = int(payload.get("iat") or 0)
            version = int(payload.get("v") or 0)
            role_value = str(payload.get("role") or "operator")
            session_id = str(payload.get("jti") or "")
            current = int(time.time() if now is None else now)
            if version not in {1, 2} or role_value not in {"operator", "guest"} or not session_id:
                return None
            role = cast(Literal["operator", "guest"], role_value)
            if issued_at <= 0 or issued_at > current + 60 or expires_at <= current:
                return None
            if role == "operator":
                if username != settings.auth_username:
                    return None
                max_ttl = settings.auth_session_ttl_seconds
                project_id = None
            else:
                if not settings.guest_access_enabled or username != settings.guest_display_name:
                    return None
                max_ttl = settings.guest_session_ttl_seconds
                raw_project_id = payload.get("project_id")
                if not raw_project_id:
                    return None
                project_id = UUID(str(raw_project_id))
                if project_id != PUBLIC_DEMO_PROJECT_ID:
                    return None
            if expires_at - issued_at > max_ttl:
                return None
            return AuthPrincipal(
                username=username,
                expires_at=expires_at,
                session_id=session_id,
                role=role,
                project_id=project_id,
            )
        except (TypeError, ValueError, json.JSONDecodeError, UnicodeDecodeError):
            return None

    @classmethod
    def retry_after_seconds(cls, key: str, *, now: float | None = None) -> int:
        current = time.monotonic() if now is None else now
        with cls._failure_lock:
            attempts = cls._failures.setdefault(key, deque())
            while attempts and current - attempts[0] >= cls.FAILURE_WINDOW_SECONDS:
                attempts.popleft()
            if len(attempts) < cls.MAX_FAILURES:
                if not attempts:
                    cls._failures.pop(key, None)
                return 0
            return max(1, int(cls.FAILURE_WINDOW_SECONDS - (current - attempts[0])))

    @classmethod
    def record_failure(cls, key: str, *, now: float | None = None) -> None:
        current = time.monotonic() if now is None else now
        with cls._failure_lock:
            attempts = cls._failures.setdefault(key, deque())
            while attempts and current - attempts[0] >= cls.FAILURE_WINDOW_SECONDS:
                attempts.popleft()
            attempts.append(current)

    @classmethod
    def clear_failures(cls, key: str) -> None:
        with cls._failure_lock:
            cls._failures.pop(key, None)

    @classmethod
    def consume_guest_request(
        cls,
        key: str,
        *,
        max_requests: int,
        window_seconds: int = 60,
        now: float | None = None,
    ) -> int:
        """Record a guest action and return Retry-After seconds when limited."""
        current = time.monotonic() if now is None else now
        with cls._guest_request_lock:
            attempts = cls._guest_requests.setdefault(key, deque())
            while attempts and current - attempts[0] >= window_seconds:
                attempts.popleft()
            if len(attempts) >= max_requests:
                return max(1, int(window_seconds - (current - attempts[0])))
            attempts.append(current)
            return 0

    @classmethod
    def try_acquire_guest_compute(cls, *, max_concurrency: int) -> bool:
        with cls._guest_compute_lock:
            if cls._active_guest_compute >= max_concurrency:
                return False
            cls._active_guest_compute += 1
            return True

    @classmethod
    def release_guest_compute(cls) -> None:
        with cls._guest_compute_lock:
            cls._active_guest_compute = max(0, cls._active_guest_compute - 1)