kingcap/openapi

1#!/usr/bin/env python32"""3Pre-start compatibility bridge for migrating from seakee/cpa-manager to4seakee/cpa-manager-plus inside the HuggingFace Space image.5 6It runs before supervisord starts cpa-manager-plus, so the new manager can see7restored settings/admin credentials on its first bootstrap pass.8"""9 10import base6411import hashlib12import json13import os14import re15import secrets16import sqlite317import stat18import time19from pathlib import Path20 21try:22    import psycopg23except Exception:  # pragma: no cover - the image installs psycopg24    psycopg = None25 26 27DB_PATH = os.getenv("USAGE_DB_PATH", "/data/usage.sqlite")28PG_DSN = os.getenv("USAGE_SYNC_PG_DSN") or os.getenv("DATABASE_URL") or os.getenv("PGSTORE_DSN", "")29SCHEMA = os.getenv("USAGE_SYNC_SCHEMA", "cpa_manager_usage")30DATA_KEY_PATH = os.getenv("CPA_MANAGER_DATA_KEY_PATH", "/data/data.key")31ADMIN_KEY_FILE = os.getenv("CPA_MANAGER_ADMIN_KEY_FILE", "/data/cpamp_admin.key")32ADMIN_CREDENTIAL_KEY = "admin_credential_v1"33 34if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", SCHEMA):35    raise SystemExit(f"invalid USAGE_SYNC_SCHEMA: {SCHEMA}")36 37 38def log(message: str):39    print(f"[prestart-migrate] {message}", flush=True)40 41 42def now_ms() -> int:43    return int(time.time() * 1000)44 45 46def ensure_parent(path: str):47    parent = Path(path).parent48    parent.mkdir(parents=True, exist_ok=True)49 50 51def write_secret_file(path: str, value: str):52    ensure_parent(path)53    p = Path(path)54    fd = os.open(str(p), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)55    with os.fdopen(fd, "w", encoding="utf-8") as f:56        f.write(value.strip())57        f.write("\n")58    try:59        p.chmod(stat.S_IRUSR | stat.S_IWUSR)60    except Exception:61        pass62 63 64def file_has_value(path: str) -> bool:65    try:66        return bool(Path(path).read_text(encoding="utf-8").strip())67    except Exception:68        return False69 70 71def ensure_data_key():72    if os.getenv("CPA_MANAGER_DATA_KEY") or os.getenv("CPA_MANAGER_DATA_KEY_FILE"):73        log("已通过环境变量/secret 管理 CPA_MANAGER_DATA_KEY，跳过 data.key 生成")74        return75    if file_has_value(DATA_KEY_PATH):76        log(f"复用已有 data key: {DATA_KEY_PATH}")77        return78 79    if PG_DSN:80        seed = f"cpamp-data-key-v1\0{SCHEMA}\0{PG_DSN}".encode("utf-8")81        raw = hashlib.sha256(seed).digest()82        source = "derived-from-pg-dsn"83    else:84        raw = secrets.token_bytes(32)85        source = "generated-random"86 87    write_secret_file(DATA_KEY_PATH, base64.b64encode(raw).decode("ascii").rstrip("="))88    log(f"已初始化 CPA-Manager-Plus data key: {DATA_KEY_PATH} ({source})")89 90 91def sqlite_connect():92    ensure_parent(DB_PATH)93    conn = sqlite3.connect(DB_PATH, timeout=30)94    conn.row_factory = sqlite3.Row95    conn.execute("PRAGMA busy_timeout = 5000")96    conn.execute("PRAGMA journal_mode = WAL")97    return conn98 99 100def ensure_settings_table(conn):101    conn.execute(102        """103        CREATE TABLE IF NOT EXISTS settings (104            key text primary key,105            value text not null,106            updated_at_ms integer not null107        )108        """109    )110    conn.commit()111 112 113def sqlite_get_setting(conn, key: str) -> str:114    ensure_settings_table(conn)115    row = conn.execute("SELECT value FROM settings WHERE key=?", (key,)).fetchone()116    return str(row["value"]) if row else ""117 118 119def sqlite_has_setting(conn, key: str) -> bool:120    ensure_settings_table(conn)121    row = conn.execute("SELECT 1 FROM settings WHERE key=?", (key,)).fetchone()122    return row is not None123 124 125def pg_connect():126    if not PG_DSN or psycopg is None:127        return None128    return psycopg.connect(PG_DSN, autocommit=True, prepare_threshold=None)129 130 131def restore_settings_from_pg(conn) -> int:132    if not PG_DSN:133        return 0134    try:135        pg = pg_connect()136        if pg is None:137            return 0138    except Exception as e:139        log(f"连接 PG 失败，跳过启动前 settings 恢复：{e}")140        return 0141 142    try:143        rows = pg.execute(144            f"""145            SELECT key, value_text, row_json, synced_at_ms146            FROM {SCHEMA}.settings147            ORDER BY key ASC148            """149        ).fetchall()150    except Exception as e:151        log(f"读取 PG settings 失败，可能是首次启动，已跳过：{e}")152        return 0153    finally:154        try:155            pg.close()156        except Exception:157            pass158 159    ensure_settings_table(conn)160    restored = 0161    for key, value_text, row_json, synced_at_ms in rows:162        key = str(key or "").strip()163        if not key:164            continue165 166        value = value_text167        updated_at_ms = int(synced_at_ms or now_ms())168 169        if isinstance(row_json, dict):170            value = row_json.get("value", value)171            try:172                updated_at_ms = int(row_json.get("updated_at_ms") or updated_at_ms)173            except Exception:174                pass175 176        if value is None:177            continue178 179        conn.execute(180            """181            INSERT INTO settings(key, value, updated_at_ms)182            VALUES(?, ?, ?)183            ON CONFLICT(key) DO UPDATE184            SET value=excluded.value,185                updated_at_ms=excluded.updated_at_ms186            """,187            (key, str(value), updated_at_ms),188        )189        restored += 1190 191    conn.commit()192    if restored:193        log(f"已从 PG 恢复 settings 到本地 SQLite：{restored} 条")194    return restored195 196 197def parse_json(raw: str):198    try:199        return json.loads(raw)200    except Exception:201        return None202 203 204def legacy_plain_management_key(conn) -> str:205    raw = sqlite_get_setting(conn, "setup")206    data = parse_json(raw)207    if not isinstance(data, dict):208        return ""209    key = str(data.get("managementKey") or data.get("ManagementKey") or "").strip()210    if key.startswith("enc:v1:"):211        return ""212    return key213 214 215def generate_admin_key() -> str:216    alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"217    return "cpamp_" + "".join(secrets.choice(alphabet) for _ in range(32))218 219 220def ensure_admin_key_file(conn):221    if os.getenv("CPA_MANAGER_ADMIN_KEY"):222        log("已通过 CPA_MANAGER_ADMIN_KEY 管理 Plus 管理员密钥")223        return224    if file_has_value(ADMIN_KEY_FILE):225        log(f"复用已有 Plus 管理员密钥文件: {ADMIN_KEY_FILE}")226        return227    if sqlite_has_setting(conn, ADMIN_CREDENTIAL_KEY):228        log("本地 SQLite 已有 Plus 管理员凭据，跳过管理员密钥文件生成")229        return230 231    legacy_key = legacy_plain_management_key(conn)232    if legacy_key:233        write_secret_file(ADMIN_KEY_FILE, legacy_key)234        log("已使用旧 /management.html 登录密码作为 CPA-Manager-Plus 管理员密钥")235        return236 237    generated = generate_admin_key()238    write_secret_file(ADMIN_KEY_FILE, generated)239    log(f"首次初始化 Plus 管理员密钥：{generated}")240    log(f"管理员密钥已写入 {ADMIN_KEY_FILE}；后续重启会复用该文件")241 242 243def main():244    ensure_parent(DB_PATH)245    ensure_data_key()246 247    conn = sqlite_connect()248    try:249        ensure_settings_table(conn)250        restore_settings_from_pg(conn)251        ensure_admin_key_file(conn)252    finally:253        conn.close()254 255 256if __name__ == "__main__":257    main()258