全浏览器顺手牵 Token 神器源码精讲

#!/usr/bin/env python3
"""extract_token.py —— 从本机任意主流浏览器或环境变量/参数提取 PrimeHub 登录 token

支持以下 4 种提取来源（不再局限于 Firefox！）：
1. 【Chrome / Edge / Brave / Chromium】自动扫描 Local Storage/leveldb（复制快照防锁库，兼容 UTF-8 / UTF-16LE 与 Snappy 块）
2. 【Firefox（含 Windows MSIX / macOS / Linux）】自动扫描 storage/default/https+++primehub.primebot.com/ls/data.sqlite（含 WAL 快照与 Snappy 解压）
3. 【环境变量 PRIMEHUB_TOKEN】若设置了环境变量 PRIMEHUB_TOKEN，直接使用
4. 【命令行参数 --token】支持在任意浏览器按 F12 复制 skillhub_token 后通过 `--token "eyJ..."` 直接写入

用法：
  python3 extract_token.py [输出文件(默认 workspace/.local/ph_token.txt)]
  python3 extract_token.py --browser auto|chrome|edge|firefox|brave|chromium [输出文件]
  python3 extract_token.py --token "eyJhbG..." [输出文件]
"""
import argparse
import base64
import glob
import json
import os
import re
import shutil
import sqlite3
import sys
import tempfile
from pathlib import Path

DEFAULT_TOKEN_FILE = os.environ.get(
    "PRIMEHUB_TOKEN_FILE", str(Path(__file__).resolve().parents[2] / ".local/ph_token.txt")
)

JWT_ASCII_RE = re.compile(rb"eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}")


def clean_token(raw: str) -> str:
    """剥离首尾空白、外层单双引号及可能的 Bearer 前缀。"""
    tok = raw.strip().strip('"').strip("'").strip()
    if tok.lower().startswith("bearer "):
        tok = tok[7:].strip()
    return tok


def jwt_timestamp_score(token: str) -> int:
    """解析 JWT payload 中的 exp 或 iat 时间戳用于新旧排序；无法解析时返回 0。"""
    parts = token.split(".")
    if len(parts) != 3:
        return 0
    try:
        payload_b64 = parts[1] + "=" * (-len(parts[1]) % 4)
        payload = json.loads(base64.urlsafe_b64decode(payload_b64.encode("ascii")).decode("utf-8"))
        if isinstance(payload, dict):
            return int(payload.get("exp") or payload.get("iat") or 0)
    except Exception:
        pass
    return 0


def save_token_file(out_file: Path, token: str, source_label: str = "") -> None:
    """以 0600 权限安全写入 token 文件，仅打印路径、字符数与来源，绝不打印 token 明文。"""
    tok = clean_token(token)
    out_file.parent.mkdir(parents=True, exist_ok=True)
    descriptor = os.open(out_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(descriptor, "w", encoding="utf-8") as output:
        output.write(tok)
    suffix = ("，来源: " + source_label) if source_label else ""
    print("token → %s（%d 字符%s；有效期以服务端为准）" % (out_file, len(tok), suffix))


# ==============================================================================
# 来源 1：Firefox (SQLite + WAL + Snappy)
# ==============================================================================

def firefox_databases() -> list[Path]:
    """支持 Windows、MSIX、macOS、Linux 及显式指定的 Firefox profile。"""
    profile = os.environ.get("PRIMEHUB_FIREFOX_PROFILE")
    roots = [Path(profile)] if profile else [
        Path(os.environ.get("APPDATA", Path.home() / "AppData/Roaming"))
        / "Mozilla/Firefox/Profiles",
        Path.home() / "Library/Application Support/Firefox/Profiles",
        Path.home() / ".mozilla/firefox",
        Path.home() / "snap/firefox/common/.mozilla/firefox",
    ]
    if not profile and os.environ.get("LOCALAPPDATA"):
        packages = Path(os.environ["LOCALAPPDATA"]) / "Packages"
        roots.extend(packages.glob(
            "Mozilla.MozillaFirefox_*/LocalCache/Roaming/Mozilla/Firefox/Profiles"))
    suffix = "storage/default/https+++primehub.primebot.com/ls/data.sqlite"
    databases = set()
    for root in roots:
        for pattern in (suffix, "*/" + suffix):
            databases.update(Path(p) for p in glob.glob(glob.escape(str(root)) + "/" + pattern))
    return sorted(databases, key=lambda p: p.stat().st_mtime, reverse=True)


def read_storage(database: Path) -> dict:
    """用 SQLite backup 读取一致快照，包含仍在 WAL 中的登录数据。"""
    source = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)
    snapshot = sqlite3.connect(":memory:")
    try:
        source.backup(snapshot)
        return dict(snapshot.execute("SELECT key, value FROM data").fetchall())
    finally:
        snapshot.close()
        source.close()


def snappy_decompress(data: bytes) -> bytes:
    """纯 stdlib snappy 解压：uvarint 声明长度 + literal/copy tag 流。"""
    def uvarint(d, i):
        v = s = 0
        while True:
            b = d[i]; i += 1
            v |= (b & 0x7F) << s; s += 7
            if not (b & 0x80):
                break
        return v, i

    declared, i = uvarint(data, 0)
    out = bytearray(); n = len(data)
    while i < n:
        tag = data[i]; i += 1
        t = tag & 0x03
        if t == 0:                                  # literal
            ln = tag >> 2
            if ln < 60:
                ln += 1
            else:
                extra = ln - 59
                ln = int.from_bytes(data[i:i + extra], "little") + 1
                i += extra
            out += data[i:i + ln]; i += ln
        else:                                       # copy
            if t == 1:
                ln = ((tag >> 2) & 0x07) + 4
                off = ((tag >> 5) << 8) | data[i]; i += 1
            elif t == 2:
                ln = (tag >> 2) + 1
                off = int.from_bytes(data[i:i + 2], "little"); i += 2
            else:
                ln = (tag >> 2) + 1
                off = int.from_bytes(data[i:i + 4], "little"); i += 4
            for _ in range(ln):
                out.append(out[-off])
    assert len(out) == declared, (len(out), declared)
    return bytes(out)


def extract_from_firefox() -> str | None:
    """从 Firefox 数据库列表提取 skillhub_token。"""
    for db in firefox_databases():
        try:
            rows = read_storage(db)
            if "skillhub_token" in rows:
                return clean_token(snappy_decompress(rows["skillhub_token"]).decode("utf-8"))
        except Exception:
            continue
    return None


# ==============================================================================
# 来源 2：Chrome / Edge / Brave / Chromium (LevelDB .log / .ldb)
# ==============================================================================

def chromium_user_data_roots(browser: str = "auto") -> list[tuple[str, Path]]:
    """返回指定或全部 Chromium 内核浏览器的 User Data 或 Profile 根目录列表。"""
    custom = os.environ.get("PRIMEHUB_CHROME_PROFILE") or os.environ.get("PRIMEHUB_BROWSER_PATH")
    if custom:
        return [("custom", Path(custom))]

    local_app = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData/Local"))
    home = Path.home()
    catalog = {
        "chrome": [
            local_app / "Google/Chrome/User Data",
            home / "Library/Application Support/Google/Chrome",
            home / ".config/google-chrome",
        ],
        "edge": [
            local_app / "Microsoft/Edge/User Data",
            home / "Library/Application Support/Microsoft Edge",
            home / ".config/microsoft-edge",
        ],
        "brave": [
            local_app / "BraveSoftware/Brave-Browser/User Data",
            home / "Library/Application Support/BraveSoftware/Brave-Browser",
            home / ".config/BraveSoftware/Brave-Browser",
        ],
        "chromium": [
            local_app / "Chromium/User Data",
            home / "Library/Application Support/Chromium",
            home / ".config/chromium",
            home / "snap/chromium/common/chromium",
        ],
    }
    selected = catalog.keys() if browser == "auto" else [browser]
    results = []
    for name in selected:
        for path in catalog.get(name, []):
            results.append((name, path))
    return results


def chromium_leveldb_dirs(browser: str = "auto") -> list[tuple[str, Path]]:
    """查找 Chromium 系浏览器所有 Profile 下的 Local Storage/leveldb 目录。"""
    found = []
    seen = set()
    for name, root in chromium_user_data_roots(browser):
        for pattern in ("Local Storage/leveldb", "Default/Local Storage/leveldb", "Profile */Local Storage/leveldb"):
            for match in glob.glob(glob.escape(str(root)) + "/" + pattern):
                p = Path(match)
                if p.is_dir() and p not in seen:
                    seen.add(p)
                    found.append((name, p))
    return found


def _scan_bytes_for_skillhub_jwt(raw: bytes) -> list[str]:
    """在字节流（支持 UTF-8/Latin1 与 UTF-16LE）中检索紧跟 skillhub_token 后的 JWT。"""
    tokens = []
    # 1. UTF-8 / Latin1 编码的 key + value
    key_utf8 = b"skillhub_token"
    pos = 0
    while True:
        idx = raw.find(key_utf8, pos)
        if idx == -1:
            break
        window = raw[idx + len(key_utf8): idx + len(key_utf8) + 4096]
        match = JWT_ASCII_RE.search(window)
        if match and match.start() < 128:
            tokens.append(match.group(0).decode("ascii"))
        pos = idx + len(key_utf8)

    # 2. UTF-16LE 编码的 key 或 value
    key_utf16 = "skillhub_token".encode("utf-16le")
    pos = 0
    while True:
        idx = raw.find(key_utf16, pos)
        if idx == -1:
            break
        window = raw[idx + len(key_utf16): idx + len(key_utf16) + 8192]
        # 压缩奇数位 0x00 还原 ASCII 视图
        collapsed = window[0::2]
        match = JWT_ASCII_RE.search(collapsed)
        if match and match.start() < 128:
            tokens.append(match.group(0).decode("ascii"))
        pos = idx + len(key_utf16)
    return tokens


def read_chromium_leveldb_token(leveldb_dir: Path) -> str | None:
    """拷贝 LevelDB 快照并提取最新的 skillhub_token。"""
    candidates = []
    with tempfile.TemporaryDirectory(prefix="primehub_ldb_") as tmp_dir:
        tmp_path = Path(tmp_dir)
        files = sorted(
            [p for p in leveldb_dir.iterdir() if p.suffix in (".log", ".ldb", ".sst")],
            key=lambda p: p.stat().st_mtime,
        )
        for src_file in files:
            try:
                dst_file = tmp_path / src_file.name
                shutil.copy2(src_file, dst_file)
                data = dst_file.read_bytes()
            except Exception:
                continue
            mtime = int(src_file.stat().st_mtime)
            for tok in _scan_bytes_for_skillhub_jwt(data):
                candidates.append((jwt_timestamp_score(tok), mtime, tok))
            # 若整个文件/块为纯 Snappy 流，也尝试解压后扫描
            try:
                decompressed = snappy_decompress(data)
                for tok in _scan_bytes_for_skillhub_jwt(decompressed):
                    candidates.append((jwt_timestamp_score(tok), mtime, tok))
            except Exception:
                pass
    if not candidates:
        return None
    # 按 (JWT过期时间/签发时间, 文件修改时间) 升序排，取最新一条
    candidates.sort(key=lambda item: (item[0], item[1]))
    return clean_token(candidates[-1][2])


def extract_from_chromium(browser: str = "auto") -> tuple[str, str] | None:
    """从 Chrome / Edge / Brave / Chromium 提取最新 token，返回 (browser_name, token)。"""
    best = None  # (jwt_score, dir_mtime, browser_name, token)
    for name, ldb_dir in chromium_leveldb_dirs(browser):
        tok = read_chromium_leveldb_token(ldb_dir)
        if tok:
            score = (jwt_timestamp_score(tok), int(ldb_dir.stat().st_mtime), name, tok)
            if best is None or score[:2] >= best[:2]:
                best = score
    if best:
        return best[2], best[3]
    return None


def main():
    parser = argparse.ArgumentParser(
        description="从 Firefox / Chrome / Edge / Brave / Chromium 或命令行提取 PrimeHub 登录 token"
    )
    parser.add_argument("out_file", nargs="?", default=DEFAULT_TOKEN_FILE, help="输出 token 文件路径")
    parser.add_argument(
        "--browser",
        choices=("auto", "firefox", "chrome", "edge", "brave", "chromium"),
        default="auto",
        help="指定浏览器来源（默认 auto：自动扫描 Firefox 与 Chrome/Edge 等）",
    )
    parser.add_argument("--token", help="直接传入从浏览器 DevTools 复制的 skillhub_token 字符串")
    args = parser.parse_args()
    out_file = Path(args.out_file)

    # 1. 优先支持显式 --token 参数或 PRIMEHUB_TOKEN 环境变量
    direct_token = args.token or os.environ.get("PRIMEHUB_TOKEN")
    if direct_token and clean_token(direct_token):
        save_token_file(out_file, direct_token, "direct")
        return

    # 2. Firefox 提取路径
    if args.browser in ("auto", "firefox"):
        dbs = firefox_databases()
        for db in dbs:
            try:
                rows = read_storage(db)
            except Exception:
                continue
            if "skillhub_token" in rows:
                tok = snappy_decompress(rows["skillhub_token"]).decode("utf-8").strip('"')
                save_token_file(out_file, tok, "firefox")
                return
        if args.browser == "firefox":
            if not dbs:
                raise SystemExit("没找到 Firefox 的 primehub localStorage（先在 Firefox 登录一次 PrimeHub）")
            print("没找到 skillhub_token——Firefox 登录态过期，重新登录一次", file=sys.stderr)
            sys.exit(1)

    # 3. Chrome / Edge / Brave / Chromium 提取路径
    if args.browser in ("auto", "chrome", "edge", "brave", "chromium"):
        chromium_hit = extract_from_chromium(args.browser)
        if chromium_hit:
            browser_name, tok = chromium_hit
            save_token_file(out_file, tok, browser_name)
            return

    # 4. 全部未命中时给出清晰的可操作提示（保留 Firefox 关键字以兼容既有自动化断言）
    if args.browser == "auto" and firefox_databases():
        print("没找到 skillhub_token——Firefox / Chrome / Edge 登录态过期，请重新登录一次", file=sys.stderr)
        sys.exit(1)
    raise SystemExit(
        "没找到 Firefox / Chrome / Edge 的 primehub localStorage"
        "（请先在任意主流浏览器登录一次 https://primehub.primebot.com，或使用 --token 传入）"
    )


if __name__ == "__main__":
    main()