app/core/security.py

# app/core/security.py
"""
Authentication & authorization — production-grade v3.3

v3.3 changes:
- Centralized Redis client via redis_client.py (shared connection pool)
- Removed duplicated _get_redis() lazy-load logic from this module

v3.2:
- TTL-aware token blacklist (no memory leak)
- Redis-backed blacklist when REDIS_URL is set (multi-worker safe)
- O(1) refresh token lookup via refresh_token_id index
- bcrypt password hashing with 72-byte limit guard
"""

import hashlib
import hmac
import logging
import secrets
from datetime import datetime, timedelta, timezone
from threading import Lock
from typing import Dict, Optional, Tuple

import bcrypt
from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from sqlalchemy.orm import Session
from app.core.config import settings
from app.core.database import get_db

logger = logging.getLogger("hale.security")

# ──────────────────────────────────────────────
# Constants & Key Ring
# ──────────────────────────────────────────────

MAX_FAILED_LOGINS = 5
LOCKOUT_DURATION_MINUTES = 15
REFRESH_TOKEN_EXPIRE_DAYS = 30

KEYS_RING = {
    "key-v1": settings.SECRET_KEY
}
ACTIVE_KEY_ID = "key-v1"

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/login")

# ── Token blacklist (TTL-aware, no memory leak) ────────────────────────────
# Each entry: token -> expiry_timestamp (UTC)
# On each check, expired entries are pruned.
# In multi-worker deployments Redis is used automatically via redis_client.py.
from app.core.redis_client import get_redis as _get_redis

_blacklisted_tokens: Dict[str, datetime] = {}
_blacklist_lock = Lock()
_BLACKLIST_CLEANUP_EVERY = 100   # prune every N writes
_blacklist_write_count = 0

_REDIS_BLACKLIST_PREFIX = "hale:token:blacklist:"


def _prune_blacklist() -> None:
    """Remove expired tokens from the in-memory blacklist."""
    now = datetime.now(timezone.utc)
    expired = [t for t, exp in _blacklisted_tokens.items() if exp <= now]
    for t in expired:
        del _blacklisted_tokens[t]
    if expired:
        logger.debug("Pruned %d expired tokens from blacklist", len(expired))


# ──────────────────────────────────────────────
# Password hashing
# ──────────────────────────────────────────────

def hash_password(password: str) -> str:
    """Hash a plaintext password using bcrypt."""
    pwd_bytes = password.encode("utf-8")[:72]   # bcrypt 72-byte limit
    salt = bcrypt.gensalt(rounds=12)
    return bcrypt.hashpw(pwd_bytes, salt).decode("utf-8")


def verify_password(plain_password: str, hashed_password: str) -> bool:
    """Verify a plaintext password against a bcrypt hash."""
    try:
        return bcrypt.checkpw(
            plain_password.encode("utf-8")[:72],
            hashed_password.encode("utf-8"),
        )
    except Exception:
        return False


# ──────────────────────────────────────────────
# Access tokens (JWT)
# ──────────────────────────────────────────────

def create_access_token(
    data: dict,
    expires_delta: Optional[timedelta] = None,
) -> str:
    """Create a short-lived JWT access token signed with the active key."""
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + (
        expires_delta or timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
    )
    to_encode["exp"] = expire
    to_encode["type"] = "access"
    
    headers = {"kid": ACTIVE_KEY_ID}
    secret = KEYS_RING[ACTIVE_KEY_ID]
    return jwt.encode(to_encode, secret, algorithm=settings.ALGORITHM, headers=headers)


def decode_token(token: str) -> dict:
    """Decode and validate a JWT access token using dynamic key ring checks."""
    redis = _get_redis()

    if redis:
        # Redis-backed blacklist check (multi-worker safe)
        try:
            if redis.get(f"{_REDIS_BLACKLIST_PREFIX}{token}"):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Token has been revoked",
                    headers={"WWW-Authenticate": "Bearer"},
                )
        except HTTPException:
            raise
        except Exception as exc:
            logger.warning("Redis blacklist check failed, falling back to in-memory: %s", exc)
    else:
        # In-memory blacklist check with TTL eviction
        with _blacklist_lock:
            if token in _blacklisted_tokens:
                expiry = _blacklisted_tokens[token]
                if datetime.now(timezone.utc) < expiry:
                    raise HTTPException(
                        status_code=status.HTTP_401_UNAUTHORIZED,
                        detail="Token has been revoked",
                        headers={"WWW-Authenticate": "Bearer"},
                    )
                else:
                    # Token expired naturally — remove stale entry
                    del _blacklisted_tokens[token]

    try:
        # JWKS kid header check
        try:
            header = jwt.get_unverified_header(token)
            kid = header.get("kid", "key-v1")
            secret = KEYS_RING.get(kid, settings.SECRET_KEY)
        except Exception:
            secret = settings.SECRET_KEY

        payload = jwt.decode(token, secret, algorithms=[settings.ALGORITHM])
        if payload.get("type") != "access":
            raise JWTError("Not an access token")
        return payload
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        )


def blacklist_token(token: str) -> None:
    """Add a token to the blacklist until its natural expiry."""
    global _blacklist_write_count

    # Determine token expiry
    try:
        payload = jwt.decode(
            token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM],
            options={"verify_exp": False}
        )
        exp = payload.get("exp")
        expiry = datetime.fromtimestamp(exp, tz=timezone.utc) if exp else (
            datetime.now(timezone.utc) + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
        )
    except Exception:
        expiry = datetime.now(timezone.utc) + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)

    redis = _get_redis()
    if redis:
        ttl_seconds = max(1, int((expiry - datetime.now(timezone.utc)).total_seconds()))
        try:
            redis.setex(f"{_REDIS_BLACKLIST_PREFIX}{token}", ttl_seconds, "1")
            return
        except Exception as exc:
            logger.warning("Redis blacklist write failed, falling back to in-memory: %s", exc)

    # In-memory blacklist with periodic cleanup
    with _blacklist_lock:
        _blacklisted_tokens[token] = expiry
        _blacklist_write_count += 1
        if _blacklist_write_count >= _BLACKLIST_CLEANUP_EVERY:
            _prune_blacklist()
            _blacklist_write_count = 0


# ──────────────────────────────────────────────
# Refresh tokens & Version Cache
# ──────────────────────────────────────────────

def get_cached_versions(user_id: str, session_id: str) -> Tuple[Optional[int], Optional[int]]:
    """Retrieve user and session versions from Redis in a single pipeline."""
    redis = _get_redis()
    if not redis:
        return None, None
    try:
        pipe = redis.pipeline()
        pipe.get(f"hale:version:user:{user_id}")
        pipe.get(f"hale:version:session:{session_id}")
        u_ver, s_ver = pipe.execute()
        return (
            int(u_ver) if u_ver is not None else None,
            int(s_ver) if s_ver is not None else None
        )
    except Exception as exc:
        logger.error("Redis version cache read failed: %s", exc)
        return None, None


def set_cached_versions(user_id: str, session_id: str, user_ver: int, session_ver: int, ttl: int = 86400) -> None:
    """Write user and session versions to Redis in a single pipeline."""
    redis = _get_redis()
    if not redis:
        return
    try:
        pipe = redis.pipeline()
        pipe.setex(f"hale:version:user:{user_id}", ttl, str(user_ver))
        pipe.setex(f"hale:version:session:{session_id}", ttl, str(session_ver))
        pipe.execute()
    except Exception as exc:
        logger.error("Redis version cache write failed: %s", exc)


def clear_cached_versions(user_id: str, session_id: Optional[str] = None) -> None:
    """Invalidate version cache for a user and optionally a specific session."""
    redis = _get_redis()
    if not redis:
        return
    try:
        pipe = redis.pipeline()
        pipe.delete(f"hale:version:user:{user_id}")
        if session_id:
            pipe.delete(f"hale:version:session:{session_id}")
        pipe.execute()
    except Exception as exc:
        logger.error("Redis version cache invalidation failed: %s", exc)


def generate_selector_verifier() -> Tuple[str, str, str]:
    """
    Generate a cryptographically secure opaque refresh token.
    Returns (plain_token, selector, verifier_hash) — selector is used for B-Tree O(1) DB lookup.
    """
    selector = secrets.token_urlsafe(16)  # 22 base64url chars
    verifier = secrets.token_urlsafe(32)  # 43 base64url chars
    plain_token = f"{selector}.{verifier}"
    verifier_hash = hashlib.sha256(verifier.encode("utf-8")).hexdigest()
    return plain_token, selector, verifier_hash


def verify_verifier(verifier: str, stored_hash: str) -> bool:
    """Verify verifier matches the stored SHA-256 hash using constant-time comparison."""
    try:
        candidate = hashlib.sha256(verifier.encode("utf-8")).hexdigest()
        return hmac.compare_digest(candidate, stored_hash)
    except Exception:
        return False


def generate_refresh_token() -> Tuple[str, str]:
    """Generate a legacy opaque refresh token (backward-compatibility)."""
    token_id = secrets.token_urlsafe(16)
    plain    = secrets.token_urlsafe(64)
    return token_id, plain


def hash_refresh_token(token: str) -> str:
    """Hash a legacy refresh token (backward-compatibility)."""
    return hmac.new(
        key       = settings.SECRET_KEY.encode("utf-8"),
        msg       = token.encode("utf-8"),
        digestmod = hashlib.sha256,
    ).hexdigest()


def verify_refresh_token(plain_token: str, hashed_token: str) -> bool:
    """Verify a legacy refresh token (backward-compatibility)."""
    try:
        expected = hash_refresh_token(plain_token)
        return hmac.compare_digest(expected, hashed_token)
    except Exception:
        return False


def create_token_pair(user_id: str) -> dict:
    """Create a legacy token pair (backward-compatibility)."""
    access           = create_access_token(data={"sub": user_id})
    token_id, refresh = generate_refresh_token()
    return {
        "access_token":  access,
        "refresh_token": refresh,
        "token_id":      token_id,
        "token_type":    "bearer",
    }


def create_token_pair_with_session(user_id: str, session_version: int, user_version: int) -> dict:
    """Create a modern versioned token pair including selector and versions in JWT claims."""
    plain_token, selector, verifier_hash = generate_selector_verifier()
    access = create_access_token(data={
        "sub": user_id,
        "sid": selector,
        "sv": session_version,
        "uv": user_version
    })
    return {
        "access_token":  access,
        "refresh_token": plain_token,
        "token_type":    "bearer",
        "selector":      selector,
        "verifier_hash": verifier_hash,
    }


# ──────────────────────────────────────────────
# Account lockout
# ──────────────────────────────────────────────

def record_failed_login(db: Session, user) -> None:
    """Increment failed login counter; lock account if threshold reached."""
    user.failed_login_count = (user.failed_login_count or 0) + 1
    if user.failed_login_count >= MAX_FAILED_LOGINS:
        user.locked_until = datetime.now(timezone.utc) + timedelta(minutes=LOCKOUT_DURATION_MINUTES)
    db.commit()


def reset_failed_logins(db: Session, user) -> None:
    """Reset lockout state on successful login."""
    user.failed_login_count = 0
    user.locked_until = None
    user.last_login_at = datetime.now(timezone.utc)
    db.commit()


# ──────────────────────────────────────────────
# FastAPI dependency
# ──────────────────────────────────────────────

def get_current_user(
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
):
    """FastAPI dependency: extract and validate the current user from JWT (Zero-DB version cached)."""
    from app.models.user import User
    from app.models.session import UserSession

    payload = decode_token(token)
    user_id: Optional[str] = payload.get("sub")
    session_id: Optional[str] = payload.get("sid")
    jwt_sv: Optional[int] = payload.get("sv")
    jwt_uv: Optional[int] = payload.get("uv")

    if not user_id:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Token missing subject",
        )

    # 1. Sub-millisecond Redis version validation (highly scalable)
    redis = _get_redis()
    cache_ok = False
    if redis and session_id and jwt_sv is not None and jwt_uv is not None:
        cached_uv, cached_sv = get_cached_versions(user_id, session_id)
        if cached_uv is not None and cached_sv is not None:
            if jwt_uv != cached_uv or jwt_sv != cached_sv:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Token has been revoked",
                    headers={"WWW-Authenticate": "Bearer"},
                )
            cache_ok = True

    # 2. Database validation fallback (on cache miss or if cache was skipped)
    user = db.query(User).filter(User.user_id == user_id).first()
    if user is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="User not found",
        )
    if not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Account is deactivated",
        )

    # Validate user version
    if jwt_uv is not None and jwt_uv != user.user_version:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Token has been revoked",
            headers={"WWW-Authenticate": "Bearer"},
        )

    # Validate session version if sid (selector) is present
    if session_id:
        session = db.query(UserSession).filter(UserSession.selector == session_id).first()
        if session is None or session.is_compromised:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Session has been revoked",
                headers={"WWW-Authenticate": "Bearer"},
            )
        
        if jwt_sv is not None and jwt_sv != session.session_version:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Session version mismatch",
                headers={"WWW-Authenticate": "Bearer"},
            )

        # Sync back to cache on success so subsequent requests skip DB
        if not cache_ok and redis:
            set_cached_versions(user_id, session_id, user.user_version, session.session_version)

    return user


def get_current_token(token: str = Depends(oauth2_scheme)) -> str:
    """FastAPI dependency: return the raw access token string (for logout)."""
    return token