backend/app/core/security.py

"""
Security utilities: password hashing, JWT tokens, refresh session management.
Access tokens: 15 minutes (short-lived, in-memory on frontend).
Refresh tokens: 30 days, stored server-side with revocation support.
"""
from datetime import datetime, timedelta, timezone
import uuid

import jwt
from pwdlib import PasswordHash
from pwdlib.hashers.argon2 import Argon2Hasher
from pwdlib.hashers.bcrypt import BcryptHasher
from sqlalchemy import Boolean, DateTime, String, func, select
from sqlalchemy.dialects.postgresql import UUID as PGUUID
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Mapped, mapped_column

from app.core.config import settings
from app.core.database import Base

password_hash = PasswordHash((Argon2Hasher(), BcryptHasher()))
ALGORITHM = "HS256"


# ===== Refresh Session Model =====

class RefreshSession(Base):
    """Server-side refresh token sessions for revocation support."""
    __tablename__ = "refresh_sessions"

    id: Mapped[uuid.UUID] = mapped_column(PGUUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    user_id: Mapped[uuid.UUID] = mapped_column(PGUUID(as_uuid=True), nullable=False, index=True)
    token_hash: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    user_agent: Mapped[str | None] = mapped_column(String(512), nullable=True)
    ip_address: Mapped[str | None] = mapped_column(String(45), nullable=True)


# ===== Token Creation =====

def create_access_token(subject: str, expires_delta: timedelta | None = None) -> str:
    """Short-lived access token (15 min default)."""
    expire = datetime.now(timezone.utc) + (expires_delta or timedelta(minutes=15))
    return jwt.encode(
        {"exp": expire, "sub": str(subject), "type": "access", "jti": uuid.uuid4().hex},
        settings.SECRET_KEY, algorithm=ALGORITHM
    )


def create_refresh_token(subject: str) -> str:
    """Long-lived refresh token (30 days). Must be stored server-side."""
    expire = datetime.now(timezone.utc) + timedelta(days=settings.REFRESH_TOKEN_EXPIRE_DAYS)
    return jwt.encode(
        {"exp": expire, "sub": str(subject), "type": "refresh", "jti": uuid.uuid4().hex},
        settings.SECRET_KEY, algorithm=ALGORITHM
    )


# ===== Token Verification =====

def verify_token(token: str, token_type: str = "access") -> dict | None:
    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[ALGORITHM])
        if payload.get("type") != token_type:
            return None
        return payload
    except jwt.PyJWTError:
        return None


# ===== Refresh Session Management =====

import hashlib

def _hash_token(token: str) -> str:
    """Hash refresh token for storage (never store raw)."""
    return hashlib.sha256(token.encode()).hexdigest()


async def create_refresh_session(
    db: AsyncSession, user_id: uuid.UUID, refresh_token: str,
    user_agent: str | None = None, ip_address: str | None = None
) -> RefreshSession:
    """Store refresh session server-side."""
    session = RefreshSession(
        user_id=user_id,
        token_hash=_hash_token(refresh_token),
        expires_at=datetime.now(timezone.utc) + timedelta(days=settings.REFRESH_TOKEN_EXPIRE_DAYS),
        user_agent=user_agent,
        ip_address=ip_address,
    )
    db.add(session)
    return session


async def validate_refresh_session(db: AsyncSession, refresh_token: str) -> RefreshSession | None:
    """Validate refresh token against server-side session. Returns None if revoked/expired."""
    token_hash = _hash_token(refresh_token)
    result = await db.execute(
        select(RefreshSession).where(
            RefreshSession.token_hash == token_hash,
            RefreshSession.revoked_at.is_(None),
            RefreshSession.expires_at > datetime.now(timezone.utc),
        )
    )
    return result.scalar_one_or_none()


async def revoke_refresh_session(db: AsyncSession, refresh_token: str) -> bool:
    """Revoke a specific refresh session (logout)."""
    token_hash = _hash_token(refresh_token)
    result = await db.execute(
        select(RefreshSession).where(RefreshSession.token_hash == token_hash)
    )
    session = result.scalar_one_or_none()
    if session:
        session.revoked_at = datetime.now(timezone.utc)
        return True
    return False


async def revoke_all_user_sessions(db: AsyncSession, user_id: uuid.UUID) -> int:
    """Revoke all refresh sessions for a user (logout everywhere)."""
    from sqlalchemy import update
    result = await db.execute(
        update(RefreshSession)
        .where(RefreshSession.user_id == user_id, RefreshSession.revoked_at.is_(None))
        .values(revoked_at=datetime.now(timezone.utc))
    )
    return result.rowcount


# ===== Password =====

def hash_password(password: str) -> str:
    return password_hash.hash(password)


def verify_password(plain_password: str, hashed_password: str) -> tuple[bool, str | None]:
    return password_hash.verify_and_update(plain_password, hashed_password)