CIAgent/apps/api/app/core/security.py

"""Password hashing and JWT helpers.

Password hashing uses Argon2 (via `argon2-cffi`) directly - it's the
currently recommended default and needs no extra abstraction layer.
JWTs are signed with `JWT_SECRET` (HS256); access tokens are short-lived,
refresh tokens are long-lived but stored server-side only as a hash so a
leaked DB row can't be replayed as a valid token by itself.
"""

from __future__ import annotations

import secrets
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from enum import StrEnum

import jwt
from argon2 import PasswordHasher
from argon2.exceptions import VerifyMismatchError
from fastapi import Request

from app.core.config import Settings

_hasher = PasswordHasher()


def get_client_ip(request: Request, settings: Settings) -> str:
    """The single source of truth for "what IP is this request from" - used
    by `is_localhost`, the IP throttle/ban engine, and anywhere else that
    needs to identify a caller. Today there's no reverse proxy in front of
    uvicorn in this stack, so the direct TCP peer (`request.client.host`) is
    the real originating address.

    Once deployed behind Cloudflare (or Nginx), the direct peer becomes the
    proxy itself, not the visitor - `settings.trusted_proxy_ip_header` (e.g.
    "CF-Connecting-IP") switches this to read the real address from that
    header instead. Only ever set this when it's actually known the proxy is
    in front and stripping/overwriting that header from untrusted clients -
    otherwise a client could simply forge it to spoof any IP. Left empty by
    default (trust the direct connection only) - see KNOWN_LIMITATIONS.md."""
    header_name = settings.trusted_proxy_ip_header
    if header_name:
        forwarded = request.headers.get(header_name)
        if forwarded:
            return forwarded.strip()
    return request.client.host if request.client is not None else "unknown"


def is_localhost(request: Request, settings: Settings) -> bool:
    """True when the request's resolved client IP (see `get_client_ip`) is
    this machine's loopback address - not merely "someone on the LAN" - or
    is explicitly listed in `settings.additional_trusted_local_ips` (empty
    by default; a narrow, opt-in escape hatch for Docker Desktop's bridge
    networking, where even host-originated traffic doesn't arrive as
    literal loopback - see KNOWN_LIMITATIONS.md)."""
    client_ip = get_client_ip(request, settings)
    if client_ip in ("127.0.0.1", "::1"):
        return True
    extra = {ip.strip() for ip in settings.additional_trusted_local_ips.split(",") if ip.strip()}
    return client_ip in extra


def hash_password(raw_password: str) -> str:
    return _hasher.hash(raw_password)


def verify_password(raw_password: str, password_hash: str) -> bool:
    try:
        return _hasher.verify(password_hash, raw_password)
    except VerifyMismatchError:
        return False


class TokenType(StrEnum):
    ACCESS = "access"
    REFRESH = "refresh"
    # Settings -> Database viewer: a short-lived token minted from a live
    # admin session, exchanged (via Nginx's db.ciagent.org /_auth route) for
    # a longer-lived DB_VIEWER_SESSION cookie. See app/api/v1/db_viewer.py.
    DB_VIEWER_BOOTSTRAP = "db_viewer_bootstrap"
    DB_VIEWER_SESSION = "db_viewer_session"


@dataclass(frozen=True)
class DecodedToken:
    user_id: uuid.UUID
    token_type: TokenType
    jti: str
    expires_at: datetime


def create_access_token(user_id: uuid.UUID, settings: Settings) -> str:
    return _encode_token(
        user_id, TokenType.ACCESS, timedelta(minutes=settings.jwt_access_token_minutes), settings
    )


def create_refresh_token(user_id: uuid.UUID, settings: Settings) -> tuple[str, str, datetime]:
    """Returns (raw_jwt, jti, expires_at). Caller stores a hash of `jti`, not the JWT itself."""
    expires_at = datetime.now(UTC) + timedelta(days=settings.jwt_refresh_token_days)
    jti = secrets.token_urlsafe(32)
    token = _encode_token(
        user_id,
        TokenType.REFRESH,
        timedelta(days=settings.jwt_refresh_token_days),
        settings,
        jti=jti,
    )
    return token, jti, expires_at


def create_db_viewer_bootstrap_token(user_id: uuid.UUID, settings: Settings) -> str:
    return _encode_token(user_id, TokenType.DB_VIEWER_BOOTSTRAP, timedelta(minutes=2), settings)


def create_db_viewer_session_token(user_id: uuid.UUID, settings: Settings) -> str:
    return _encode_token(user_id, TokenType.DB_VIEWER_SESSION, timedelta(minutes=60), settings)


def _encode_token(
    user_id: uuid.UUID,
    token_type: TokenType,
    expires_in: timedelta,
    settings: Settings,
    jti: str | None = None,
) -> str:
    now = datetime.now(UTC)
    payload = {
        "sub": str(user_id),
        "type": token_type.value,
        "iat": now,
        "exp": now + expires_in,
        "jti": jti or secrets.token_urlsafe(16),
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


class InvalidTokenError(Exception):
    pass


def decode_token(token: str, settings: Settings, expected_type: TokenType) -> DecodedToken:
    try:
        payload = jwt.decode(token, settings.jwt_secret, algorithms=["HS256"])
    except jwt.PyJWTError as exc:
        raise InvalidTokenError(str(exc)) from exc

    if payload.get("type") != expected_type.value:
        raise InvalidTokenError(f"Expected a {expected_type.value} token")

    try:
        user_id = uuid.UUID(payload["sub"])
    except (KeyError, ValueError) as exc:
        raise InvalidTokenError("Malformed token subject") from exc

    return DecodedToken(
        user_id=user_id,
        token_type=TokenType(payload["type"]),
        jti=payload["jti"],
        expires_at=datetime.fromtimestamp(payload["exp"], tz=UTC),
    )


def hash_token_identifier(jti: str) -> str:
    """One-way hash of a refresh token's `jti` for storage/comparison (not the JWT itself)."""
    import hashlib

    return hashlib.sha256(jti.encode("utf-8")).hexdigest()


def generate_email_code() -> str:
    """A 6-digit numeric code for email verification / password reset -
    short-lived and IP-throttled (app/services/ip_throttle_service.py), so
    it doesn't need Argon2's cost the way a password does."""
    return f"{secrets.randbelow(1_000_000):06d}"


def hash_email_code(code: str) -> str:
    """One-way hash of an email code for storage/comparison (never the raw
    code) - same precedent as hash_token_identifier."""
    import hashlib

    return hashlib.sha256(code.encode("utf-8")).hexdigest()