"""在 lldb 上下文中运行，动态枚举内存区域搜索密钥"""
import hashlib, hmac, struct, os

DB_BASE = '/Users/yibo/Library/Containers/com.tencent.xinWeChat/Data/Documents/xwechat_files/yibooo_ac07/db_storage/'

def load_header(relpath):
    path = DB_BASE + relpath
    if os.path.exists(path):
        with open(path, 'rb') as f:
            return f.read(4096)
    return None

def make_verifier(header):
    salt = header[:16]
    mac_salt = bytes(b ^ 0x3A for b in salt)
    stored = header[4064:4096]
    content = header[16:4048]
    iv = header[4048:4064]
    def verify_raw(key):
        if len(set(key)) < 10:
            return False
        hk = hashlib.pbkdf2_hmac('sha256', key, mac_salt, 2, dklen=32)
        h = hmac.new(hk, digestmod='sha256')
        h.update(content)
        h.update(iv)
        h.update(struct.pack('<I', 1))
        return h.digest() == stored
    return verify_raw

# 加载多个数据库
verifiers = {}
for name, path in [('contact', 'contact/contact.db'),
                   ('msg0',    'message/message_0.db'),
                   ('general', 'general/general.db')]:
    h = load_header(path)
    if h:
        verifiers[name] = (bytes.fromhex(
            {'contact': 'f8e768b1b996c8b84441eea996b2335e',
             'msg0':    'c6e5830c56967341a101f3c756202a3e',
             'general': 'd560d66c5fcbc3ae406939ef178ffcc9'}[name]
        ), make_verifier(h))

process = lldb.debugger.GetSelectedTarget().GetProcess()
if not process or not process.IsValid():
    print("[ERR] 无有效进程")
else:
    print(f"[OK] PID={process.GetProcessID()}, state={process.GetState()}")

    addr = 0
    total_scanned = 0
    found = False
    region_count = 0
    readable_count = 0

    while not found:
        region_info = lldb.SBMemoryRegionInfo()
        err = process.GetMemoryRegionInfo(addr, region_info)
        if not err.Success():
            break

        next_addr = region_info.GetRegionEnd()
        if next_addr <= addr:
            break

        region_size = next_addr - addr

        # 只处理可读写、非可执行的区域
        if region_info.IsReadable() and region_info.IsWritable() and not region_info.IsExecutable():
            readable_count += 1

            # 分块读取（每块 4MB）
            chunk_size = 4 * 1024 * 1024
            for chunk_off in range(0, region_size, chunk_size):
                read_sz = min(chunk_size, region_size - chunk_off)
                err2 = lldb.SBError()
                data = process.ReadMemory(addr + chunk_off, read_sz, err2)
                if err2.Fail() or not data:
                    continue

                total_scanned += len(data)

                # 搜索 salt 字节
                for dbname, (salt, verify) in verifiers.items():
                    idx = 0
                    while True:
                        pos = data.find(salt, idx)
                        if pos == -1:
                            break
                        idx = pos + 1
                        print(f"[SALT] {dbname} @ {addr+chunk_off+pos:#x}")
                        # 在 salt 附近 ±1KB 搜索密钥
                        s = max(0, pos - 1024)
                        e = min(len(data), pos + 16 + 1024)
                        for i in range(s, e - 32, 1):
                            cand = data[i:i+32]
                            if verify(cand):
                                print(f"[KEY] {dbname}={cand.hex()}")
                                found = True
                                with open('/tmp/wechat_key.hex', 'w') as f:
                                    f.write(cand.hex())
                                break
                    if found:
                        break
                if found:
                    break

            # 进度
            region_count += 1
            if region_count % 50 == 0:
                print(f"[PROG] {region_count} regions, {total_scanned//1024//1024}MB, readable={readable_count}")

        addr = next_addr
        if addr > 0x7fffffffffff:  # user space limit
            break

    print(f"[DONE] regions={region_count}, readable={readable_count}, scanned={total_scanned//1024//1024}MB, found={found}")
    if not found:
        print("[FAIL] Salt not found in any memory region")