taiji-AI-PAD/services/mcp-server/app/auth.py

"""Authentication helpers (API Key + JWT) for MCP server."""

from __future__ import annotations

import hashlib
import os
from datetime import datetime, timedelta
from typing import Any, Dict, Optional

from fastapi import Depends, HTTPException, Request, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwt
from passlib.context import CryptContext
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from config import settings
from database import get_db
from models import APIKey, User

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
bearer_scheme = HTTPBearer(auto_error=False)


def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)


def get_password_hash(password: str) -> str:
    return pwd_context.hash(password)


def create_access_token(data: Dict[str, Any], expires_delta: Optional[timedelta] = None) -> str:
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=settings.jwt_expire_minutes))
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, settings.secret_key, algorithm=settings.jwt_algorithm)
    return encoded_jwt


async def _get_user_by_email(email: str, db: AsyncSession) -> Optional[User]:
    result = await db.execute(select(User).where(User.email == email))
    return result.scalar_one_or_none()


async def authenticate_user(email: str, password: str, db: AsyncSession) -> Optional[User]:
    user = await _get_user_by_email(email, db)
    if not user or not user.hashed_password:
        return None
    if not verify_password(password, user.hashed_password):
        return None
    return user


async def ensure_user(email: str, password: str, db: AsyncSession) -> User:
    user = await _get_user_by_email(email, db)
    if user:
        return user
    hashed = get_password_hash(password or os.urandom(8).hex())
    username = email.split("@")[0]
    user = User(
        name=username,  # 添加name字段
        username=username,
        email=email,
        password_hash=hashed,  # 使用password_hash
        hashed_password=hashed,  # 兼容字段
        full_name=email,
        role="user",  # 添加role字段
        is_active=True,
        is_admin=False,
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


async def _check_api_key(key: str, db: AsyncSession) -> Optional[APIKey]:
    hashed = hashlib.sha256(key.encode()).hexdigest()
    result = await db.execute(select(APIKey).where(APIKey.prefix == key[:8]))
    api_keys = result.scalars().all()
    for api_key in api_keys:
        try:
            if verify_password(key, api_key.key_hash) or api_key.key_hash == hashed:
                if api_key.is_active:
                    return api_key
        except Exception:
            continue
    return None


async def require_auth(
    request: Request,
    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
    db: AsyncSession = Depends(get_db),
) -> Dict[str, Any]:
    """Require either Bearer JWT or X-API-Key header."""

    path = request.url.path
    allow_paths = {
        "/health",
        "/metrics",
        "/docs",
        "/redoc",
        "/openapi.json",
        "/api/channel/auth/login",
        "/api/admin/auth/login",
        "/api/providers/auth/login",
    }
    if path in allow_paths or not path.startswith("/api"):
        return {}

    api_key_header = request.headers.get("X-API-Key")
    if api_key_header:
        api_key = await _check_api_key(api_key_header, db)
        if api_key:
            principal = {"type": "api_key", "user_id": str(api_key.user_id), "scopes": api_key.scopes}
            request.state.principal = principal
            return principal

    if credentials and credentials.scheme.lower() == "bearer":
        token = credentials.credentials
        try:
            payload = jwt.decode(token, settings.secret_key, algorithms=[settings.jwt_algorithm])
            user_id: str | None = payload.get("sub")
            email: str | None = payload.get("email")
            if user_id is None:
                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token payload")
            request.state.principal = {"type": "jwt", "user_id": user_id, "email": email, "claims": payload}
            return request.state.principal
        except JWTError:
            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token") from None

    raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Not authenticated")


async def authenticate_request(request: Request, db: AsyncSession) -> Optional[Dict[str, Any]]:
    """Authenticate a request without FastAPI dependency injection (middleware use)."""

    allow_paths = {
        "/health",
        "/metrics",
        "/docs",
        "/redoc",
        "/openapi.json",
        "/api/channel/auth/login",
        "/api/admin/auth/login",
        "/api/providers/auth/login",
    }
    if request.url.path in allow_paths or not request.url.path.startswith("/api"):
        return {}

    api_key_header = request.headers.get("X-API-Key")
    if api_key_header:
        api_key = await _check_api_key(api_key_header, db)
        if api_key:
            principal = {"type": "api_key", "user_id": str(api_key.user_id), "scopes": api_key.scopes}
            return principal

    auth_header = request.headers.get("Authorization")
    if auth_header and auth_header.lower().startswith("bearer "):
        token = auth_header.split(" ", 1)[1]
        try:
            payload = jwt.decode(token, settings.secret_key, algorithms=[settings.jwt_algorithm])
            user_id: str | None = payload.get("sub")
            email: str | None = payload.get("email")
            if user_id is None:
                return None
            return {"type": "jwt", "user_id": user_id, "email": email, "claims": payload}
        except JWTError:
            return None

    return None