Cognitive_Agent/src/cognitive_agent/auth.py

"""Authentication and abuse controls for the public demo server."""

from __future__ import annotations

import argparse
import base64
import getpass
import hashlib
import hmac
import json
import os
import secrets
import time
from collections import defaultdict, deque
from dataclasses import dataclass
from threading import RLock


PASSWORD_SCHEME = "pbkdf2_sha256"
PASSWORD_ITERATIONS = 600_000


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


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


def hash_password(
    password: str,
    *,
    salt: bytes | None = None,
    iterations: int = PASSWORD_ITERATIONS,
) -> str:
    if len(password) < 12:
        raise ValueError("demo password must contain at least 12 characters")
    salt = salt or secrets.token_bytes(16)
    digest = hashlib.pbkdf2_hmac(
        "sha256",
        password.encode("utf-8"),
        salt,
        iterations,
    )
    return f"{PASSWORD_SCHEME}${iterations}${_b64encode(salt)}${_b64encode(digest)}"


def verify_password(password: str, encoded: str) -> bool:
    try:
        scheme, iteration_text, salt_text, digest_text = encoded.split("$", 3)
        if scheme != PASSWORD_SCHEME:
            return False
        iterations = int(iteration_text)
        if not 100_000 <= iterations <= 2_000_000:
            return False
        expected = _b64decode(digest_text)
        actual = hashlib.pbkdf2_hmac(
            "sha256",
            password.encode("utf-8"),
            _b64decode(salt_text),
            iterations,
        )
        return hmac.compare_digest(actual, expected)
    except (ValueError, TypeError):
        return False


def _env_bool(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


@dataclass(frozen=True)
class AuthConfig:
    enabled: bool
    username: str
    password_hash: str
    session_secret: bytes
    session_ttl_seconds: int
    login_attempts: int
    login_window_seconds: int
    requests_per_minute: int
    trust_proxy_headers: bool
    cookie_name: str = "cognitive_agent_session"

    @classmethod
    def from_env(cls) -> AuthConfig:
        enabled = _env_bool("DEMO_AUTH_ENABLED")
        username = os.getenv("DEMO_AUTH_USERNAME", "").strip()
        password_hash = os.getenv("DEMO_AUTH_PASSWORD_HASH", "").strip()
        secret_text = os.getenv("DEMO_AUTH_SESSION_SECRET", "").strip()
        if enabled:
            if not username or len(username) > 64:
                raise RuntimeError("DEMO_AUTH_USERNAME must contain 1-64 characters")
            if not password_hash.startswith(f"{PASSWORD_SCHEME}$"):
                raise RuntimeError("DEMO_AUTH_PASSWORD_HASH is missing or unsupported")
            if len(secret_text) < 43:
                raise RuntimeError("DEMO_AUTH_SESSION_SECRET must contain at least 43 characters")
        return cls(
            enabled=enabled,
            username=username,
            password_hash=password_hash,
            session_secret=secret_text.encode("utf-8"),
            session_ttl_seconds=max(
                300,
                min(int(os.getenv("DEMO_AUTH_SESSION_TTL_SECONDS", "28800")), 604800),
            ),
            login_attempts=max(
                1,
                min(int(os.getenv("DEMO_AUTH_LOGIN_ATTEMPTS", "5")), 20),
            ),
            login_window_seconds=max(
                60,
                min(int(os.getenv("DEMO_AUTH_LOGIN_WINDOW_SECONDS", "300")), 3600),
            ),
            requests_per_minute=max(
                1,
                min(int(os.getenv("DEMO_REQUESTS_PER_MINUTE", "10")), 120),
            ),
            trust_proxy_headers=_env_bool("DEMO_TRUST_PROXY_HEADERS"),
        )


class SessionCodec:
    def __init__(self, config: AuthConfig) -> None:
        self.config = config

    def issue(self, username: str, now: int | None = None) -> str:
        now = int(time.time()) if now is None else now
        payload = {
            "sub": username,
            "iat": now,
            "exp": now + self.config.session_ttl_seconds,
            "nonce": secrets.token_urlsafe(12),
        }
        encoded = _b64encode(
            json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
        )
        signature = _b64encode(
            hmac.new(self.config.session_secret, encoded.encode("ascii"), hashlib.sha256).digest()
        )
        return f"{encoded}.{signature}"

    def verify(self, token: str, now: int | None = None) -> str | None:
        if not token or len(token) > 2048:
            return None
        try:
            encoded, signature = token.split(".", 1)
            expected = _b64encode(
                hmac.new(
                    self.config.session_secret,
                    encoded.encode("ascii"),
                    hashlib.sha256,
                ).digest()
            )
            if not hmac.compare_digest(signature, expected):
                return None
            payload = json.loads(_b64decode(encoded))
            current_time = int(time.time()) if now is None else now
            if payload.get("sub") != self.config.username:
                return None
            if not int(payload.get("iat", 0)) <= current_time < int(payload.get("exp", 0)):
                return None
            return str(payload["sub"])
        except (ValueError, TypeError, json.JSONDecodeError):
            return None


class SlidingWindowLimiter:
    def __init__(self, limit: int, window_seconds: int) -> None:
        self.limit = limit
        self.window_seconds = window_seconds
        self._events: dict[str, deque[float]] = defaultdict(deque)
        self._lock = RLock()

    def retry_after(self, key: str, now: float | None = None) -> int:
        current = time.monotonic() if now is None else now
        with self._lock:
            events = self._events[key]
            cutoff = current - self.window_seconds
            while events and events[0] <= cutoff:
                events.popleft()
            if len(events) < self.limit:
                return 0
            return max(1, int(self.window_seconds - (current - events[0]) + 0.999))

    def record(self, key: str, now: float | None = None) -> None:
        current = time.monotonic() if now is None else now
        with self._lock:
            self._events[key].append(current)

    def clear(self, key: str) -> None:
        with self._lock:
            self._events.pop(key, None)


class AuthRuntime:
    def __init__(self, config: AuthConfig) -> None:
        self.config = config
        self.sessions = SessionCodec(config)
        self.login_failures = SlidingWindowLimiter(
            config.login_attempts,
            config.login_window_seconds,
        )
        self.api_requests = SlidingWindowLimiter(config.requests_per_minute, 60)


def main() -> None:
    parser = argparse.ArgumentParser(description="Cognitive Agent authentication utilities")
    parser.add_argument("command", choices=["hash-password", "generate-secret"])
    args = parser.parse_args()
    if args.command == "generate-secret":
        print(secrets.token_urlsafe(48))
        return
    password = getpass.getpass("Demo password: ")
    confirmation = getpass.getpass("Confirm password: ")
    if password != confirmation:
        raise SystemExit("Passwords do not match")
    print(hash_password(password))


if __name__ == "__main__":
    main()