ForgeFlow-ERP/backend/app/services/auth.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Annotated, Any

from fastapi import Depends, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.db.session import get_db
from app.models.org import Department, Employee, Permission, Role, RolePermission, User, UserRole
from app.schemas.database import AuthSession, LoginUser

PBKDF2_ITERATIONS = 600000

ROLE_TITLE_MAP = {
    "ADMIN": "总控驾驶舱",
    "SALES": "订单与客户协同",
    "PLANNER": "MRP与排产控制",
    "PURCHASER": "采购执行中心",
    "WAREHOUSE": "仓储作业中心",
    "SUPERVISOR": "车间调度中心",
    "OPERATOR": "现场报工中心",
    "FINANCE": "成本与财务中心",
    "EQUIP_ADMIN": "设备维护中心",
}


@dataclass
class AuthContext:
    user: User
    employee: Employee
    department: Department | None
    role_codes: list[str]
    role_names: list[str]
    permission_codes: list[str]


bearer_scheme = HTTPBearer(auto_error=False, description="登录接口返回的 Bearer Token")


def hash_password(password: str, salt: str | None = None) -> str:
    salt_bytes = os.urandom(16) if salt is None else salt.encode("utf-8")
    digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt_bytes, PBKDF2_ITERATIONS)
    return (
        f"pbkdf2_sha256${PBKDF2_ITERATIONS}$"
        f"{base64.urlsafe_b64encode(salt_bytes).decode('utf-8')}$"
        f"{base64.urlsafe_b64encode(digest).decode('utf-8')}"
    )


def verify_password(password: str, stored_hash: str) -> bool:
    try:
        algorithm, iterations, salt_b64, digest_b64 = stored_hash.split("$", 3)
    except ValueError:
        return False

    if algorithm != "pbkdf2_sha256":
        return False

    salt = base64.urlsafe_b64decode(salt_b64.encode("utf-8"))
    expected = base64.urlsafe_b64decode(digest_b64.encode("utf-8"))
    actual = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, int(iterations))
    return hmac.compare_digest(actual, expected)


def get_auth_context(db: Session, username: str) -> AuthContext:
    stmt = (
        select(User, Employee, Department)
        .join(Employee, Employee.id == User.employee_id)
        .outerjoin(Department, Department.id == User.dept_id)
        .where(User.username == username, User.status == "ACTIVE")
    )
    row = db.execute(stmt).first()
    if not row:
        raise HTTPException(status_code=401, detail="用户名或密码错误")

    user, employee, department = row

    role_rows = db.execute(
        select(Role.role_code, Role.role_name)
        .join(UserRole, UserRole.role_id == Role.id)
        .where(UserRole.user_id == user.id, Role.status == "ACTIVE")
        .order_by(Role.id)
    ).all()
    role_codes = [row[0] for row in role_rows]
    role_names = [row[1] for row in role_rows]

    permission_rows = db.execute(
        select(Permission.permission_code)
        .join(RolePermission, RolePermission.permission_id == Permission.id)
        .join(UserRole, UserRole.role_id == RolePermission.role_id)
        .where(
            UserRole.user_id == user.id,
            Permission.status == "ACTIVE",
        )
        .group_by(Permission.permission_code)
        .order_by(Permission.permission_code)
    ).all()
    permission_codes = [row[0] for row in permission_rows]
    if user.is_super_admin or "ADMIN" in role_codes:
        for permission_code in ("MENU_SYSTEM_EXTENSION", "MENU_SYSTEM_PERMISSION"):
            if permission_code not in permission_codes:
                permission_codes.append(permission_code)

    return AuthContext(
        user=user,
        employee=employee,
        department=department,
        role_codes=role_codes,
        role_names=role_names,
        permission_codes=permission_codes,
    )


def login_with_password(db: Session, username: str, password: str) -> AuthSession:
    context = get_auth_context(db, username)
    if not verify_password(password, context.user.password_hash):
        raise HTTPException(status_code=401, detail="用户名或密码错误")
    context.user.last_login_at = datetime.now(timezone.utc).replace(tzinfo=None)
    db.add(context.user)
    db.commit()
    return build_auth_session(context)


def build_login_user(context: AuthContext) -> LoginUser:
    title = ROLE_TITLE_MAP.get(context.role_codes[0], "工业风管理台") if context.role_codes else "工业风管理台"
    return LoginUser(
        user_id=context.user.id,
        username=context.user.username,
        employee_id=context.employee.id,
        mobile=context.employee.mobile,
        display_name=context.user.nickname or context.employee.employee_name,
        title=title,
        dept_name=context.department.dept_name if context.department else None,
        role_codes=context.role_codes,
        role_names=context.role_names,
        permission_codes=context.permission_codes,
    )


def build_auth_session(context: AuthContext) -> AuthSession:
    access_token, expires_at = create_access_token(context)
    return AuthSession(
        access_token=access_token,
        token_type="bearer",
        expires_at=expires_at,
        **build_login_user(context).model_dump(),
    )


def next_employee_code(db: Session) -> str:
    total = db.scalar(select(func.count(Employee.id))) or 0
    return f"EMP-{int(total) + 1:04d}"


def create_access_token(context: AuthContext) -> tuple[str, datetime]:
    settings = get_settings()
    now = datetime.now(timezone.utc)
    expires_at = now + timedelta(minutes=settings.auth_token_expire_minutes)
    header = {"alg": "HS256", "typ": "JWT"}
    payload = {
        "sub": context.user.username,
        "uid": context.user.id,
        "iat": int(now.timestamp()),
        "exp": int(expires_at.timestamp()),
    }
    encoded_header = _urlsafe_b64encode(json.dumps(header, separators=(",", ":")).encode("utf-8"))
    encoded_payload = _urlsafe_b64encode(json.dumps(payload, separators=(",", ":")).encode("utf-8"))
    signing_input = f"{encoded_header}.{encoded_payload}"
    signature = _sign_token(signing_input, settings.auth_secret_key)
    return f"{signing_input}.{signature}", expires_at


def decode_access_token(token: str) -> dict[str, Any]:
    settings = get_settings()
    try:
        encoded_header, encoded_payload, signature = token.split(".", 2)
    except ValueError as exc:
        raise HTTPException(status_code=401, detail="登录凭证无效") from exc

    signing_input = f"{encoded_header}.{encoded_payload}"
    expected_signature = _sign_token(signing_input, settings.auth_secret_key)
    if not hmac.compare_digest(expected_signature, signature):
        raise HTTPException(status_code=401, detail="登录凭证无效")

    try:
        payload = json.loads(_urlsafe_b64decode(encoded_payload).decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc:
        raise HTTPException(status_code=401, detail="登录凭证无效") from exc

    expires_at = payload.get("exp")
    username = payload.get("sub")
    if not isinstance(expires_at, int) or not username:
        raise HTTPException(status_code=401, detail="登录凭证无效")
    if expires_at <= int(datetime.now(timezone.utc).timestamp()):
        raise HTTPException(status_code=401, detail="登录已过期，请重新登录")

    return payload


def get_current_auth_context(
    credentials: Annotated[HTTPAuthorizationCredentials | None, Security(bearer_scheme)],
    db: Session = Depends(get_db),
) -> AuthContext:
    if credentials is None or credentials.scheme.lower() != "bearer":
        raise HTTPException(status_code=401, detail="未登录或登录已过期")

    payload = decode_access_token(credentials.credentials)
    return get_auth_context(db, str(payload["sub"]))


def require_authenticated_user(
    context: Annotated[AuthContext, Depends(get_current_auth_context)],
) -> AuthContext:
    return context


def require_system_admin(
    context: Annotated[AuthContext, Depends(get_current_auth_context)],
) -> AuthContext:
    if context.user.is_super_admin or "ADMIN" in context.role_codes:
        return context
    raise HTTPException(status_code=403, detail="仅系统管理员可访问系统管理")


def _sign_token(signing_input: str, secret_key: str) -> str:
    digest = hmac.new(secret_key.encode("utf-8"), signing_input.encode("utf-8"), hashlib.sha256).digest()
    return _urlsafe_b64encode(digest)


def _urlsafe_b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")


def _urlsafe_b64decode(data: str) -> bytes:
    padding = "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(f"{data}{padding}".encode("utf-8"))