WechatOnCloud/bridge/tools/extract_key_standalone.py

"""独立的 Linux 微信进程内存 key 提取工具（调试用）。

用法（在容器内执行）：
    python3 /opt/woc-bridge/tools/extract_key_standalone.py

设计目标：
- 与 woc-bridge 主服务隔离，扫描失败不影响 bridge
- 自动检测所有微信进程并按内存大小排序扫描
- 匹配 x'<hex>' 模式提取 enc_key
- 输出 salt -> enc_key 映射，供 bridge 注入使用

注：生产环境请通过 /api/db/init 接口触发自动提取，本脚本仅用于调试与故障排查。
"""

from __future__ import annotations

import argparse
import logging
import os
import sys

# 添加父目录到 sys.path，便于从 tools/ 子目录 import 父级 woc_bridge 包
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from woc_bridge.db.key_extractor import KeyExtractor


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
_logger = logging.getLogger("extract_key")


def _find_encrypted_db_dir() -> str:
    """尝试自动定位微信数据根目录。"""
    db_root = "/config"
    if os.path.isdir(db_root):
        return db_root
    return "/"


def main() -> int:
    parser = argparse.ArgumentParser(description="从 Linux 微信进程内存提取 SQLCipher DB enc_key")
    parser.add_argument("--pid", type=int, default=0, help="指定微信进程 PID（不指定则自动检测）")
    parser.add_argument("--db", type=str, default="", help="微信 DB 文件或目录（不指定则自动查找 /config）")
    parser.add_argument("--full", action="store_true", help="完整扫描所有可读内存区域（默认优先堆）")
    args = parser.parse_args()

    db_path = args.db or _find_encrypted_db_dir()
    _logger.info("目标 DB: %s", db_path)

    extractor = KeyExtractor(args.pid, db_path)
    key_map = extractor.extract_all_keys(full_scan=args.full)

    if not key_map:
        print("\n[-] 未能提取到任何有效 enc_key")
        return 1

    print(f"\n[+] 成功提取 {len(key_map)} 个 enc_key:")
    print("-" * 80)
    for salt_hex, enc_key_hex in key_map.items():
        print(f"salt:  {salt_hex}")
        print(f"key:   {enc_key_hex}")
        print(f"       前 8 位: {enc_key_hex[:8]}... 后 8 位: ...{enc_key_hex[-8:]}")
        print("-" * 80)

    # 如果只有一个 key，同时输出可直接注入 bridge 的 curl 命令
    if len(key_map) == 1:
        enc_key_hex = next(iter(key_map.values()))
        print("\n[+] 可通过以下命令注入 bridge:")
        print(
            f"    curl -X POST http://localhost:8088/api/db/decrypt "
            f"-H 'Content-Type: application/json' -d '{{\"key\":\"{enc_key_hex}\"}}'"
        )
    else:
        print("\n[+] 多个 salt 时，建议直接启用 bridge 自动提取（WOC_DB_KEY_AUTO_EXTRACT=true）")

    return 0


if __name__ == "__main__":
    sys.exit(main())