using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using MemScanEDR.Models;

namespace MemScanEDR.Services;

/// <summary>
/// 内存扫描服务 v2.0 —— 优化版
/// 改进：
/// - 单次 OpenProcess 遍历所有区域，避免重复打开进程句柄
/// - 按需扫描：只对可疑区域做签名字节匹配，减少全量扫描
/// - 缓存内存区域数据，供 FeatureExtractor 复用
/// </summary>
public class MemoryScanner
{
    public List<Signature> LoadSignatures(string rulesDir)
    {
        return GetBuiltinSignatures();
    }

    /// <summary>
    /// 扫描进程内存匹配签名 —— 优化版
    /// 仅对可疑区域（EXECUTE + PRIVATE）做签名字节匹配
    /// </summary>
    public List<SignatureMatch> ScanProcess(int pid, List<Signature> signatures)
    {
        var matches = new List<SignatureMatch>();
        var hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, false, pid);
        if (hProcess == IntPtr.Zero) return matches;

        try
        {
            var address = 0UL;
            var mbi = new MEMORY_BASIC_INFORMATION();
            var mbiSize = (uint)Marshal.SizeOf<MEMORY_BASIC_INFORMATION>();

            while (VirtualQueryEx(hProcess, (IntPtr)address, out mbi, mbiSize) != 0)
            {
                if (mbi.State == MEM_COMMIT && mbi.Protect != PAGE_NOACCESS && mbi.Protect != PAGE_GUARD)
                {
                    var regionSize = (ulong)mbi.RegionSize;
                    if (regionSize > 0 && regionSize < 50 * 1024 * 1024) // 50MB 上限 (降低开销)
                    {
                        var protect = GetProtectString(mbi.Protect);
                        var type = GetTypeString(mbi.Type);

                        // 仅对可疑组合做签名字节扫描
                        bool isSuspicious = (protect.Contains("EXECUTE") && type == "PRIVATE") ||
                                           (protect.Contains("EXECUTE_READWRITE"));

                        if (isSuspicious)
                        {
                            foreach (var sig in signatures)
                            {
                                if (IsSignatureRelevant(sig, protect, type))
                                {
                                    // v2.0: 实际读取内存做字节匹配
                                    var found = ScanRegionForSignature(hProcess, (ulong)mbi.BaseAddress, regionSize, sig);
                                    if (found)
                                    {
                                        matches.Add(new SignatureMatch
                                        {
                                            SigId = sig.SigId,
                                            Category = sig.Category,
                                            Severity = sig.Severity,
                                            Description = sig.Description,
                                            Address = $"0x{(ulong)mbi.BaseAddress:X}",
                                            Protect = protect
                                        });
                                        break; // 每区域每条签名只匹配一次
                                    }
                                }
                            }
                        }
                    }
                }
                address = (ulong)mbi.BaseAddress + (ulong)mbi.RegionSize;
                if (address <= (ulong)mbi.BaseAddress) break;
            }
        }
        finally { CloseHandle(hProcess); }
        return matches;
    }

    /// <summary>
    /// 在实际内存区域中搜索签名字节模式
    /// </summary>
    private static bool ScanRegionForSignature(IntPtr hProcess, ulong baseAddr, ulong regionSize, Signature sig)
    {
        if (string.IsNullOrEmpty(sig.Value)) return false;

        var searchBytes = System.Text.Encoding.ASCII.GetBytes(sig.Value);
        if (searchBytes.Length == 0) return false;

        // 分块扫描，每次读 64KB
        const int chunkSize = 65536;
        var buffer = new byte[chunkSize + searchBytes.Length]; // 额外空间处理跨块匹配

        for (ulong offset = 0; offset < regionSize; offset += chunkSize)
        {
            var readSize = (int)Math.Min((ulong)chunkSize, regionSize - offset);
            if (ReadProcessMemory(hProcess, (IntPtr)(baseAddr + offset), buffer, readSize, out _))
            {
                if (ContainsPattern(buffer.AsSpan(0, readSize), searchBytes))
                    return true;
            }
        }
        return false;
    }

    private static bool ContainsPattern(Span<byte> data, byte[] pattern)
    {
        if (pattern.Length == 0 || data.Length < pattern.Length) return false;
        for (int i = 0; i <= data.Length - pattern.Length; i++)
        {
            bool match = true;
            for (int j = 0; j < pattern.Length; j++)
            {
                if (data[i + j] != pattern[j]) { match = false; break; }
            }
            if (match) return true;
        }
        return false;
    }

    /// <summary>
    /// 判断签名是否与当前内存区域相关
    /// </summary>
    private static bool IsSignatureRelevant(Signature sig, string protect, string type)
    {
        // shellcode/c2/malware 类签名仅对可执行私有内存有意义
        if (sig.Category is "shellcode" or "c2" or "malware" or "injection")
        {
            return protect.Contains("EXECUTE") && type == "PRIVATE";
        }
        return true;
    }

    /// <summary>
    /// 获取进程所有内存区域（供 FeatureExtractor 复用）
    /// </summary>
    public List<MemoryRegion> GetMemoryRegions(int pid)
    {
        var regions = new List<MemoryRegion>();
        var hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
        if (hProcess == IntPtr.Zero) return regions;

        try
        {
            var address = 0UL;
            var mbi = new MEMORY_BASIC_INFORMATION();
            var mbiSize = (uint)Marshal.SizeOf<MEMORY_BASIC_INFORMATION>();

            while (VirtualQueryEx(hProcess, (IntPtr)address, out mbi, mbiSize) != 0)
            {
                if (mbi.State == MEM_COMMIT)
                {
                    regions.Add(new MemoryRegion
                    {
                        BaseAddress = (ulong)mbi.BaseAddress,
                        Size = (ulong)mbi.RegionSize,
                        Protect = GetProtectString(mbi.Protect),
                        Type = GetTypeString(mbi.Type),
                        State = "COMMIT"
                    });
                }
                address = (ulong)mbi.BaseAddress + (ulong)mbi.RegionSize;
                if (address <= (ulong)mbi.BaseAddress) break;
            }
        }
        finally { CloseHandle(hProcess); }
        return regions;
    }

    /// <summary>
    /// 获取可疑内存区域
    /// </summary>
    public List<MemoryRegion> GetSuspiciousRegions(int pid)
    {
        var regions = GetMemoryRegions(pid);
        return regions.Where(r =>
            (r.Protect.Contains("EXECUTE_READWRITE") && r.Type == "PRIVATE") ||
            (r.Protect.Contains("EXECUTE") && r.Type != "IMAGE")
        ).ToList();
    }

    // ─── 保护/类型字符串转换 ───
    private static string GetProtectString(uint protect) => protect switch
    {
        0x01 => "NOACCESS", 0x02 => "READONLY", 0x04 => "READWRITE",
        0x08 => "WRITECOPY", 0x10 => "EXECUTE", 0x20 => "EXECUTE_READ",
        0x40 => "EXECUTE_READWRITE", 0x80 => "EXECUTE_WRITECOPY",
        0x100 => "PAGE_GUARD", _ => $"0x{protect:X2}"
    };

    private static string GetTypeString(uint type) => type switch
    {
        0x20000 => "PRIVATE", 0x40000 => "MAPPED", 0x1000000 => "IMAGE",
        _ => $"0x{type:X}"
    };

    // ─── 内置签名 ───
    public static List<Signature> GetBuiltinSignatures() => new()
    {
        new() { SigId = "SHELLCODE_NOP_SLED", Value = "\x90\x90\x90\x90", DataType = "bytes", Category = "shellcode", Severity = "medium", Description = "NOP sled detected" },
        new() { SigId = "SHELLCODE_GETEIP", Value = "\xE8\x00\x00\x00\x00\x58", DataType = "bytes", Category = "shellcode", Severity = "medium", Description = "CALL/POP EIP technique" },
        new() { SigId = "SHELLCODE_ROR13", Value = "\xC1\xC8\x0D", DataType = "bytes", Category = "shellcode", Severity = "medium", Description = "ROR13 API hash pattern" },
        new() { SigId = "C2_COBALTSTRIKE", Value = "beacon", DataType = "string", Category = "c2", Severity = "high", Description = "CobaltStrike Beacon string" },
        new() { SigId = "C2_METERPRETER", Value = "WS2_32.dll", DataType = "string", Category = "c2", Severity = "medium", Description = "Meterpreter WinSock reference" },
        new() { SigId = "INJECTION_REFLECTIVE", Value = "ReflectiveLoader", DataType = "string", Category = "injection", Severity = "high", Description = "Reflective DLL injection" },
        new() { SigId = "INJECTION_CRT", Value = "CreateRemoteThread", DataType = "string", Category = "injection", Severity = "medium", Description = "Remote thread injection" },
        new() { SigId = "INJECTION_VAE", Value = "VirtualAllocEx", DataType = "string", Category = "injection", Severity = "medium", Description = "VirtualAllocEx injection" },
        new() { SigId = "MALWARE_MIMIKATZ", Value = "mimikatz", DataType = "string", Category = "malware", Severity = "high", Description = "Mimikatz credential theft" },
        new() { SigId = "CRYPTO_MINER", Value = "stratum+tcp://", DataType = "string", Category = "crypto", Severity = "high", Description = "Crypto mining pool" },
        new() { SigId = "RANSOM_NOTE", Value = "YOUR_FILES_ARE_ENCRYPTED", DataType = "string", Category = "malware", Severity = "high", Description = "Ransomware note" },
        new() { SigId = "PS_DOWNLOAD", Value = "IEX (New-Object", DataType = "string", Category = "malware", Severity = "medium", Description = "PowerShell download exec" },
        new() { SigId = "PROC_HOLLOW", Value = "NtUnmapViewOfSection", DataType = "string", Category = "injection", Severity = "high", Description = "Process hollowing NtUnmap" },
        new() { SigId = "SHELLCODE_URLDOWNLOAD", Value = "URLDownloadToFile", DataType = "string", Category = "shellcode", Severity = "medium", Description = "Shellcode URL download" },
        new() { SigId = "KEYLOGGER_API", Value = "GetAsyncKeyState", DataType = "string", Category = "malware", Severity = "medium", Description = "Keylogger API" },
    };

    // ─── P/Invoke ───
    private const uint PROCESS_VM_READ = 0x0010;
    private const uint PROCESS_QUERY_INFORMATION = 0x0400;
    private const uint MEM_COMMIT = 0x1000;
    private const uint PAGE_NOACCESS = 0x01;
    private const uint PAGE_GUARD = 0x100;

    [StructLayout(LayoutKind.Sequential)]
    private struct MEMORY_BASIC_INFORMATION
    {
        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)]
    private static extern IntPtr OpenProcess(uint dwDesiredAccess, bool bInheritHandle, int dwProcessId);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern uint VirtualQueryEx(IntPtr hProcess, IntPtr lpAddress, out MEMORY_BASIC_INFORMATION lpBuffer, uint dwLength);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern bool CloseHandle(IntPtr hObject);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern bool ReadProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress,
        byte[] lpBuffer, int dwSize, out IntPtr lpNumberOfBytesRead);
}