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"))