SupplyAgent/api/services/auth.py

from __future__ import annotations

import hashlib
import hmac
import logging
import secrets
from datetime import datetime, timedelta
from typing import Any

from sqlalchemy.exc import IntegrityError

from supply_infra.config import get_infra_settings
from supply_infra.db import get_session
from supply_infra.db.models.auth_user import AuthUser
from supply_infra.db.repositories.auth_repo import AuthRepository
from supply_infra.pipeline.dates import CHINA_TIMEZONE

logger = logging.getLogger(__name__)

SESSION_COOKIE_NAME = "supply_session"
ADMIN_ROLE = "admin"
USER_ROLE = "user"
ACTIVE_STATUS = "active"
DISABLED_STATUS = "disabled"
VALID_ROLES = {ADMIN_ROLE, USER_ROLE}
VALID_STATUSES = {ACTIVE_STATUS, DISABLED_STATUS}

_SCRYPT_N = 2**14
_SCRYPT_R = 8
_SCRYPT_P = 1
_DUMMY_PASSWORD_HASH: str | None = None


class InvalidCredentialsError(Exception):
    """Raised for every login rejection to avoid exposing account state."""


class DuplicateUsernameError(Exception):
    """Raised when a local username already exists."""


def _now() -> datetime:
    return datetime.now(CHINA_TIMEZONE).replace(tzinfo=None)


def normalize_username(username: str) -> str:
    return username.strip().lower()


def hash_password(password: str) -> str:
    salt = secrets.token_bytes(16)
    derived = hashlib.scrypt(
        password.encode("utf-8"),
        salt=salt,
        n=_SCRYPT_N,
        r=_SCRYPT_R,
        p=_SCRYPT_P,
        dklen=32,
    )
    return (
        f"scrypt${_SCRYPT_N}${_SCRYPT_R}${_SCRYPT_P}$"
        f"{salt.hex()}${derived.hex()}"
    )


def verify_password(password: str, encoded: str) -> bool:
    try:
        algorithm, n, r, p, salt_hex, expected_hex = encoded.split("$", 5)
        if algorithm != "scrypt":
            return False
        derived = hashlib.scrypt(
            password.encode("utf-8"),
            salt=bytes.fromhex(salt_hex),
            n=int(n),
            r=int(r),
            p=int(p),
            dklen=len(bytes.fromhex(expected_hex)),
        )
        return hmac.compare_digest(derived.hex(), expected_hex)
    except (TypeError, ValueError):
        return False


def _dummy_password_hash() -> str:
    global _DUMMY_PASSWORD_HASH
    if _DUMMY_PASSWORD_HASH is None:
        _DUMMY_PASSWORD_HASH = hash_password("supply-agent-invalid-password")
    return _DUMMY_PASSWORD_HASH


def hash_session_token(token: str) -> str:
    return hashlib.sha256(token.encode("utf-8")).hexdigest()


def serialize_user(user: AuthUser) -> dict[str, Any]:
    return {
        "id": user.id,
        "username": user.username,
        "display_name": user.display_name,
        "role": user.role,
        "status": user.status,
        "last_login_at": user.last_login_at.isoformat() if user.last_login_at else None,
        "created_at": user.created_at.isoformat() if user.created_at else None,
        "updated_at": user.updated_at.isoformat() if user.updated_at else None,
    }


def authenticate(
    *,
    username: str,
    password: str,
    ip_address: str | None,
    user_agent: str | None,
) -> tuple[dict[str, Any], str]:
    now = _now()
    normalized = normalize_username(username)
    settings = get_infra_settings()

    with get_session() as db:
        repo = AuthRepository(db)
        repo.delete_expired_sessions(now)
        user = repo.get_user_by_username(normalized)

        if user is None:
            verify_password(password, _dummy_password_hash())
            raise InvalidCredentialsError

        if user.locked_until and user.locked_until > now:
            verify_password(password, user.password_hash)
            raise InvalidCredentialsError

        if user.locked_until and user.locked_until <= now:
            user.locked_until = None
            user.failed_login_count = 0

        password_valid = verify_password(password, user.password_hash)
        if not password_valid or user.status != ACTIVE_STATUS:
            if password_valid:
                raise InvalidCredentialsError
            user.failed_login_count += 1
            if user.failed_login_count >= 5:
                user.failed_login_count = 0
                user.locked_until = now + timedelta(minutes=15)
            raise InvalidCredentialsError

        user.failed_login_count = 0
        user.locked_until = None
        user.last_login_at = now
        token = secrets.token_urlsafe(32)
        repo.create_session(
            token_hash=hash_session_token(token),
            user_id=user.id,
            expires_at=now + timedelta(hours=settings.auth_session_hours),
            ip_address=(ip_address or "")[:64] or None,
            user_agent=(user_agent or "")[:512] or None,
        )
        db.flush()
        return serialize_user(user), token


def resolve_session(token: str) -> dict[str, Any] | None:
    now = _now()
    with get_session() as db:
        repo = AuthRepository(db)
        user = repo.get_active_session(
            token_hash=hash_session_token(token),
            now=now,
        )
        if user is None:
            return None
        return serialize_user(user)


def revoke_session(token: str) -> None:
    with get_session() as db:
        AuthRepository(db).delete_session_by_hash(hash_session_token(token))


def list_users() -> list[dict[str, Any]]:
    with get_session() as db:
        return [serialize_user(user) for user in AuthRepository(db).list_users()]


def create_user(
    *,
    username: str,
    password: str,
    display_name: str,
    role: str,
) -> dict[str, Any]:
    if role not in VALID_ROLES:
        raise ValueError("invalid role")
    try:
        with get_session() as db:
            user = AuthRepository(db).create_user(
                username=normalize_username(username),
                password_hash=hash_password(password),
                display_name=display_name.strip(),
                role=role,
            )
            db.flush()
            return serialize_user(user)
    except IntegrityError as exc:
        raise DuplicateUsernameError from exc


def update_user(
    user_id: int,
    *,
    display_name: str | None = None,
    role: str | None = None,
    status: str | None = None,
    password: str | None = None,
) -> dict[str, Any] | None:
    if role is not None and role not in VALID_ROLES:
        raise ValueError("invalid role")
    if status is not None and status not in VALID_STATUSES:
        raise ValueError("invalid status")

    with get_session() as db:
        repo = AuthRepository(db)
        user = repo.get_user(user_id)
        if user is None:
            return None
        revoke_existing_sessions = False
        if display_name is not None:
            user.display_name = display_name.strip()
        if role is not None and role != user.role:
            user.role = role
            revoke_existing_sessions = True
        if status is not None and status != user.status:
            user.status = status
            revoke_existing_sessions = True
        if password is not None:
            user.password_hash = hash_password(password)
            user.failed_login_count = 0
            user.locked_until = None
            revoke_existing_sessions = True
        if revoke_existing_sessions:
            repo.delete_user_sessions(user.id)
        db.flush()
        return serialize_user(user)


def delete_user(user_id: int) -> bool:
    with get_session() as db:
        repo = AuthRepository(db)
        user = repo.get_user(user_id)
        if user is None:
            return False
        repo.delete_user_sessions(user.id)
        repo.delete_user(user)
        return True


def ensure_bootstrap_admin() -> None:
    settings = get_infra_settings()
    username = normalize_username(settings.auth_bootstrap_admin_username)
    password = settings.auth_bootstrap_admin_password.get_secret_value()
    if not username and not password:
        return
    if not username or not password:
        raise RuntimeError(
            "AUTH_BOOTSTRAP_ADMIN_USERNAME and AUTH_BOOTSTRAP_ADMIN_PASSWORD "
            "must be configured together"
        )
    if len(password) < 8:
        raise RuntimeError("AUTH_BOOTSTRAP_ADMIN_PASSWORD must contain at least 8 characters")

    with get_session() as db:
        repo = AuthRepository(db)
        if repo.get_user_by_username(username) is not None:
            return
        repo.create_user(
            username=username,
            password_hash=hash_password(password),
            display_name=settings.auth_bootstrap_admin_display_name.strip() or "系统管理员",
            role=ADMIN_ROLE,
        )
        logger.info("Created bootstrap administrator: username=%s", username)