using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;

namespace MemScanEDR.Services;

/// <summary>
/// PE 头分析器 —— 验证内存中的 PE 镜像与磁盘文件是否一致
/// 用于检测 Process Hollowing（进程掏空）攻击
/// </summary>
public class PEHeaderAnalyzer
{
    private const uint PROCESS_VM_READ = 0x0010;
    private const uint PROCESS_QUERY_INFORMATION = 0x0400;

    // PE 签名 "MZ"
    private const ushort MZ_SIGNATURE = 0x5A4D;
    // PE\0\0 签名
    private const uint PE_SIGNATURE = 0x00004550;

    /// <summary>
    /// 检测进程内存中是否存在 PE 头 (MZ 签名)
    /// </summary>
    public bool HasInMemoryPE(IntPtr hProcess, ulong baseAddress)
    {
        var buffer = new byte[2];
        if (ReadProcessMemory(hProcess, (IntPtr)baseAddress, buffer, 2, out _))
        {
            var mz = BitConverter.ToUInt16(buffer, 0);
            return mz == MZ_SIGNATURE;
        }
        return false;
    }

    /// <summary>
    /// 检测内存 PE 与磁盘文件是否一致 (Process Hollowing 检测)
    /// 比较内存中 PE 头的关键字段与磁盘文件
    /// </summary>
    public PeIntegrityResult CheckPeIntegrity(string diskPath, IntPtr hProcess, ulong imageBase)
    {
        var result = new PeIntegrityResult();

        if (string.IsNullOrEmpty(diskPath) || !File.Exists(diskPath))
        {
            result.IsValid = false;
            result.Reason = "磁盘文件不存在";
            return result;
        }

        try
        {
            // 读取磁盘 PE 头
            var diskPe = File.ReadAllBytes(diskPath);

            // 读取内存 PE 头 (前 4096 字节)
            var memBuffer = new byte[Math.Min(4096, diskPe.Length)];
            if (!ReadProcessMemory(hProcess, (IntPtr)imageBase, memBuffer, memBuffer.Length, out _))
            {
                result.IsValid = false;
                result.Reason = "无法读取进程内存";
                return result;
            }

            // 比较 DOS 头
            for (int i = 0; i < Math.Min(64, Math.Min(diskPe.Length, memBuffer.Length)); i++)
            {
                if (diskPe[i] != memBuffer[i])
                {
                    result.MismatchCount++;
                }
            }

            // 比较 PE 签名位置
            if (diskPe.Length > 0x3C + 4 && memBuffer.Length > 0x3C + 4)
            {
                var diskPeOffset = BitConverter.ToInt32(diskPe, 0x3C);
                var memPeOffset = BitConverter.ToInt32(memBuffer, 0x3C);

                if (diskPeOffset != memPeOffset)
                {
                    result.MismatchCount++;
                }
                else if (diskPeOffset > 0 && diskPeOffset + 4 <= Math.Min(diskPe.Length, memBuffer.Length))
                {
                    var diskSig = BitConverter.ToUInt32(diskPe, diskPeOffset);
                    var memSig = BitConverter.ToUInt32(memBuffer, diskPeOffset);
                    if (diskSig != memSig && memSig != PE_SIGNATURE)
                    {
                        result.MismatchCount++;
                    }
                }
            }

            result.IsValid = result.MismatchCount <= 5; // 容忍少量差异
            result.Reason = result.IsValid
                ? "内存PE与磁盘一致"
                : $"内存PE与磁盘不一致 ({result.MismatchCount} 处差异) —— 可能存在 Process Hollowing";
        }
        catch (Exception ex)
        {
            result.IsValid = false;
            result.Reason = $"PE 完整性检查异常: {ex.Message}";
        }

        return result;
    }

    /// <summary>
    /// 从内存中提取可疑字符串 (C2域名、挖矿池、凭证路径等)
    /// </summary>
    public List<string> ExtractSuspiciousStrings(IntPtr hProcess, ulong baseAddress, ulong size)
    {
        var result = new List<string>();
        var buffer = new byte[Math.Min((int)size, 65536)];

        if (!ReadProcessMemory(hProcess, (IntPtr)baseAddress, buffer, buffer.Length, out _))
            return result;

        // 提取ASCII字符串（最小4字符）
        var current = new StringBuilder();
        for (int i = 0; i < buffer.Length; i++)
        {
            if (buffer[i] >= 0x20 && buffer[i] <= 0x7E)
            {
                current.Append((char)buffer[i]);
            }
            else
            {
                if (current.Length >= 4)
                {
                    var s = current.ToString();
                    if (IsSuspiciousString(s))
                        result.Add(s);
                }
                current.Clear();
            }
        }

        return result;
    }

    /// <summary>
    /// 检测内存中的可疑 API 字符串
    /// </summary>
    public List<string> ExtractSuspiciousApis(IntPtr hProcess, ulong baseAddress, ulong size)
    {
        var result = new List<string>();
        var buffer = new byte[Math.Min((int)size, 65536)];

        if (!ReadProcessMemory(hProcess, (IntPtr)baseAddress, buffer, buffer.Length, out _))
            return result;

        var text = Encoding.ASCII.GetString(buffer);
        foreach (var api in DangerousApis)
        {
            if (text.Contains(api, StringComparison.OrdinalIgnoreCase))
                result.Add(api);
        }

        return result;
    }

    private static bool IsSuspiciousString(string s) => s.Length switch
    {
        >= 4 when SuspiciousPatterns.Any(p => s.Contains(p, StringComparison.OrdinalIgnoreCase)) => true,
        _ => false
    };

    private static readonly string[] SuspiciousPatterns = new[]
    {
        "stratum+tcp://", "xmr.", "cryptonight", "coinmine", "minergate",
        ".onion", "c2.", "beacon", "meterpreter", "shellcode",
        "mimikatz", "lsass", "SAM\\", "SYSTEM\\",
        "YOUR_FILES_ARE_ENCRYPTED", "decrypt", "ransom", "bitcoin",
        "powershell -enc", "IEX (New-Object", "Invoke-Expression",
        "keylogger", "keylog", "GetAsyncKeyState",
        "cmd.exe /c", "schtasks /create", "reg add",
        "\\AppData\\Roaming\\", "\\Temp\\", "Startup\\"
    };

    private static readonly string[] DangerousApis = new[]
    {
        "VirtualAlloc", "VirtualAllocEx", "VirtualProtect", "VirtualProtectEx",
        "WriteProcessMemory", "CreateRemoteThread", "NtCreateThreadEx",
        "RtlCreateUserThread", "QueueUserAPC", "SetThreadContext",
        "OpenProcess", "NtOpenProcess", "AdjustTokenPrivileges",
        "CryptEncrypt", "CryptDecrypt", "CryptAcquireContext",
        "WSASocket", "connect", "send", "recv", "InternetConnect",
        "URLDownloadToFile", "WinHttpOpen", "WinHttpConnect",
        "CreateToolhelp32Snapshot", "Process32First", "Module32First",
        "GetProcAddress", "LoadLibrary", "LdrLoadDll",
        "NtQuerySystemInformation", "NtQueryInformationProcess",
        "NtReadVirtualMemory", "NtWriteVirtualMemory",
        "MiniDumpWriteDump", "SamOpenUser", "SamQueryInformationUser"
    };

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

/// <summary>
/// PE 完整性检查结果
/// </summary>
public class PeIntegrityResult
{
    public bool IsValid { get; set; }
    public string Reason { get; set; } = "";
    public int MismatchCount { get; set; }
}