vacuum-wall/lib/auth.py

"""JWT authentication module for Vacuum Wall.

Handles token creation, validation, refresh, and blacklisting.
Configuration comes from config/auth/config.json.
"""

from __future__ import annotations

import logging
import secrets
import time
import uuid
from pathlib import Path
from typing import Any

import jwt

from lib.common import load_json, save_json
from lib.db import (
    Q_DELETE_EXPIRED_BLACKLIST,
    Q_DELETE_REFRESH_TOKEN,
    Q_INSERT_BLACKLIST,
    Q_SELECT_BLACKLIST,
    Q_SELECT_REFRESH_TOKEN,
    Q_UPSERT_REFRESH_TOKEN,
    get_db,
)

logger = logging.getLogger(__name__)

AUTH_CONFIG_PATH = (
    Path(__file__).resolve().parent.parent / "config" / "auth" / "config.json"
)

_DEFAULT_JWT_CONFIG = {
    "access_token_ttl": 900,
    "refresh_token_ttl": 604800,
    "algorithm": "HS256",
    "secret": "",
}


def _get_jwt_config() -> dict[str, Any]:
    """Load JWT configuration from auth config."""
    raw = load_json(AUTH_CONFIG_PATH)
    return raw.get("jwt", _DEFAULT_JWT_CONFIG)


def get_secret() -> str | None:
    """Return the JWT signing secret, or ``None`` if not configured."""
    secret = _get_jwt_config().get("secret")
    return secret if secret else None


def get_access_ttl() -> int:
    """Return access token TTL in seconds."""
    return _get_jwt_config().get("access_token_ttl", 900)


def get_refresh_ttl() -> int:
    """Return refresh token TTL in seconds."""
    return _get_jwt_config().get("refresh_token_ttl", 604800)


def get_algorithm() -> str:
    """Return the JWT algorithm."""
    return _get_jwt_config().get("algorithm", "HS256")


def generate_access_token(username: str, permissions: dict[str, str]) -> str:
    """Generate a new access token for *username*.

    Args:
        username: The authenticated username.
        permissions: Dict mapping subsystem names to permission levels.

    Returns:
        JWT token string.

    Raises:
        RuntimeError: If JWT secret is not configured.
    """
    secret = get_secret()
    if not secret:
        raise RuntimeError("JWT secret is not configured")
    algorithm = get_algorithm()
    now = int(time.time())
    payload = {
        "sub": username,
        "exp": now + get_access_ttl(),
        "iat": now,
        "jti": str(uuid.uuid4()),
        "type": "access",
        "permissions": permissions,
    }
    return jwt.encode(payload, secret, algorithm=algorithm)


def generate_refresh_token(username: str) -> str:
    """Generate a new refresh token for *username*.

    Args:
        username: The authenticated username.

    Returns:
        JWT refresh token string.

    Raises:
        RuntimeError: If JWT secret is not configured.
    """
    secret = get_secret()
    if not secret:
        raise RuntimeError("JWT secret is not configured")
    algorithm = get_algorithm()
    now = int(time.time())
    payload = {
        "sub": username,
        "exp": now + get_refresh_ttl(),
        "iat": now,
        "jti": str(uuid.uuid4()),
        "type": "refresh",
    }
    return jwt.encode(payload, secret, algorithm=algorithm)


def generate_tokens(username: str, permissions: dict[str, str]) -> dict[str, str]:
    """Generate both access and refresh tokens.

    Args:
        username: The authenticated username.
        permissions: Dict mapping subsystem names to permission levels.

    Returns:
        Dict with ``access_token`` and ``refresh_token`` keys.
    """
    access_token = generate_access_token(username, permissions)
    refresh_token = generate_refresh_token(username)
    _persist_refresh_token(username, refresh_token)
    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
    }


def _persist_refresh_token(username: str, refresh_token: str) -> None:
    """Persist the active refresh token JTI for *username* in the database.

    One row per user — replaces any existing entry on upsert.

    Args:
        username: The username.
        refresh_token: The JWT refresh token string.
    """
    payload = decode_token(refresh_token)
    if payload is None:
        return
    jti = payload.get("jti")
    if not jti:
        return
    issued_at = payload.get("iat", int(time.time()))
    db = get_db()
    db.run(Q_UPSERT_REFRESH_TOKEN, (username, jti, issued_at))


def blacklist_active_refresh_token(username: str) -> None:
    """Blacklist the user's currently active refresh token from the database.

    Looks up the stored JTI for *username*, blacklists it, and removes the
    database entry. Safe to call when no token is registered — the query
    will simply return no rows.

    Args:
        username: The username.
    """
    db = get_db()
    rows = db.query(Q_SELECT_REFRESH_TOKEN, (username,))
    if rows:
        blacklist_token(rows[0]["jti"], token_type="refresh")
    db.run(Q_DELETE_REFRESH_TOKEN, (username,))


def clear_active_refresh_token(username: str) -> None:
    """Remove the user's stored refresh token entry without blacklisting.

    Used when the refresh token has already been blacklisted (e.g., during
    a successful refresh rotation).

    Args:
        username: The username.
    """
    db = get_db()
    db.run(Q_DELETE_REFRESH_TOKEN, (username,))


def decode_token(token_string: str) -> dict[str, Any] | None:
    """Decode and validate a JWT token.

    Args:
        token_string: The JWT token string (without Bearer prefix).

    Returns:
        Payload dict if valid, None if invalid/expired or secret not configured.
    """
    secret = get_secret()
    if not secret:
        return None
    algorithm = get_algorithm()
    try:
        payload = jwt.decode(token_string, secret, algorithms=[algorithm])
        return payload
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
        return None


def validate_token(
    token_string: str, token_type: str = "access"
) -> dict[str, Any] | None:
    """Validate a JWT token and check it against the blacklist.

    Args:
        token_string: The JWT token string.
        token_type: Expected token type ("access" or "refresh").

    Returns:
        Payload dict including permissions, or None if invalid/blacklisted.
    """
    payload = decode_token(token_string)
    if payload is None:
        return None
    if payload.get("type") != token_type:
        return None

    jti = payload.get("jti")
    if jti and is_blacklisted(jti):
        return None

    return payload


def blacklist_token(jti: str, token_type: str = "access") -> None:
    """Add a JTI to the blacklist.

    Args:
        jti: Token UUID to blacklist.
        token_type: Token type ("access" or "refresh"). Defaults to "access".
            The blacklist expiry is set to current time + the token type's TTL.
    """
    db = get_db()
    ttl = get_refresh_ttl() if token_type == "refresh" else get_access_ttl()
    db.run(Q_INSERT_BLACKLIST, (jti, token_type, int(time.time()) + ttl))


def is_blacklisted(jti: str) -> bool:
    """Check if a JTI is blacklisted.

    Args:
        jti: Token UUID to check.

    Returns:
        True if the token has been blacklisted.
    """
    db = get_db()
    rows = db.query(Q_SELECT_BLACKLIST, (jti,))
    return len(rows) > 0


def blacklist_expired() -> None:
    """Remove expired entries from the blacklist."""
    db = get_db()
    now = int(time.time())
    db.run(Q_DELETE_EXPIRED_BLACKLIST, (now,))


def rotate_secret() -> None:
    """Rotate the JWT secret, invalidating all existing tokens.

    Used when a user's password is changed to ensure all prior sessions
    are immediately terminated regardless of token expiration.
    """
    raw = load_json(AUTH_CONFIG_PATH)
    jwt_config = raw.get("jwt", _DEFAULT_JWT_CONFIG)
    jwt_config["secret"] = secrets.token_urlsafe(48)
    raw["jwt"] = jwt_config
    save_json(AUTH_CONFIG_PATH, raw)
    logger.warning("JWT secret rotated — all existing tokens are now invalid")


class RateLimiter:
    """Simple sliding-window rate limiter for login attempts.

    Maintains per-key attempt timestamps and rejects requests that exceed
    the allowed count within the window.
    """

    def __init__(self, max_attempts: int = 5, window_seconds: int = 300) -> None:
        self.max_attempts = max_attempts
        self.window = window_seconds
        self.attempts: dict[str, list[float]] = {}

    def is_allowed(self, key: str) -> bool:
        """Check if a request from *key* is allowed.

        Args:
            key: Identifier for the rate limit bucket (e.g., username or IP).

        Returns:
            True if the request is allowed, False if rate limited.
        """
        now = time.time()
        cutoff = now - self.window
        timestamps = self.attempts.get(key, [])

        # Clean old entries
        self.attempts[key] = [t for t in timestamps if t > cutoff]

        if len(self.attempts[key]) >= self.max_attempts:
            return False

        self.attempts[key].append(now)
        return True

    def cleanup(self) -> None:
        """Remove expired entries from all buckets."""
        now = time.time()
        cutoff = now - self.window
        for key in list(self.attempts):
            self.attempts[key] = [t for t in self.attempts[key] if t > cutoff]
            if not self.attempts[key]:
                del self.attempts[key]


# Global rate limiters
_login_limiter = RateLimiter(max_attempts=10, window_seconds=300)
_webauthn_limiter = RateLimiter(max_attempts=5, window_seconds=600)


def check_login_rate(username: str) -> bool:
    """Check if login is rate-limited for the given username.

    Args:
        username: The login attempt username.

    Returns:
        True if the attempt is allowed, False if rate limited.
    """
    return _login_limiter.is_allowed(username)


def check_webauthn_rate(username: str) -> bool:
    """Check if WebAuthn authentication is rate-limited for the given username.

    Args:
        username: The WebAuthn attempt username.

    Returns:
        True if the attempt is allowed, False if rate limited.
    """
    return _webauthn_limiter.is_allowed(username)