32aef14ddb
- Introduce RuntimeModelResource and RuntimeKnowledgeBase classes to manage workflow resources. - Update AssistantConfig to include workflow_model_resources and workflow_knowledge_bases for better integration. - Refactor validation and processing logic in routes and services to accommodate workflow types. - Implement dynamic variable support for workflow assistants and enhance graph normalization. - Add ToolExecutor for reusable tool execution across different assistant types. - Update various services to ensure compatibility with new workflow features and improve error handling.
137 lines
4.1 KiB
Python
137 lines
4.1 KiB
Python
"""Single-admin authentication helpers.
|
|
|
|
The project currently needs one deployment-level super admin, configured by
|
|
environment variables. Tokens are signed with HMAC and stored in an HttpOnly
|
|
cookie so the frontend does not handle bearer tokens directly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import settings
|
|
from fastapi import HTTPException, Request, Response, WebSocket, status
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AdminUser:
|
|
username: str
|
|
display_name: str = "超级管理员"
|
|
role: str = "super_admin"
|
|
|
|
|
|
def _b64encode(data: bytes) -> str:
|
|
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _b64decode(data: str) -> bytes:
|
|
padding = "=" * (-len(data) % 4)
|
|
return base64.urlsafe_b64decode(f"{data}{padding}".encode("ascii"))
|
|
|
|
|
|
def _sign(payload: str) -> str:
|
|
digest = hmac.new(
|
|
settings.AUTH_SECRET_KEY.encode("utf-8"),
|
|
payload.encode("ascii"),
|
|
hashlib.sha256,
|
|
).digest()
|
|
return _b64encode(digest)
|
|
|
|
|
|
def _admin_user() -> AdminUser:
|
|
return AdminUser(username=settings.ADMIN_USERNAME)
|
|
|
|
|
|
def create_admin_token() -> str:
|
|
now = int(time.time())
|
|
payload = {
|
|
"sub": settings.ADMIN_USERNAME,
|
|
"role": "super_admin",
|
|
"iat": now,
|
|
"exp": now + settings.AUTH_TOKEN_EXPIRE_MINUTES * 60,
|
|
}
|
|
encoded_payload = _b64encode(
|
|
json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
|
)
|
|
return f"{encoded_payload}.{_sign(encoded_payload)}"
|
|
|
|
|
|
def verify_admin_token(token: str | None) -> AdminUser | None:
|
|
if not token or "." not in token:
|
|
return None
|
|
encoded_payload, signature = token.rsplit(".", 1)
|
|
if not hmac.compare_digest(_sign(encoded_payload), signature):
|
|
return None
|
|
try:
|
|
payload: dict[str, Any] = json.loads(_b64decode(encoded_payload))
|
|
except (ValueError, json.JSONDecodeError):
|
|
return None
|
|
if payload.get("sub") != settings.ADMIN_USERNAME:
|
|
return None
|
|
if int(payload.get("exp", 0)) < int(time.time()):
|
|
return None
|
|
return _admin_user()
|
|
|
|
|
|
def authenticate_admin(username: str, password: str) -> AdminUser | None:
|
|
username_ok = hmac.compare_digest(username, settings.ADMIN_USERNAME)
|
|
password_ok = hmac.compare_digest(password, settings.ADMIN_PASSWORD)
|
|
return _admin_user() if username_ok and password_ok else None
|
|
|
|
|
|
def set_auth_cookie(response: Response, token: str) -> None:
|
|
response.set_cookie(
|
|
settings.AUTH_COOKIE_NAME,
|
|
token,
|
|
max_age=settings.AUTH_TOKEN_EXPIRE_MINUTES * 60,
|
|
httponly=True,
|
|
secure=settings.AUTH_COOKIE_SECURE,
|
|
samesite=settings.AUTH_COOKIE_SAMESITE,
|
|
path="/",
|
|
)
|
|
|
|
|
|
def clear_auth_cookie(response: Response) -> None:
|
|
response.delete_cookie(
|
|
settings.AUTH_COOKIE_NAME,
|
|
httponly=True,
|
|
secure=settings.AUTH_COOKIE_SECURE,
|
|
samesite=settings.AUTH_COOKIE_SAMESITE,
|
|
path="/",
|
|
)
|
|
|
|
|
|
async def require_admin(request: Request) -> AdminUser:
|
|
user = verify_admin_token(request.cookies.get(settings.AUTH_COOKIE_NAME))
|
|
if not user:
|
|
authorization = request.headers.get("authorization", "")
|
|
scheme, _, encoded_credentials = authorization.partition(" ")
|
|
if scheme.lower() == "basic" and encoded_credentials:
|
|
try:
|
|
credentials = base64.b64decode(encoded_credentials).decode("utf-8")
|
|
username, separator, password = credentials.partition(":")
|
|
if separator:
|
|
user = authenticate_admin(username, password)
|
|
except (ValueError, UnicodeDecodeError):
|
|
user = None
|
|
if not user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="未登录或登录已过期",
|
|
)
|
|
return user
|
|
|
|
|
|
async def require_admin_websocket(websocket: WebSocket) -> AdminUser | None:
|
|
user = verify_admin_token(websocket.cookies.get(settings.AUTH_COOKIE_NAME))
|
|
if not user:
|
|
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
|
return None
|
|
return user
|