wx-win-agent/node-agent/WxAgent.Windows/ProcessMemoryScanner.cs

using System.Buffers.Binary;
using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Text;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.System.Memory;
using Windows.Win32.System.Threading;
using WxAgent.Core;

namespace WxAgent.Windows;

public sealed record ProcessAccessProbe(bool Success, string? Error);

public static unsafe class ProcessMemoryScanner
{
    private const int ChunkSize = 1024 * 1024;
    private const ulong MaximumRegionSize = 500UL * 1024 * 1024;
    private const ulong MaximumUserAddress = 0x00007FFFFFFEFFFF;

    public static ProcessAccessProbe Probe(int processId)
    {
        var process = PInvoke.OpenProcess(PROCESS_ACCESS_RIGHTS.PROCESS_QUERY_INFORMATION | PROCESS_ACCESS_RIGHTS.PROCESS_VM_READ, false, (uint)processId);
        if (process.IsNull)
        {
            return new ProcessAccessProbe(false, new Win32Exception(Marshal.GetLastWin32Error()).Message);
        }

        PInvoke.CloseHandle(process);
        return new ProcessAccessProbe(true, null);
    }

    public static IReadOnlyList<WcdbKeyCandidate> Scan(int processId, CancellationToken cancellationToken)
    {
        var process = PInvoke.OpenProcess(PROCESS_ACCESS_RIGHTS.PROCESS_QUERY_INFORMATION | PROCESS_ACCESS_RIGHTS.PROCESS_VM_READ, false, (uint)processId);
        if (process.IsNull)
        {
            throw new WxAgentException(WxAgentErrorCode.ProcessAccessDenied, $"Cannot read Weixin process {processId}.", new Win32Exception(Marshal.GetLastWin32Error()));
        }

        try
        {
            var found = new List<WcdbKeyCandidate>(ScanLegacyHex(process, cancellationToken));
            var seen = new HashSet<WcdbKeyCandidate>(found, WcdbKeyCandidateComparer.Instance);
            foreach (var candidate in ScanConfigCipher(process, cancellationToken))
            {
                if (seen.Add(candidate)) found.Add(candidate);
            }

            return found;
        }
        finally
        {
            PInvoke.CloseHandle(process);
        }
    }

    private static IReadOnlyList<WcdbKeyCandidate> ScanLegacyHex(HANDLE process, CancellationToken cancellationToken)
    {
        WcdbKeyCandidateScanner[] scanners = [new(), new(utf16: true)];
        var found = new List<WcdbKeyCandidate>();
        ulong? contiguousEnd = null;
        ulong address = 0;
        while (address < MaximumUserAddress)
        {
            cancellationToken.ThrowIfCancellationRequested();
            MEMORY_BASIC_INFORMATION region;
            if (PInvoke.VirtualQueryEx(process, (void*)address, &region, (nuint)sizeof(MEMORY_BASIC_INFORMATION)) == 0)
            {
                break;
            }

            var baseAddress = (ulong)region.BaseAddress;
            var regionSize = (ulong)region.RegionSize;
            if (regionSize == 0 || baseAddress > ulong.MaxValue - regionSize)
            {
                break;
            }

            var next = baseAddress + regionSize;
            if (next <= address)
            {
                break;
            }

            if (!IsReadable(region) || regionSize > MaximumRegionSize)
            {
                BreakRun(scanners, found);
                contiguousEnd = null;
            }
            else
            {
                if (contiguousEnd != baseAddress)
                {
                    BreakRun(scanners, found);
                }

                contiguousEnd = ReadRegion(process, baseAddress, regionSize, scanners, found, cancellationToken)
                    ? next
                    : null;
            }

            address = next;
        }

        BreakRun(scanners, found);
        return found;
    }

    private static IReadOnlyList<WcdbKeyCandidate> ScanConfigCipher(HANDLE process, CancellationToken cancellationToken)
    {
        var typeName = Encoding.ASCII.GetBytes(WcdbConfigCipher.TypeName);
        var needles = new HashSet<ulong>();
        WalkRegions(process, cancellationToken, (buffer, length, address) =>
        {
            using var scanner = new MemoryPatternScanner(typeName);
            foreach (var match in scanner.Feed(buffer.AsSpan(0, length), address)) needles.Add(match);
        });
        if (needles.Count == 0) return [];

        var nodeBases = new HashSet<ulong>();
        foreach (var needle in needles)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var pair = new byte[16];
            BinaryPrimitives.WriteUInt64LittleEndian(pair.AsSpan(0, 8), needle);
            BinaryPrimitives.WriteUInt64LittleEndian(pair.AsSpan(8, 8), (ulong)typeName.Length);
            using var pairScanner = new MemoryPatternScanner(pair);
            WalkRegions(process, cancellationToken, (buffer, length, address) =>
            {
                foreach (var match in pairScanner.Feed(buffer.AsSpan(0, length), address)) nodeBases.Add(match - 0x10);
            });
        }

        var found = new List<WcdbKeyCandidate>();
        foreach (var nodeBase in nodeBases)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var node = ReadBytes(process, nodeBase, 0x50);
            if (node is null) continue;
            try
            {
                var stringPointer = ReadUInt64(node, 0x10);
                var stringLength = ReadUInt64(node, 0x18);
                if (!needles.Contains(stringPointer) || stringLength != (ulong)typeName.Length) continue;
                var configPointer = ReadUInt64(node, 0x28);
                if (configPointer is < 0x10000 or > MaximumUserAddress) continue;
                var obj = ReadBytes(process, configPointer + 0x88, 0x28);
                if (obj is null) continue;
                var dataPointer = ReadUInt64(obj, 0x8);
                var dataLength = ReadUInt64(obj, 0x10);
                if (dataLength == 0 || dataLength > WcdbConfigCipher.MaximumBlobLength
                    || dataPointer is < 0x10000 or > MaximumUserAddress) continue;
                var blob = ReadBytes(process, dataPointer, (int)dataLength);
                if (blob is null || blob.Length != (int)dataLength) continue;
                foreach (var candidate in WcdbConfigCipher.Decode(blob)) found.Add(candidate);
            }
            finally
            {
                CryptographicOperations.ZeroMemory(node);
            }
        }

        return found;
    }

    private static void WalkRegions(HANDLE process, CancellationToken cancellationToken, Action<byte[], int, ulong> visit)
    {
        ulong address = 0;
        while (address < MaximumUserAddress)
        {
            cancellationToken.ThrowIfCancellationRequested();
            MEMORY_BASIC_INFORMATION region;
            if (PInvoke.VirtualQueryEx(process, (void*)address, &region, (nuint)sizeof(MEMORY_BASIC_INFORMATION)) == 0)
            {
                break;
            }

            var baseAddress = (ulong)region.BaseAddress;
            var regionSize = (ulong)region.RegionSize;
            if (regionSize == 0 || baseAddress > ulong.MaxValue - regionSize)
            {
                break;
            }

            var next = baseAddress + regionSize;
            if (next <= address)
            {
                break;
            }

            if (IsReadable(region) && regionSize <= MaximumRegionSize)
            {
                var offset = 0UL;
                while (offset < regionSize)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    var count = (int)Math.Min((ulong)ChunkSize, regionSize - offset);
                    var buffer = GC.AllocateUninitializedArray<byte>(count);
                    nuint bytesRead = 0;
                    var ok = false;
                    fixed (byte* destination = buffer)
                    {
                        ok = PInvoke.ReadProcessMemory(process, (void*)(baseAddress + offset), destination, (nuint)count, &bytesRead);
                    }

                    try
                    {
                        if (ok && bytesRead == (nuint)count)
                            visit(buffer, count, baseAddress + offset);
                    }
                    finally
                    {
                        CryptographicOperations.ZeroMemory(buffer);
                    }
                    if (!ok || bytesRead != (nuint)count) break;
                    offset += (ulong)count;
                }
            }

            address = next;
        }
    }

    private static byte[]? ReadBytes(HANDLE process, ulong address, int length)
    {
        if (length <= 0 || address > MaximumUserAddress - (ulong)length) return null;
        var buffer = GC.AllocateUninitializedArray<byte>(length);
        nuint bytesRead = 0;
        var ok = false;
        fixed (byte* destination = buffer)
        {
            ok = PInvoke.ReadProcessMemory(process, (void*)address, destination, (nuint)length, &bytesRead);
        }

        if (!ok || bytesRead != (nuint)length)
        {
            CryptographicOperations.ZeroMemory(buffer);
            return null;
        }

        return buffer;
    }

    private static ulong ReadUInt64(ReadOnlySpan<byte> buffer, int offset) =>
        offset + 8 <= buffer.Length ? BinaryPrimitives.ReadUInt64LittleEndian(buffer.Slice(offset, 8)) : 0;

    private static bool ReadRegion(
        HANDLE process,
        ulong baseAddress,
        ulong regionSize,
        IReadOnlyList<WcdbKeyCandidateScanner> scanners,
        List<WcdbKeyCandidate> found,
        CancellationToken cancellationToken)
    {
        var tailIsContiguous = false;
        var offset = 0UL;
        while (offset < regionSize)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var count = (int)Math.Min((ulong)ChunkSize, regionSize - offset);
            var buffer = GC.AllocateUninitializedArray<byte>(count);
            nuint bytesRead = 0;
            var success = false;
            fixed (byte* destination = buffer)
            {
                success = PInvoke.ReadProcessMemory(process, (void*)(baseAddress + offset), destination, (nuint)count, &bytesRead);
            }

            try
            {
                if (bytesRead > 0)
                {
                    var length = checked((int)Math.Min(bytesRead, (nuint)count));
                    foreach (var scanner in scanners)
                        found.AddRange(scanner.Feed(buffer.AsSpan(0, length)));
                }
            }
            finally
            {
                CryptographicOperations.ZeroMemory(buffer);
            }

            if (!success || bytesRead != (nuint)count)
            {
                BreakRun(scanners, found);
                tailIsContiguous = false;
            }
            else
            {
                tailIsContiguous = true;
            }

            offset += (ulong)count;
        }

        return tailIsContiguous;
    }

    private static void BreakRun(IEnumerable<WcdbKeyCandidateScanner> scanners, List<WcdbKeyCandidate> found)
    {
        foreach (var scanner in scanners)
            found.AddRange(scanner.Feed([], finalBlock: true));
    }

    private static bool IsReadable(MEMORY_BASIC_INFORMATION region)
    {
        const uint memCommit = 0x1000;
        const uint pageNoAccess = 0x01;
        const uint pageGuard = 0x100;
        const uint readableMask = 0x02 | 0x04 | 0x08 | 0x20 | 0x40 | 0x80;
        var state = (uint)region.State;
        var protection = (uint)region.Protect;
        return state == memCommit && (protection & (pageNoAccess | pageGuard)) == 0 && (protection & readableMask) != 0;
    }

    private sealed class WcdbKeyCandidateComparer : IEqualityComparer<WcdbKeyCandidate>
    {
        public static readonly WcdbKeyCandidateComparer Instance = new();

        public bool Equals(WcdbKeyCandidate? x, WcdbKeyCandidate? y) =>
            x is not null && y is not null && string.Equals(x.EncKey, y.EncKey, StringComparison.Ordinal);

        public int GetHashCode(WcdbKeyCandidate obj) => StringComparer.Ordinal.GetHashCode(obj.EncKey);
    }
}