# extract-weixin-key.ps1
# 输入参数: -SeedDir <path>  (可选, 默认自动找 Documents/xwechat_files/ 最近目录)
# 输出 (stdout): JSON { wxid: ..., dbKeys: { "message_0.db": "<hex>", ... } }
# 退出码: 0 OK / 1 无 Weixin 进程 / 2 seed 目录不存在 / 3 内存读取失败

param([string]$SeedDir = "")

$ErrorActionPreference = 'Stop'

# 1. 找 Weixin 主进程
$weixin = Get-Process Weixin -ErrorAction SilentlyContinue | Where-Object { $_.MainWindowHandle -ne 0 } | Select-Object -First 1
if (-not $weixin) { Write-Error "No Weixin process"; exit 1 }

# 2. 确定 seed 目录
if (-not $SeedDir) {
    $root = Join-Path ([Environment]::GetFolderPath('MyDocuments')) 'xwechat_files'
    if (-not (Test-Path $root)) { Write-Error "xwechat_files not found"; exit 2 }
    $SeedDir = Get-ChildItem $root -Directory | Where-Object { $_.Name -ne 'all_users' -and $_.Name -ne 'Backup' } |
        Sort-Object LastWriteTime -Descending | Select-Object -First 1 -ExpandProperty FullName
}

# 3. 收集各 DB 的 salt (前 16B)
$dbStorage = Join-Path $SeedDir 'db_storage'
$dbFiles = Get-ChildItem $dbStorage -Recurse -Filter '*.db' -File |
    Where-Object { $_.Name -notmatch '-wal$|-shm$' }
$saltMap = @{}
foreach ($f in $dbFiles) {
    $bytes = New-Object byte[] 16
    try {
        # Weixin 运行时持有 DB 句柄 (exclusive write), 必须 FileShare=ReadWrite 才能并发读
        $fs = [System.IO.File]::Open($f.FullName, [System.IO.FileMode]::Open, [System.IO.FileAccess]::Read, [System.IO.FileShare]::ReadWrite)
        try { $fs.Read($bytes, 0, 16) | Out-Null } finally { $fs.Close() }
        $saltMap[$f.Name] = [BitConverter]::ToString($bytes).Replace('-', '').ToLowerInvariant()
    } catch {
        # 个别 DB 读不到不算致命 (例如临时锁), 跳过该 DB; 主流程仍可工作
        continue
    }
}

# 4. P/Invoke — 枚举 + 读内存
Add-Type -TypeDefinition @"
using System;
using System.Runtime.InteropServices;
public static class W {
    [StructLayout(LayoutKind.Sequential)] public struct MBI {
        public IntPtr BaseAddress; public IntPtr AllocationBase;
        public uint AllocationProtect; public IntPtr RegionSize;
        public uint State; public uint Protect; public uint Type;
    }
    [DllImport("kernel32.dll", SetLastError=true)] public static extern IntPtr OpenProcess(uint a, bool i, int p);
    [DllImport("kernel32.dll", SetLastError=true)] public static extern bool CloseHandle(IntPtr h);
    [DllImport("kernel32.dll", SetLastError=true)] public static extern int VirtualQueryEx(IntPtr h, IntPtr a, out MBI m, uint l);
    [DllImport("kernel32.dll", SetLastError=true)] public static extern bool ReadProcessMemory(IntPtr h, IntPtr a, byte[] b, IntPtr s, out IntPtr r);
}
"@

$handle = [W]::OpenProcess(0x410, $false, $weixin.Id)
if ($handle -eq [IntPtr]::Zero) { Write-Error "OpenProcess failed"; exit 3 }

# 扫描 + 提取 96-char hex literal
$addr = [IntPtr]::Zero
$max = 0x7FFFFFFEFFFF
$regex = [regex]::new("x'([a-f0-9]{96})'", 'IgnoreCase')
$literals = New-Object System.Collections.Generic.HashSet[string]
$latin1 = [System.Text.Encoding]::GetEncoding('iso-8859-1')
try {
    while ([Int64]$addr -lt $max) {
        $mbi = New-Object W+MBI
        if ([W]::VirtualQueryEx($handle, $addr, [ref]$mbi, 48) -eq 0) { break }
        $isCommit = ($mbi.State -band 0x1000) -ne 0
        $isPrivate = ($mbi.Type -band 0x20000) -ne 0
        $isRW = ($mbi.Protect -eq 0x04) -or ($mbi.Protect -eq 0x40)
        $noGuard = ($mbi.Protect -band 0x100) -eq 0
        if ($isCommit -and $isPrivate -and $isRW -and $noGuard) {
            $size = [Int64]$mbi.RegionSize
            if ($size -gt 0 -and $size -lt 1GB) {
                $buf = New-Object byte[] $size
                $read = [IntPtr]::Zero
                $ok = [W]::ReadProcessMemory($handle, $mbi.BaseAddress, $buf, [IntPtr]::new($size), [ref]$read)
                if ($ok) {
                    $text = $latin1.GetString($buf, 0, [Int64]$read)
                    foreach ($m in $regex.Matches($text)) {
                        [void]$literals.Add($m.Groups[1].Value.ToLowerInvariant())
                    }
                }
            }
        }
        $addr = [IntPtr]::new([Int64]$mbi.BaseAddress + [Int64]$mbi.RegionSize)
    }
} finally { [W]::CloseHandle($handle) | Out-Null }

# 5. 反向匹配 salt → key
$dbKeys = @{}
foreach ($lit in $literals) {
    $keyHex = $lit.Substring(0, 64)
    $saltHex = $lit.Substring(64)
    foreach ($entry in $saltMap.GetEnumerator()) {
        if ($entry.Value -eq $saltHex) {
            $dbKeys[$entry.Key] = $keyHex
            break
        }
    }
}

# 6. 输出 JSON
$result = @{
    seedDir       = $SeedDir
    wxid          = (Split-Path $SeedDir -Leaf)
    wechatVersion = $weixin.MainModule.FileVersionInfo.FileVersion
    pid           = $weixin.Id
    dbKeys        = $dbKeys
}
$result | ConvertTo-Json -Compress