WechatOnCloud/bridge/woc_bridge/models/db.py

"""DB 域模型：密钥注入 / 密钥状态 / DB 初始化。"""

from __future__ import annotations

from typing import Optional

from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# DB 解密接口（POST /api/db/decrypt, GET /api/db/key/status）
# ---------------------------------------------------------------------------
class DbDecryptRequest(BaseModel):
    """手动 key 注入请求。

    外部系统（如 ForcePilot）或管理员通过 POST /api/db/decrypt 传入 64 位
    十六进制密钥，bridge 验证并缓存，后续所有 DB 查询接口立即恢复可用。
    """

    key: str = Field(description="64 位十六进制 SQLCipher 密钥")
    salt: Optional[str] = Field(
        default=None,
        description="可选：该 key 对应的 DB salt（32 位十六进制）。未指定时 bridge 自动读取当前 DB 的 salt",
    )


class DbDecryptResponse(BaseModel):
    """key 注入结果。

    - verified=true 时 key 已缓存，后续 DB 查询接口立即可用
    - verified=false 时 key 不匹配，未缓存
    """

    success: bool = Field(description="请求是否处理成功（不代表 key 正确）")
    verified: bool = Field(default=False, description="key 是否通过验证并缓存")
    key_mode: Optional[str] = Field(
        default=None,
        description="key 形态：enc_key（已派生密钥）/ key_material（原始密钥材料）；验证失败时为 None",
    )
    error: Optional[str] = Field(default=None, description="验证失败或特殊状态：key_mismatch / file_too_small / db_not_encrypted")


class DbKeyStatusResponse(BaseModel):
    """key 缓存状态查询响应。

    供调用方判断是否需要注入 key 或等待自动提取。
    """

    cached: bool = Field(default=False, description="是否有缓存的 key")
    source: Optional[str] = Field(
        default=None,
        description="key 来源：env（环境变量）/ api（POST 注入）/ auto_extract（内存扫描）/ file（持久化文件）",
    )
    verified: bool = Field(default=False, description="缓存 key 是否已通过验证")
    key_prefix: Optional[str] = Field(
        default=None,
        description="key 前 4 + ... + 后 4 字符，用于确认是哪个 key（不泄露完整 key）",
    )


class DbInitRequest(BaseModel):
    """DB 初始化请求（显式触发密钥提取）。"""

    pid: Optional[int] = Field(default=None, description="可选：指定微信进程 PID")
    db_dir: Optional[str] = Field(default=None, description="可选：指定微信数据目录，默认自动检测")
    force: bool = Field(default=False, description="是否强制重新提取，忽略已有 keys 文件")


class DbInitResponse(BaseModel):
    """DB 初始化响应。"""

    success: bool = Field(description="请求是否受理")
    state: str = Field(description="状态：started / already_done / in_progress / failed")
    message: str = Field(description="阶段描述")
    key_count: Optional[int] = Field(default=None, description="已提取到的 salt→key 映射数量")


class DbInitStatusResponse(BaseModel):
    """DB 初始化后台任务状态。"""

    state: str = Field(description="状态：idle / running / success / failed")
    progress_pct: Optional[float] = Field(default=None, description="进度百分比 0-100")
    message: Optional[str] = Field(default=None, description="当前阶段描述")
    key_count: Optional[int] = Field(default=None, description="成功提取的密钥数")
    error: Optional[str] = Field(default=None, description="失败原因")