DataOps-platform/app/core/system/auth.py

"""PostgreSQL-backed authentication using Argon2id and short-lived bearer tokens."""

from __future__ import annotations

import logging
import threading
import time
from collections import defaultdict, deque
from typing import Any

from argon2 import PasswordHasher
from argon2.exceptions import InvalidHashError, VerificationError, VerifyMismatchError
from flask import current_app
from sqlalchemy import text

from app import db
from app.core.system.tokens import TokenError, decode_access_token, issue_access_token


logger = logging.getLogger(__name__)
_password_hasher = PasswordHasher()
_failure_lock = threading.Lock()
_login_failures: dict[str, deque[float]] = defaultdict(deque)
_FAILURE_WINDOW_SECONDS = 300
_FAILURE_LIMIT = 5


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


def verify_password(encoded: str, password: str) -> bool:
    try:
        return bool(_password_hasher.verify(encoded, password))
    except (VerifyMismatchError, VerificationError, InvalidHashError):
        return False


def _rate_key(username: str, ip_address: str | None) -> str:
    return f"{username.strip().lower()}|{ip_address or '-'}"


def _check_rate_limit(key: str, *, now: float | None = None) -> None:
    now = now or time.time()
    with _failure_lock:
        failures = _login_failures[key]
        while failures and failures[0] <= now - _FAILURE_WINDOW_SECONDS:
            failures.popleft()
        if len(failures) >= _FAILURE_LIMIT:
            raise PermissionError("too many login attempts")


def _record_failure(key: str, *, now: float | None = None) -> None:
    with _failure_lock:
        _login_failures[key].append(now or time.time())


def _clear_failures(key: str) -> None:
    with _failure_lock:
        _login_failures.pop(key, None)


def _audit(
    *,
    user_id: str | None,
    username: str,
    event_type: str,
    success: bool,
    ip_address: str | None,
    user_agent: str | None,
    detail: str | None = None,
) -> None:
    db.session.execute(
        text(
            "INSERT INTO public.auth_audit_events "
            "(user_id, username, event_type, success, ip_address, user_agent, detail) "
            "VALUES (CAST(:user_id AS uuid), :username, :event_type, :success, "
            ":ip_address, :user_agent, :detail)"
        ),
        {
            "user_id": user_id,
            "username": username[:64],
            "event_type": event_type,
            "success": success,
            "ip_address": (ip_address or "")[:64] or None,
            "user_agent": (user_agent or "")[:300] or None,
            "detail": (detail or "")[:500] or None,
        },
    )


def authenticate_user(
    username: str,
    password: str,
    *,
    ip_address: str | None = None,
    user_agent: str | None = None,
) -> dict[str, Any] | None:
    key = _rate_key(username, ip_address)
    _check_rate_limit(key)
    row = db.session.execute(
        text(
            "SELECT u.id::text, u.username, u.display_name, u.password_hash, "
            "u.status, COALESCE(array_agg(r.name ORDER BY r.name) "
            "FILTER (WHERE r.name IS NOT NULL), ARRAY[]::varchar[]) AS roles "
            "FROM public.users u "
            "LEFT JOIN public.user_roles ur ON ur.user_id = u.id "
            "LEFT JOIN public.roles r ON r.id = ur.role_id "
            "WHERE lower(u.username) = lower(:username) "
            "GROUP BY u.id"
        ),
        {"username": username.strip()},
    ).one_or_none()

    valid = bool(row and row[4] == "active" and verify_password(row[3], password))
    if not valid:
        _record_failure(key)
        _audit(
            user_id=row[0] if row else None,
            username=username,
            event_type="login",
            success=False,
            ip_address=ip_address,
            user_agent=user_agent,
            detail="invalid credentials or disabled account",
        )
        db.session.commit()
        return None

    db.session.execute(
        text("UPDATE public.users SET last_login_at = CURRENT_TIMESTAMP WHERE id = CAST(:id AS uuid)"),
        {"id": row[0]},
    )
    _audit(
        user_id=row[0],
        username=row[1],
        event_type="login",
        success=True,
        ip_address=ip_address,
        user_agent=user_agent,
    )
    db.session.commit()
    _clear_failures(key)
    roles = list(row[5])
    return {
        "id": row[0],
        "username": row[1],
        "display_name": row[2],
        "roles": roles,
        "token": issue_access_token(
            user_id=row[0], roles=roles, secret=current_app.config["SECRET_KEY"]
        ),
    }


def load_identity(user_id: str) -> dict[str, Any] | None:
    row = db.session.execute(
        text(
            "SELECT u.id::text, u.username, u.display_name, u.status, "
            "COALESCE(array_agg(r.name ORDER BY r.name) FILTER "
            "(WHERE r.name IS NOT NULL), ARRAY[]::varchar[]) "
            "FROM public.users u "
            "LEFT JOIN public.user_roles ur ON ur.user_id = u.id "
            "LEFT JOIN public.roles r ON r.id = ur.role_id "
            "WHERE u.id = CAST(:id AS uuid) GROUP BY u.id"
        ),
        {"id": user_id},
    ).one_or_none()
    if not row or row[3] != "active":
        return None
    return {"id": row[0], "username": row[1], "display_name": row[2], "roles": list(row[4])}


def load_identity_from_token(token: str, *, secret: str) -> dict[str, Any] | None:
    try:
        claims = decode_access_token(token, secret=secret)
    except TokenError:
        return None
    return load_identity(claims["sub"])


# Compatibility aliases retained only for imports during the route migration.
def login_user(username: str, password: str):
    result = authenticate_user(username, password)
    return (True, result) if result else (False, "用户名或密码错误")


def get_user_by_username(username: str):
    row = db.session.execute(
        text("SELECT id::text FROM public.users WHERE lower(username) = lower(:username)"),
        {"username": username},
    ).scalar_one_or_none()
    return load_identity(row) if row else None


def init_db() -> bool:
    logger.warning("init_db is deprecated; use Alembic migrations")
    return True


def require_auth(view):
    from app.core.system.permissions import require_permissions, READ_GOVERNANCE

    return require_permissions(READ_GOVERNANCE)(view)