workbuddy-portal/manage.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli

"""WorkBuddy Portal —— 统一命令行入口。

采集 / 存储 / 呈现三件事都由本项目承担，不再依赖外部计划任务或自动化。

多用户说明：所有涉及「数据」或「凭证」的子命令都作用于**某一个账号**。
用 `-u/--user <用户名>` 指定；不指定时取「管理员优先、其次 id 最小」的那个
（老库升级后数据都在首个账号名下，所以不指定也能沿用旧习惯）。
唯独 `collect` 不带 `-u` 时会**逐个账号**跑一遍，与进程内调度线程的行为一致。

常用：
  python manage.py init                       初始化数据库（建表 + 默认配置 + 管理员）
  python manage.py serve                      启动 Web（0.0.0.0:8848，进程内含调度线程）
  python manage.py serve --port 9000 --debug  开发模式（reloader 下调度只启动一份）
  python manage.py collect                    为**所有已启用账号**各跑一次增量采集
  python manage.py collect -u alice           只为 alice 采集
  python manage.py migrate-csv [文件]          从旧版 CSV 存档导入（默认自动探测路径）
  python manage.py import-xlsx <文件>          从官网导出的 xlsx 合入
  python manage.py fill-prompt                补全缺失的 User Prompt
  python manage.py export-csv [路径]           导出与官网同构的 CSV（文件名带账号名）
  python manage.py stats [-u 账号]             看存档概况，不联网
  python manage.py users                      列出所有账号及其数据量 / 凭证状态
  python manage.py passwd <用户名> [新密码]     重置密码；账号不存在则创建
  python manage.py passwd <用户名> --role admin  新建或提权为管理员
  python manage.py status                     查看各账号的调度与最近采集状态
  python manage.py vacuum                     整理数据库（checkpoint + VACUUM）
"""
import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from workbuddy_portal import client, collect, config, db, query, scheduler   # noqa: E402
from workbuddy_portal import security                                        # noqa: E402


def _p(*a):
    print(*a)


def _warn(*a):
    print(*a)


# ---------------- 账号解析 ----------------
def _default_uid(conn):
    """没显式指定 `-u` 时的目标账号：管理员优先，其次 id 最小。

    老库升级后全部数据都归到首个账号，所以这个默认值正好等价于旧行为；
    全空库（只有实例级配置）返回 0，即「实例作用域」。
    """
    row = conn.execute("SELECT id FROM users ORDER BY is_admin DESC, id LIMIT 1").fetchone()
    return row["id"] if row else 0


def _resolve_uid(conn, name):
    """把 `-u` 的取值（用户名或数字 id）解析成 uid；解析不到返回 None。"""
    if name is None or name == "":
        return _default_uid(conn)
    row = db.user_by_name(conn, str(name))
    if row is None and str(name).isdigit():
        row = db.user_by_id(conn, int(name))
    if row is None:
        _p("[error] 没有这个账号：%s（用 manage.py users 查看）" % name)
        return None
    return row["id"]


def _uid_or_fail(conn, name):
    """解析失败时返回 None，调用方自行 return 2。"""
    uid = _resolve_uid(conn, name)
    if uid is None:
        return None
    if uid == 0:
        _warn("[warn] 库里还没有任何账号，本次按“实例作用域”执行（先跑 manage.py init）")
    return uid


def _ua_of(conn, uid):
    """人话描述某个账号的 Cookie 状态（只看密文可解性，不碰明文）。"""
    st = db.secret_state(conn, "cookie", uid)
    if st["broken"]:
        return "损坏（密钥换过，需重新粘贴）"
    if not st["set"]:
        return "未配置"
    return "已配置 %d 字符" % st["chars"]


# ---------------- 初始化 / 服务 ----------------
def cmd_init(args):
    db.init_db(admin_user=args.user, admin_password=args.password)
    conn = db.connect()
    try:
        uid = _default_uid(conn)
        n = query.totals(conn, uid)
        _p("数据库已就绪：%s" % config.SQLITE_PATH)
        _p("  账号数：%d" % db.user_count(conn))
        _p("  主账号存档 %d 条 / %.2f 积分 / %d 个活跃日" % (n["records"], n["credits"], n["days"]))
        _p("  操作账号：%s" % args.user)
        st = db.secret_state(conn, "cookie", uid)
        if st["broken"]:
            _p("  [warn] Cookie 密文无法解开（cookie_key 与写入时不一致），请登录后重新粘贴")
    finally:
        conn.close()


def cmd_serve(args):
    from workbuddy_portal import create_app
    app = create_app(start_scheduler=not args.no_scheduler)
    host = args.host or config.DEFAULT_HOST
    port = args.port or config.DEFAULT_PORT
    if args.debug:
        app.run(host=host, port=port, debug=True)
        return
    try:
        from waitress import serve
        _p("生产模式（waitress）监听 http://%s:%d" % (host, port))
        serve(app, host=host, port=port, threads=8, ident="workbuddy-portal")
    except ImportError:
        _p("[warn] 未安装 waitress，回退到 Flask 内置服务器（生产建议 pip install waitress）")
        app.run(host=host, port=port, threaded=True)


# ---------------- 采集 / 导入 / 导出 ----------------
def cmd_collect(args):
    """不带 -u 时逐个已启用账号采集；带 -u 时只采一个。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        if args.user:
            uid = _resolve_uid(conn, args.user)
            if uid is None:
                return 2
            targets = [uid]
        else:
            targets = [r["id"] for r in db.active_users(conn)]
            if not targets:
                _p("[error] 没有任何启用中的账号")
                return 2
    finally:
        conn.close()

    rc = 0
    for uid in targets:
        conn = db.connect()
        try:
            row = db.user_by_id(conn, uid)
            who = row["username"] if row else "uid=%s" % uid
            st = db.secret_state(conn, "cookie", uid)
            if not st["set"] or st["broken"]:
                _p("— %s：跳过（Cookie %s）" % (who, _ua_of(conn, uid)))
                continue
        finally:
            conn.close()
        _p("— %s：" % who)
        try:
            r = collect.run_sync(trigger="cli", uid=uid)
        except collect.Busy as e:
            _p("  [busy] %s" % e)
            rc = rc or 1
            continue
        except collect.NotReady as e:
            _p("  [skip] %s" % e)
            continue
        except db.SecretUnreadable as e:
            _p("  [error] %s（请重新粘贴 Cookie）" % e)
            rc = rc or 4
            continue
        except collect.ApiError as e:
            _p("  [error] %s" % e)
            rc = rc or (3 if e.cookie_expired else 5)
            continue
        for line in r["lines"]:
            _p("  " + line)
    return rc


def cmd_migrate_csv(args):
    db.init_db(create_admin=False)
    path = args.path
    if not path:
        for c in config.LEGACY_CSV_CANDIDATES:
            if os.path.exists(c):
                path = c
                break
    if not path:
        _p("[error] 找不到旧存档 CSV，请显式指定路径")
        return 2
    conn = db.connect()
    try:
        uid = _uid_or_fail(conn, args.user)
        if uid is None:
            return 2
        collect.migrate_from_csv(conn, uid, path, log=_p)
        n = query.totals(conn, uid)
        _p("账号 uid=%s 当前存档：%d 条 / %.2f 积分 / %s ~ %s"
           % (uid, n["records"], n["credits"], n["firstDay"], n["lastDay"]))
    finally:
        conn.close()
    return 0


def cmd_import_xlsx(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        uid = _uid_or_fail(conn, args.user)
        if uid is None:
            return 2
        collect.import_xlsx(conn, uid, args.path, log=_p)
    finally:
        conn.close()
    return 0


def cmd_fill_prompt(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        uid = _uid_or_fail(conn, args.user)
        if uid is None:
            return 2
        collect.fill_prompt(conn, uid, log=_p)
    finally:
        conn.close()
    return 0


def cmd_export_csv(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        uid = _uid_or_fail(conn, args.user)
        if uid is None:
            return 2
        row = db.user_by_id(conn, uid)
        path, n = collect.export_csv(conn, uid, args.path,
                                     username=(row["username"] if row else None))
        _p("已导出 %d 条 -> %s" % (n, path))
    finally:
        conn.close()
    return 0


# ---------------- 统计 ----------------
def cmd_stats(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        # 先给一张全局概览：多用户下最常问的就是「一共多少、谁占多少」
        rows = conn.execute(
            "SELECT u.id, u.username, u.display_name, u.is_admin, u.status,"
            "       COUNT(r.request_id) AS records, COALESCE(SUM(r.credits),0) AS credits"
            "  FROM users u LEFT JOIN usage_records r ON r.user_id=u.id"
            " GROUP BY u.id ORDER BY records DESC, u.id").fetchall()
        total = conn.execute("SELECT COUNT(*) AS c, COALESCE(SUM(credits),0) AS s"
                             " FROM usage_records").fetchone()
        _p("全库存档：%d 条 / %.2f 积分 / %d 个账号"
           % (total["c"], total["s"], db.user_count(conn)))
        if rows:
            _p("")
            _p("%-4s %-16s %-10s %-6s %-8s %8s %12s" %
               ("id", "用户名", "角色", "状态", "Cookie", "调用", "积分"))
            for r in rows:
                _p("%-4d %-16s %-10s %-6s %-8s %8d %12.2f" %
                   (r["id"], r["username"], "管理员" if r["is_admin"] else "普通",
                    "启用" if r["status"] == "active" else "停用",
                    _ua_of(conn, r["id"]), r["records"], r["credits"]))

        uid = _resolve_uid(conn, args.user)
        if uid is None:
            return 2
        t = query.totals(conn, uid)
        if not t["records"]:
            _p("")
            _p("（uid=%s 没有数据；换个 -u，或先跑 manage.py migrate-csv / collect）" % uid)
            return 0
        _p("")
        _p("== uid=%s 明细 ==" % uid)
        _p("存档：%d 条 / %.2f 积分 / %d 个活跃日（%s ~ %s）"
           % (t["records"], t["credits"], t["days"], t["firstDay"], t["lastDay"]))
        _p("计费调用 %d · 免费调用 %d · 模型 %d · 客户端 %d"
           % (t["billableCalls"], t["freeCalls"], t["models"], t["clients"]))
        _p("")
        _p("%-24s %8s %12s %10s %8s" % ("模型", "调用", "积分", "单次均价", "免费占比"))
        for m in query.dims(conn, uid)["model"]:
            _p("%-24s %8d %12.2f %10.2f %7.0f%%"
               % (m["name"], m["calls"], m["credits"], m["avgPerCall"], m["freeRate"] * 100))
        runs = conn.execute("SELECT id,trigger,status,started_at,added,dup,total,message"
                            " FROM collect_runs WHERE user_id=? ORDER BY id DESC LIMIT 5",
                            (uid,)).fetchall()
        if runs:
            _p("")
            _p("最近采集：")
            for r in runs:
                _p("  #%d %s %s +%d/%d → %d  %s"
                   % (r["id"], r["started_at"], r["status"], r["added"], r["dup"],
                      r["total"], r["message"] or ""))
    finally:
        conn.close()
    return 0


def cmd_users(args):
    """列出账号：角色 / 状态 / 数据量 / 凭证状态 / 最近登录。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        rows = conn.execute(
            "SELECT u.*, COUNT(r.request_id) AS records,"
            "       COALESCE(SUM(r.credits),0) AS credits,"
            "       MAX(r.day) AS last_day"
            "  FROM users u LEFT JOIN usage_records r ON r.user_id=u.id"
            " GROUP BY u.id ORDER BY u.id").fetchall()
        if not rows:
            _p("还没有任何账号。跑 manage.py init 建管理员，或让用户自助注册。")
            return 0
        _p("%-4s %-16s %-12s %-6s %-6s %10s %8s %12s  %s" %
           ("id", "用户名", "显示名", "角色", "状态", "Cookie", "调用", "积分", "最近登录 IP"))
        for r in rows:
            _p("%-4d %-16s %-12s %-6s %-6s %10s %8d %12.2f  %s" %
               (r["id"], r["username"], r["display_name"] or "",
                "管理员" if r["is_admin"] else "普通",
                "启用" if r["status"] == "active" else "停用",
                _ua_of(conn, r["id"]), r["records"], r["credits"],
                r["last_login_ip"] or "—"))
        _p("")
        _p("提示：cookie_key 或 secret_key 可在 data/instance.json 里找到，"
           "二者权限等同管理员口令，切勿随仓库分发。")
    finally:
        conn.close()
    return 0


# ---------------- 凭证 / 口令 ----------------
def cmd_import_creds(args):
    """把 VSCode 设置里的 cookie / userAgent 接管进数据库（一次性迁移用）。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        uid = _uid_or_fail(conn, args.user)
        if uid is None:
            return 2
        found = client.read_vscode_creds()
        if not found:
            _p("[error] 没找到 VSCode 系编辑器的 settings.json")
            return 2
        hit = None
        for path, cookie, ua, note in found:
            _p("  %-70s %s" % (path, note))
            if cookie and hit is None:
                hit = (cookie, ua)
        if not hit:
            _p("[error] 这些文件里都没有 codebuddyUsage.cookie，请到「配置管理」页手工粘贴")
            return 2
        cookie, ua = hit
        # 一律走 set_secret（内部就是 set_setting），值在落库前完成加密
        db.set_secret(conn, "cookie", cookie, uid)
        if ua:
            db.set_setting(conn, "user_agent", ua, uid)
        row = db.user_by_id(conn, uid)
        db.audit(conn, "import_creds", (row["username"] if row else "cli"),
                 "从 VSCode 设置导入凭证（%d 字符）" % len(cookie), "127.0.0.1", uid)
        _p("已把 Cookie（%d 字符）与 User-Agent（%s）写入账号 uid=%s"
           % (len(cookie), "有" if ua else "无", uid))
    finally:
        conn.close()
    return 0


def cmd_passwd(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        row = db.user_by_name(conn, args.user)
        pwd = args.password
        is_admin = 1 if args.role == "admin" else 0
        if row:
            sets, vals = [], []
            if pwd:
                sets.append("password_hash=?")
                vals.append(security.hash_password(pwd))
            if args.role:
                sets.append("is_admin=?")
                vals.append(is_admin)
            if args.activate:
                sets.append("status='active'")
            if not sets:
                _p("没给新密码也没给 --role，什么都没改")
                return 0
            vals.append(row["id"])
            conn.execute("UPDATE users SET %s WHERE id=?" % ",".join(sets), vals)
            _p("已更新账号 %s（%s）" % (args.user, "，".join(
                x.split("=")[0] for x in sets)))
        else:
            if not pwd:
                _p("[error] 新账号必须给出密码")
                return 2
            conn.execute(
                "INSERT INTO users(username,password_hash,display_name,is_admin,status,"
                " created_at) VALUES(?,?,?,?,'active',?)",
                (args.user, security.hash_password(pwd), args.user, is_admin, db.now_str()))
            uid = conn.execute("SELECT id FROM users WHERE username=?",
                               (args.user,)).fetchone()["id"]
            db.audit(conn, "user_create", "cli", "命令行创建账号 %s" % args.user,
                     "127.0.0.1", uid)
            _p("已创建账号 %s（uid=%s，%s）"
               % (args.user, uid, "管理员" if is_admin else "普通"))
        if pwd:
            _p("密码：%s" % pwd)
    finally:
        conn.close()
    return 0


# ---------------- 状态 / 维护 ----------------
def cmd_status(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        _p("服务器时间：%s" % db.now_str())
        _p("互斥锁：%s" % ("存在（有采集在跑）" if os.path.exists(collect.LOCK_PATH) else "不存在"))
        _p("调度总开关（实例级）：%s"
           % ("启用" if db.get_bool(conn, "schedule_enabled", True) else "停用"))
        users = db.active_users(conn)
        if not users:
            _p("（没有任何启用中的账号。先 manage.py init 建管理员，"
               "或在「用户管理」页启用一个账号）")
        for u in users:
            uid = u["id"]
            _p("")
            _p("== uid=%d %s%s ==" % (uid, u["username"],
                                      "（管理员）" if u["is_admin"] else ""))
            _p("  调度：%s / 每日 %s" %
               ("启用" if db.get_bool(conn, "schedule_enabled", True, uid) else "停用",
                ", ".join(scheduler.slots(conn, uid)) or "—"))
            nxt = scheduler.next_run_at(conn, uid)
            _p("  下次执行：%s" % (nxt.strftime("%Y-%m-%d %H:%M:%S") if nxt else "—"))
            _p("  Cookie：%s" % _ua_of(conn, uid))
            last = conn.execute("SELECT * FROM collect_runs WHERE user_id=?"
                                " ORDER BY id DESC LIMIT 1", (uid,)).fetchone()
            if last:
                _p("  最近采集：#%d %s %s %s" % (last["id"], last["started_at"],
                                            last["status"], last["message"] or ""))
            else:
                _p("  最近采集：无")
        _p("")
        _p("（注意：调度线程只在 manage.py serve 进程内运行；"
           "多实例部署时其余实例要设 WB_DISABLE_SCHEDULER=1）")
    finally:
        conn.close()


def cmd_vacuum(args):
    """整理数据库：checkpoint WAL + VACUUM 回收空间。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        before = os.path.getsize(config.SQLITE_PATH) if os.path.exists(config.SQLITE_PATH) else 0
        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
        conn.execute("VACUUM")
        conn.execute("PRAGMA optimize")
        after = os.path.getsize(config.SQLITE_PATH) if os.path.exists(config.SQLITE_PATH) else 0
        _p("数据库整理完成：%s → %s（%+d 字节）" % (_human(before), _human(after), after - before))
        n = conn.execute("SELECT COUNT(*) FROM usage_records").fetchone()[0]
        _p("全库存档 %d 条记录" % n)
    finally:
        conn.close()
    return 0


def _human(n):
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024 or unit == "GB":
            return ("%d B" % n) if unit == "B" else ("%.1f %s" % (n, unit))
        n /= 1024.0


# ---------------- 参数表 ----------------
def _add_user_opt(p, help_text="作用账号（用户名或 uid），默认取管理员 / 最小 id"):
    p.add_argument("-u", "--user", default=None, help=help_text)


def main():
    ap = argparse.ArgumentParser(description="WorkBuddy Portal（workbuddy-portal）",
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 epilog=__doc__)
    sub = ap.add_subparsers(dest="cmd")

    s = sub.add_parser("init", help="初始化数据库")
    s.add_argument("--user", default="admin", help="首个管理员用户名（仅库为空时生效）")
    s.add_argument("--password", default=None)
    s.set_defaults(func=cmd_init)

    s = sub.add_parser("serve", help="启动 Web 服务")
    s.add_argument("--host", default=None)
    s.add_argument("--port", type=int, default=None)
    s.add_argument("--debug", action="store_true")
    s.add_argument("--no-scheduler", action="store_true", help="不启动进程内调度线程")
    s.set_defaults(func=cmd_serve)

    s = sub.add_parser("collect", help="执行一次增量采集（默认所有启用账号）")
    _add_user_opt(s, "只采这一个账号；不传则逐个启用账号采集")
    s.set_defaults(func=cmd_collect)

    s = sub.add_parser("migrate-csv", help="从旧版 CSV 导入")
    s.add_argument("path", nargs="?")
    _add_user_opt(s, "这份老存档算谁的")
    s.set_defaults(func=cmd_migrate_csv)

    s = sub.add_parser("import-xlsx", help="从官网 xlsx 导入")
    s.add_argument("path")
    _add_user_opt(s)
    s.set_defaults(func=cmd_import_xlsx)

    s = sub.add_parser("fill-prompt", help="补全缺失的 User Prompt")
    _add_user_opt(s)
    s.set_defaults(func=cmd_fill_prompt)

    s = sub.add_parser("export-csv", help="导出 CSV")
    s.add_argument("path", nargs="?")
    _add_user_opt(s)
    s.set_defaults(func=cmd_export_csv)

    s = sub.add_parser("stats", help="存档概况（先全库概览，再给指定账号明细）")
    _add_user_opt(s)
    s.set_defaults(func=cmd_stats)

    s = sub.add_parser("users", help="列出所有账号及其数据量 / 凭证状态")
    s.set_defaults(func=cmd_users)

    s = sub.add_parser("import-creds", help="从 VSCode 设置导入 cookie / UA 到数据库")
    _add_user_opt(s)
    s.set_defaults(func=cmd_import_creds)

    s = sub.add_parser("passwd", help="重置 / 创建登录账号")
    s.add_argument("user")
    s.add_argument("password", nargs="?")
    s.add_argument("--role", choices=["admin", "user"], default=None,
                   help="不提则保持原角色；新建时默认普通账号")
    s.add_argument("--activate", action="store_true", help="顺便把状态改回启用")
    s.set_defaults(func=cmd_passwd)

    s = sub.add_parser("status", help="各账号的调度与最近采集状态")
    s.set_defaults(func=cmd_status)

    s = sub.add_parser("vacuum", help="整理数据库（checkpoint + VACUUM）")
    s.set_defaults(func=cmd_vacuum)

    args = ap.parse_args()
    if not getattr(args, "func", None):
        ap.print_help()
        return 0
    config.ensure_dirs()
    return args.func(args) or 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        _p("\n已中断")
        sys.exit(130)