live-hub-py/web/backend/schemas.py

"""Pydantic 请求/响应模型"""

from datetime import datetime, timezone
from typing import Optional, Any
from pydantic import BaseModel, Field, ConfigDict, model_serializer


def _ensure_tz(dt: Optional[datetime]) -> Optional[datetime]:
    """确保 datetime 带有 UTC 时区信息，无时区的视为 UTC。"""
    if dt is None:
        return None
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt


# ---- 认证 ----
class LoginRequest(BaseModel):
    username: str
    password: str


class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"
    role: str
    username: str
    permissions: list[str]


class UserInfo(BaseModel):
    id: int
    username: str
    role: str
    is_active: bool
    remark: str = ""
    created_at: Optional[datetime] = None
    permissions: list[str] = []
    custom_permissions: Optional[list[str]] = None

    model_config = ConfigDict(from_attributes=True)

    @model_serializer
    def _serialize(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "username": self.username,
            "role": self.role,
            "is_active": self.is_active,
            "remark": self.remark,
            "created_at": _ensure_tz(self.created_at).isoformat() if self.created_at else None,
            "permissions": self.permissions,
            "custom_permissions": self.custom_permissions,
        }


# ---- 用户管理 ----
class UserCreate(BaseModel):
    username: str = Field(..., min_length=2, max_length=64)
    password: str = Field(..., min_length=6, max_length=128)
    role: str = Field("support", pattern="^(super_admin|operation|support)$")
    remark: str = ""


class UserUpdate(BaseModel):
    password: Optional[str] = None
    role: Optional[str] = None
    is_active: Optional[bool] = None
    remark: Optional[str] = None
    custom_permissions: Optional[list[str]] = None


# ---- 账号 ----
class AccountImport(BaseModel):
    """批量导入，文本格式：用户名|密码|邮箱|邮箱密码"""
    text: str


class AccountAssign(BaseModel):
    assigned_to: Optional[int] = None


class BatchAssign(BaseModel):
    account_ids: list[int]
    assigned_to: Optional[int] = None  # None=取消分配


class AccountTag(BaseModel):
    tag: Optional[str] = None
    account_ids: Optional[list[int]] = None


class AccountOut(BaseModel):
    id: int
    username: str
    # 敏感字段根据角色决定是否返回
    password: Optional[str] = None
    email: Optional[str] = None
    email_password: Optional[str] = None
    tag: str = ""
    assigned_to: Optional[int] = None
    assigned_username: Optional[str] = None
    remark: str = ""
    created_at: Optional[datetime] = None

    model_config = ConfigDict(from_attributes=True)

    @model_serializer
    def _serialize(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "username": self.username,
            "password": self.password,
            "email": self.email,
            "email_password": self.email_password,
            "tag": self.tag,
            "assigned_to": self.assigned_to,
            "assigned_username": self.assigned_username,
            "remark": self.remark,
            "created_at": _ensure_tz(self.created_at).isoformat() if self.created_at else None,
        }


# ---- 登录任务 ----
class LoginBatchRequest(BaseModel):
    account_ids: list[int]
    max_geetest_retries: int = 5
    max_proxy_retries: int = 10  # 代理切换次数，0=无限切换直到成功
    concurrency: int = 3  # 并发数，1-10


class LoginTaskOut(BaseModel):
    id: int
    batch_id: str
    account_id: int
    account_username: str = ""
    status: str
    cookie: str = ""
    message: str = ""
    created_by: int
    created_at: Optional[datetime] = None
    finished_at: Optional[datetime] = None

    model_config = ConfigDict(from_attributes=True)

    @model_serializer
    def _serialize(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "batch_id": self.batch_id,
            "account_id": self.account_id,
            "account_username": self.account_username,
            "status": self.status,
            "cookie": self.cookie,
            "message": self.message,
            "created_by": self.created_by,
            "created_at": _ensure_tz(self.created_at).isoformat() if self.created_at else None,
            "finished_at": _ensure_tz(self.finished_at).isoformat() if self.finished_at else None,
        }


# ---- 代理配置 ----
class ProxyConfigOut(BaseModel):
    enabled: bool = False
    api_url: str = ""
    http: str = ""
    https: str = ""
    whitelist_enabled: bool = False
    whitelist_uid: str = ""
    whitelist_ukey: str = ""


class ProxyConfigUpdate(ProxyConfigOut):
    pass


# ---- 通用 ----
class MessageResponse(BaseModel):
    message: str
    success: bool = True