"""通过 salt 定位快速探测内存中的 SQLCipher enc_key。

策略：
1. 读取加密 DB 的 salt
2. 在微信进程内存中搜索该 salt
3. 对每个 salt 命中点，在 ±2KB 范围内枚举 32 字节候选
4. 用 HMAC-SHA512 验证候选是否为 enc_key
"""

from __future__ import annotations

import hashlib
import hmac
import os
import struct
import sys
from typing import Optional

PAGE_SZ = 4096
KEY_SZ = 32
SALT_SZ = 16
RESERVE_SIZE = 80
IV_SIZE = 16
HMAC_SIZE = 64
MAC_SALT_XOR = 0x3A


def _verify_enc_key(enc_key: bytes, page1: bytes) -> bool:
    salt = page1[:SALT_SZ]
    mac_salt = bytes(b ^ MAC_SALT_XOR for b in salt)
    mac_key = hashlib.pbkdf2_hmac("sha512", enc_key, mac_salt, 2, dklen=KEY_SZ)
    hmac_data = page1[SALT_SZ: PAGE_SZ - RESERVE_SIZE + IV_SIZE]
    stored_hmac = page1[PAGE_SZ - HMAC_SIZE: PAGE_SZ]
    hm = hmac.new(mac_key, hmac_data, hashlib.sha512)
    hm.update(struct.pack("<I", 1))
    return hm.digest() == stored_hmac


def _get_regions(pid: int) -> list[tuple[int, int, str]]:
    regions = []
    try:
        with open(f"/proc/{pid}/maps", "r", encoding="utf-8", errors="replace") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 2:
                    continue
                perms = parts[1]
                if "r" not in perms:
                    continue
                name = parts[5] if len(parts) >= 6 else ""
                if name in ("[vdso]", "[vsyscall]", "[vvar]"):
                    continue
                start_s, end_s = parts[0].split("-")
                start = int(start_s, 16)
                end = int(end_s, 16)
                regions.append((start, end - start, name))
    except Exception as e:
        print(f"读取 maps 失败: {e}")
    return regions


def _read_chunk(mem, base: int, offset: int, size: int) -> bytes:
    try:
        mem.seek(base + offset)
        return mem.read(size)
    except Exception:
        return b""


def probe_pid(pid: int, salt: bytes, page1: bytes, window: int = 2048) -> Optional[bytes]:
    regions = _get_regions(pid)
    print(f"PID {pid}: {len(regions)} 个可读区域，salt={salt.hex()}")
    found = 0
    with open(f"/proc/{pid}/mem", "rb") as mem:
        for base, size, name in regions:
            pos = 0
            chunk_size = 4 * 1024 * 1024
            while pos < size:
                to_read = min(chunk_size, size - pos)
                data = _read_chunk(mem, base, pos, to_read)
                if not data:
                    break
                idx = 0
                while True:
                    hit = data.find(salt, idx)
                    if hit < 0:
                        break
                    abs_addr = base + pos + hit
                    found += 1
                    if found % 100 == 0:
                        print(f"  已找到 {found} 个 salt 命中...")
                    # 读取 ±window 区域
                    start = max(0, hit - window)
                    end = min(len(data), hit + SALT_SZ + window)
                    region = data[start:end]
                    # 枚举 32 字节候选（16 字节对齐）
                    for cand_off in range(0, len(region) - KEY_SZ + 1):
                        cand = region[cand_off:cand_off + KEY_SZ]
                        if _verify_enc_key(cand, page1):
                            print(f"  !! 找到 enc_key @ PID={pid} addr=0x{abs_addr - window + cand_off:016x} name={name}")
                            return cand
                    idx = hit + 1
                pos += len(data)
    print(f"PID {pid}: salt 命中 {found} 次，未找到 enc_key")
    return None


def main():
    db_path = sys.argv[1] if len(sys.argv) > 1 else "/config/xwechat_files/wxid_zfj3oc1hn5no22_43c4/db_storage/contact/contact.db"
    pid = int(sys.argv[2]) if len(sys.argv) > 2 else 348

    with open(db_path, "rb") as f:
        page1 = f.read(PAGE_SZ)
    salt = page1[:SALT_SZ]
    print(f"DB: {db_path}, salt={salt.hex()}")

    key = probe_pid(pid, salt, page1)
    if key:
        print(f"enc_key={key.hex()}")
        return 0

    print("未找到 enc_key")
    return 1


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