"""
企业微信 (WXWork) 进程内存密钥提取模块
通过 Windows API 读取 WXWork.exe 进程内存, 扫描并提取 wxSQLite3 AES-128 原始密钥

策略:
1. 找到所有 WXWork 相关进程
2. 枚举可读内存区域
3. 搜索已知模式 (cipher 结构体特征 / "sAlT" / "aes128" 字符串)
4. 从匹配位置附近提取 16 字节候选密钥
5. 用数据库第 1 页验证候选密钥
"""

import ctypes
import ctypes.wintypes as wintypes
import hashlib
import os
import struct
import sys
import time

# ---- Windows API 常量 ----
PROCESS_VM_READ = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400
TH32CS_SNAPPROCESS = 0x00000002

MEM_COMMIT = 0x1000
PAGE_READWRITE = 0x04
PAGE_READONLY = 0x02
PAGE_EXECUTE_READWRITE = 0x40
PAGE_EXECUTE_READ = 0x20

READABLE_PROTECTIONS = {
    PAGE_READWRITE,
    PAGE_READONLY,
    PAGE_EXECUTE_READWRITE,
    PAGE_EXECUTE_READ,
}


class MEMORY_BASIC_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("BaseAddress", ctypes.c_void_p),
        ("AllocationBase", ctypes.c_void_p),
        ("AllocationProtect", wintypes.DWORD),
        ("PartitionId", wintypes.WORD),
        ("RegionSize", ctypes.c_size_t),
        ("State", wintypes.DWORD),
        ("Protect", wintypes.DWORD),
        ("Type", wintypes.DWORD),
    ]


class PROCESSENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes.DWORD),
        ("cntUsage", wintypes.DWORD),
        ("th32ProcessID", wintypes.DWORD),
        ("th32DefaultHeapID", ctypes.POINTER(ctypes.c_ulong)),
        ("th32ModuleID", wintypes.DWORD),
        ("cntThreads", wintypes.DWORD),
        ("th32ParentProcessID", wintypes.DWORD),
        ("pcPriClassBase", ctypes.c_long),
        ("dwFlags", wintypes.DWORD),
        ("szExeFile", ctypes.c_char * 260),
    ]


def get_wxwork_pids():
    """查找所有 WXWork 相关进程 ID"""
    pids = []
    kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)

    CreateToolhelp32Snapshot = kernel32.CreateToolhelp32Snapshot
    CreateToolhelp32Snapshot.argtypes = [wintypes.DWORD, wintypes.DWORD]
    CreateToolhelp32Snapshot.restype = wintypes.HANDLE

    Process32First = kernel32.Process32First
    Process32First.argtypes = [wintypes.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
    Process32First.restype = wintypes.BOOL

    Process32Next = kernel32.Process32Next
    Process32Next.argtypes = [wintypes.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
    Process32Next.restype = wintypes.BOOL

    CloseHandle = kernel32.CloseHandle
    CloseHandle.argtypes = [wintypes.HANDLE]

    snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
    if snapshot == -1:
        return pids

    pe = PROCESSENTRY32()
    pe.dwSize = ctypes.sizeof(PROCESSENTRY32)

    if Process32First(snapshot, ctypes.byref(pe)):
        while True:
            exe_name = pe.szExeFile.lower()
            if b'wxwork' in exe_name or b'wework' in exe_name:
                pids.append(pe.th32ProcessID)
            if not Process32Next(snapshot, ctypes.byref(pe)):
                break

    CloseHandle(snapshot)
    return pids


def read_process_memory_regions(pid, max_total=512 * 1024 * 1024):
    """读取进程的可读内存区域, 返回 [(address, data), ...]"""
    kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)

    OpenProcess = kernel32.OpenProcess
    OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
    OpenProcess.restype = wintypes.HANDLE

    VirtualQueryEx = kernel32.VirtualQueryEx
    VirtualQueryEx.argtypes = [wintypes.HANDLE, ctypes.c_void_p,
                               ctypes.POINTER(MEMORY_BASIC_INFORMATION), ctypes.c_size_t]
    VirtualQueryEx.restype = ctypes.c_size_t

    ReadProcessMemory = kernel32.ReadProcessMemory
    ReadProcessMemory.argtypes = [wintypes.HANDLE, ctypes.c_void_p,
                                  ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)]
    ReadProcessMemory.restype = wintypes.BOOL

    CloseHandle = kernel32.CloseHandle
    CloseHandle.argtypes = [wintypes.HANDLE]

    h_process = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
    if not h_process:
        return []

    regions = []
    total_read = 0
    addr = 0
    mbi = MEMORY_BASIC_INFORMATION()

    try:
        while addr < 0x7FFFFFFF0000:  # 用户空间上限
            if VirtualQueryEx(h_process, ctypes.c_void_p(addr), ctypes.byref(mbi),
                              ctypes.sizeof(mbi)) == 0:
                break

            # c_void_p 值为 0 时在 Python 中解析为 None, 需转成整数
            base_addr = mbi.BaseAddress
            if base_addr is None:
                base_addr = 0
            base_addr = int(base_addr)

            size = mbi.RegionSize
            if size > 0 and mbi.State == MEM_COMMIT and mbi.Protect in READABLE_PROTECTIONS:
                # 限制单块读取大小
                read_size = min(size, 64 * 1024 * 1024)
                buffer = ctypes.create_string_buffer(read_size)
                bytes_read = ctypes.c_size_t(0)

                if ReadProcessMemory(h_process, ctypes.c_void_p(base_addr),
                                     buffer, read_size, ctypes.byref(bytes_read)):
                    data = buffer.raw[:bytes_read.value]
                    if len(data) > 0:
                        regions.append((base_addr, data))
                        total_read += len(data)
                        if total_read >= max_total:
                            break

            addr = base_addr + size
            if addr <= 0:
                break
    finally:
        CloseHandle(h_process)

    return regions


def find_key_candidates_in_data(data, base_addr=0):
    """
    在内存数据中搜索 wxSQLite3 cipher 结构体特征, 提取候选密钥
    返回候选密钥列表 [(key_bytes, address, reason), ...]
    """
    candidates = []

    # 模式 1: 搜索 "sAlT" 字符串, 其前面的 20 字节可能是 [key(16) + page_no(4)]
    salt_offset = 0
    while True:
        idx = data.find(b'sAlT', salt_offset)
        if idx == -1:
            break
        # 检查前 20 字节: 可能是 raw_key + page_no(LE)
        if idx >= 20:
            before = data[idx - 20:idx]
            key_candidate = before[:16]
            page_no_bytes = before[16:20]
            # page_no 应该是一个小整数 (1, 2, 3, ...)
            page_no = struct.unpack("<I", page_no_bytes)[0]
            if 0 <= page_no <= 10000:
                candidates.append((key_candidate, base_addr + idx - 20,
                                   f"sAlT附近 page_no={page_no}"))
        # 也检查前面 16 字节 (直接是 key)
        if idx >= 16:
            key_candidate = data[idx - 16:idx]
            candidates.append((key_candidate, base_addr + idx - 16, "sAlT前16字节"))
        salt_offset = idx + 4

    # 模式 2: 搜索 "aes128cbc" / "aes128" / "AES-128" 字符串
    for pattern in [b'aes128cbc', b'aes-128-cbc', b'AES-128-CBC', b'aes128',
                    b'wxSQLite3', b'wxsqlite3']:
        pat_offset = 0
        while True:
            idx = data.find(pattern, pat_offset)
            if idx == -1:
                break
            # 在模式附近搜索 16 字节密钥
            start = max(0, idx - 256)
            end = min(len(data), idx + len(pattern) + 256)
            nearby = data[start:end]
            # 提取所有可能的 16 字节候选 (步进 4 对齐以减少噪音)
            for i in range(0, len(nearby) - 16, 4):
                cand = nearby[i:i + 16]
                # 过滤: 密钥不应包含大量 0x00
                zero_ratio = cand.count(0) / 16
                if zero_ratio < 0.5:
                    candidates.append((cand, base_addr + start + i, f"{pattern}附近"))
            pat_offset = idx + len(pattern)

    # 模式 3: 搜索 16 字节长度标记 + 指针结构
    # wxSQLite3 cipher 结构通常包含: int nKey = 16; void *zKey = ptr; void *db = ptr;
    # 搜索 "10000000" (LE uint32 16) 后面跟指针, 指针指向的地址可能存有密钥
    # 这个模式比较复杂, 暂时用简单方式: 搜索 16 字节序列中连续出现 0x10 0x00 0x00 0x00
    idx16 = 0
    while True:
        idx = data.find(b'\x10\x00\x00\x00', idx16)
        if idx == -1:
            break
        # 检查后面 8 字节是否像指针 (高 4 字节 0x00 或 0x7F)
        if idx + 12 <= len(data):
            ptr_bytes = data[idx + 4:idx + 12]
            # 检查 4-8 字节的指针值
            ptr_val = struct.unpack("<Q", ptr_bytes)[0] if len(ptr_bytes) == 8 else 0
            if 0x10000 < ptr_val < 0x7FFFFFFF0000:
                candidates.append((data[idx:idx + 16] if idx + 16 <= len(data) else b'',
                                   base_addr + idx, "nKey=16结构"))
        idx16 = idx + 4

    # 模式 4: 搜索 SQLite 相关字符串附近
    for pattern in [b'SQLite format', b'message.db', b'session.db']:
        pat_offset = 0
        while True:
            idx = data.find(pattern, pat_offset)
            if idx == -1:
                break
            start = max(0, idx - 128)
            end = min(len(data), idx + len(pattern) + 128)
            nearby = data[start:end]
            for i in range(0, len(nearby) - 16, 8):
                cand = nearby[i:i + 16]
                zero_ratio = cand.count(0) / 16
                if zero_ratio < 0.5:
                    candidates.append((cand, base_addr + start + i, f"{pattern}附近"))
            pat_offset = idx + len(pattern)

    # 模式 5: 搜索精确的密钥派生材料 [key(16) + page_no(4) + "sAlT"]
    # 页面 1 的派生材料: raw_key + b'\x01\x00\x00\x00' + b'sAlT'
    page1_material = b'\x01\x00\x00\x00sAlT'
    mat_offset = 0
    while True:
        idx = data.find(page1_material, mat_offset)
        if idx == -1:
            break
        if idx >= 16:
            key_candidate = data[idx - 16:idx]
            # 过滤: 密钥不应全零或全相同
            if len(set(key_candidate)) > 4:
                candidates.append((key_candidate, base_addr + idx - 16,
                                   "page1派生材料: key+page_no+sAlT"))
        mat_offset = idx + len(page1_material)

    # 模式 6: wxSQLite3 cipher 结构体检测
    # 结构体通常包含: [ptr] [nKey=16 LE] [ptr zKey] [nIV=16 LE] [ptr] ...
    # 搜索 nKey=16 后面跟指针, 然后指针指向的地址附近可能有密钥
    # 在数据中搜索 16 字节的密钥模式: 高熵 + 指针结构
    # 搜索特征: \x10\x00\x00\x00 + 指针 (8字节) + \x10\x00\x00\x00
    struct_offset = 0
    while True:
        idx = data.find(b'\x10\x00\x00\x00', struct_offset)
        if idx == -1:
            break
        # 检查后面是否跟指针
        if idx + 12 <= len(data):
            ptr_bytes = data[idx + 4:idx + 12]
            ptr_val = struct.unpack("<Q", ptr_bytes)[0]
            # 指针指向用户空间
            if 0x10000 < ptr_val < 0x7FFFFFFF0000:
                # 检查是否还有第二个 nKey=16
                if idx + 16 <= len(data) and data[idx + 12:idx + 16] == b'\x10\x00\x00\x00':
                    # 这是一个 cipher 结构体! 在结构体前后的内存中搜索密钥
                    start = max(0, idx - 512)
                    end = min(len(data), idx + 512)
                    nearby = data[start:end]
                    for i in range(0, len(nearby) - 16, 4):
                        cand = nearby[i:i + 16]
                        zero_ratio = cand.count(0) / 16
                        if zero_ratio < 0.3:  # 密钥通常不会是大量零
                            candidates.append((cand, base_addr + start + i,
                                               "cipher结构体: nKey=16+ptr+nKey=16"))
        struct_offset = idx + 4

    # 模式 7: 搜索密钥可能所在的堆内存 (特征: 16字节高熵数据 + 附近有常见字符串)
    # wxSQLite3 的密钥通常会以特定格式存在, 如连续的 0x10 (16) + 密钥
    # 搜索: \x10\x00\x00\x00 后跟 16 字节高熵数据
    key_pat_offset = 0
    while True:
        idx = data.find(b'\x10\x00\x00\x00', key_pat_offset)
        if idx == -1:
            break
        if idx + 20 <= len(data):
            key_candidate = data[idx + 4:idx + 20]
            # 检查是否像密钥 (非文本, 非全零)
            zero_ratio = key_candidate.count(0) / 16
            printable = sum(1 for b in key_candidate if 32 <= b <= 126) / 16
            if zero_ratio < 0.4 and printable < 0.7:
                candidates.append((key_candidate, base_addr + idx + 4,
                                   "nKey=16后密钥缓冲区"))
        key_pat_offset = idx + 4

    return candidates


def verify_key(raw_key, page1):
    """验证密钥: 解密第 1 页并检查是否为合法 SQLite 页面"""
    from wxwork_crypto import verify_wxsqlite3_aes128_key
    return verify_wxsqlite3_aes128_key(raw_key, page1)


def extract_wxwork_key(db_page1_path=None, db_base=None):
    """
    从企业微信进程内存提取 raw key
    返回: (key_bytes, pid, 来源描述) 或 None
    db_base: 企业微信数据目录 (WXWork), None 使用默认 Documents\\WXWork
    """
    # 准备验证用页面数据
    if db_base is None:
        db_base = os.path.join(os.path.expanduser("~"), "Documents", "WXWork")
    if db_page1_path is None:
        # 自动查找第一个加密数据库
        wxwork_base = db_base
        page1 = None
        if os.path.isdir(wxwork_base):
            for d in sorted(os.listdir(wxwork_base)):
                msg_path = os.path.join(wxwork_base, d, "Data", "message.db")
                if os.path.exists(msg_path):
                    with open(msg_path, "rb") as f:
                        page1 = f.read(4096)
                    if page1 and page1[:16] != b'SQLite format 3\x00':
                        break
                    page1 = None
    else:
        with open(db_page1_path, "rb") as f:
            page1 = f.read(4096)

    if not page1 or page1[:16] == b'SQLite format 3\x00':
        print("[-] 未找到加密的数据库文件")
        return None

    pids = get_wxwork_pids()
    if not pids:
        print("[-] 未找到运行中的企业微信进程, 请先启动并登录企业微信")
        return None

    print(f"[+] 找到企业微信进程: {pids}")
    print("[*] 正在读取进程内存并扫描密钥...")

    seen_keys = set()
    found_key = None
    found_info = None

    for pid in pids:
        regions = read_process_memory_regions(pid)
        print(f"  PID {pid}: 读取了 {len(regions)} 个内存区域")
        for base_addr, data in regions:
            candidates = find_key_candidates_in_data(data, base_addr)
            for key_cand, addr, reason in candidates:
                if len(key_cand) != 16:
                    continue
                key_bytes = bytes(key_cand)
                if key_bytes in seen_keys:
                    continue
                seen_keys.add(key_bytes)

                if verify_key(key_bytes, page1):
                    found_key = key_bytes
                    found_info = f"PID={pid}, 地址=0x{addr:X}, 来源={reason}"
                    print(f"\n[+] 找到有效密钥! {found_info}")
                    print(f"    Key: {found_key.hex()}")
                    return found_key, pid, found_info

    print("[-] 未能从进程内存中找到密钥")
    return None


if __name__ == "__main__":
    result = extract_wxwork_key()
    if result:
        key, pid, info = result
        print(f"\n密钥提取成功!")
        print(f"  Raw Key: {key.hex()}")
        print(f"  来源: {info}")
    else:
        print("\n密钥提取失败, 请确保:")
        print("  1. 企业微信已启动并登录")
        print("  2. 以管理员权限运行此脚本")