#include "Common.h"
#include "VmDetect.h"

#include <intrin.h>
#include <iphlpapi.h>
#include <string>

#pragma comment(lib, "iphlpapi.lib")

static std::string ToLowerAscii(const std::string& in)
{
    std::string out = in;

    for (auto& c : out)
    {
        if (c >= 'A' && c <= 'Z')
            c = (char)(c + 32);
    }

    return out;
}

static std::string NarrowAscii(const wchar_t* w)
{
    std::string out;

    if (!w)
        return out;

    for (; *w; ++w)
    {
        if (*w < 128)
            out.push_back((char)*w);
    }

    return out;
}

static bool ContainsAny(const std::string& haystackLower, std::initializer_list<const char*> needles)
{
    for (const char* n : needles)
    {
        if (haystackLower.find(n) != std::string::npos)
            return true;
    }

    return false;
}

// ---------------------------------------------------------------- CPUID

struct CpuidResult
{
    bool hypervisorPresent = false;
    std::string vendor;   // raw 12-byte hypervisor vendor
    std::string brand;
};

static CpuidResult QueryCpuid()
{
    CpuidResult r;

    int regs[4]{};
    __cpuid(regs, 0);
    int maxLeaf = regs[0];

    __cpuid(regs, 1);
    r.hypervisorPresent = (regs[2] & (1 << 31)) != 0;

    if (maxLeaf >= 0x40000000)
    {
        int hv[4]{};
        __cpuid(hv, 0x40000000);

        char vendor[13]{};
        memcpy(vendor + 0, &hv[1], 4);
        memcpy(vendor + 4, &hv[2], 4);
        memcpy(vendor + 8, &hv[3], 4);
        r.vendor = vendor;

        if (maxLeaf >= 0x40000001)
        {
            int b[4]{};
            __cpuid(b, 0x40000001);

            char brand[13]{};
            memcpy(brand + 0, &b[0], 4);
            memcpy(brand + 4, &b[1], 4);
            memcpy(brand + 8, &b[2], 4);
            r.brand = brand;
        }
    }

    return r;
}

// ---------------------------------------------------------------- SMBIOS

static std::string ReadSmbiosStrings()
{
    std::string collected;

    DWORD size = GetSystemFirmwareTable('RSMB', 0, nullptr, 0);

    if (size < 8 || size > 1024 * 1024)
        return collected;

    std::vector<BYTE> buf(size);
    DWORD got = GetSystemFirmwareTable('RSMB', 0, buf.data(), size);

    if (got < 8)
        return collected;

    const BYTE* data = buf.data();
    DWORD length = got;

    // Raw SMBIOS table: 8-byte header, then structure table.
    DWORD tableLen = *(const DWORD*)(data + 4);

    if (tableLen == 0 || tableLen > length - 8)
        tableLen = length - 8;

    const BYTE* p = data + 8;
    const BYTE* end = data + 8 + tableLen;

    while (p + 4 <= end)
    {
        BYTE structLen = p[1];

        if (structLen < 4 || p + structLen > end)
            break;

        const BYTE* strings = p + structLen;

        // Strings: NUL-terminated ASCII, double NUL ends the list.
        while (strings + 1 < end && !(strings[0] == 0 && strings[1] == 0))
        {
            const char* s = (const char*)strings;
            size_t len = strlen(s);

            collected += s;
            collected += ' ';

            strings += len + 1;
        }

        p = strings + 2;
    }

    return collected;
}

// ---------------------------------------------------------------- registry

static bool RegKeyExists(HKEY root, const wchar_t* subKey)
{
    HKEY h = nullptr;

    if (RegOpenKeyExW(root, subKey, 0, KEY_READ, &h) != ERROR_SUCCESS)
        return false;

    RegCloseKey(h);
    return true;
}

static std::string ReadDiskEnumValue()
{
    std::string out;

    HKEY h = nullptr;

    if (RegOpenKeyExW(
        HKEY_LOCAL_MACHINE,
        L"SYSTEM\\CurrentControlSet\\Services\\Disk\\Enum",
        0,
        KEY_READ,
        &h) != ERROR_SUCCESS)
    {
        return out;
    }

    wchar_t buf[1024]{};
    DWORD size = sizeof(buf);
    DWORD type = 0;

    if (RegQueryValueExW(h, L"0", nullptr, &type, (LPBYTE)buf, &size) == ERROR_SUCCESS &&
        (type == REG_SZ || type == REG_EXPAND_SZ))
    {
        out = NarrowAscii(buf);
    }

    RegCloseKey(h);
    return out;
}

// ---------------------------------------------------------------- network

static bool HasVirtualMac()
{
    // Vendor OUIs that only appear on virtual adapters.
    static const BYTE oui[][3] =
    {
        { 0x00, 0x0C, 0x29 }, // VMware
        { 0x00, 0x50, 0x56 }, // VMware
        { 0x00, 0x05, 0x69 }, // VMware
        { 0x00, 0x1C, 0x14 }, // VMware
        { 0x08, 0x00, 0x27 }, // VirtualBox
        { 0x52, 0x54, 0x00 }, // QEMU / KVM
        { 0x00, 0x16, 0x3E }, // Xen
        { 0x00, 0x15, 0x5D }, // Hyper-V
        { 0x00, 0x03, 0xFF }, // Microsoft virtual
        { 0x00, 0x1C, 0x42 }, // Parallels
    };

    ULONG size = 16 * 1024;
    std::vector<BYTE> buf(size);

    ULONG ret = GetAdaptersInfo((PIP_ADAPTER_INFO)buf.data(), &size);

    if (ret == ERROR_BUFFER_OVERFLOW)
    {
        buf.resize(size);
        ret = GetAdaptersInfo((PIP_ADAPTER_INFO)buf.data(), &size);
    }

    if (ret != ERROR_SUCCESS)
        return false;

    for (PIP_ADAPTER_INFO a = (PIP_ADAPTER_INFO)buf.data(); a; a = a->Next)
    {
        for (const auto& prefix : oui)
        {
            if (a->AddressLength >= 3 &&
                a->Address[0] == prefix[0] &&
                a->Address[1] == prefix[1] &&
                a->Address[2] == prefix[2])
            {
                return true;
            }
        }
    }

    return false;
}

// ---------------------------------------------------------------- artifacts

static bool HasVmServices()
{
    static const wchar_t* services[] =
    {
        L"SYSTEM\\CurrentControlSet\\Services\\VBoxService",
        L"SYSTEM\\CurrentControlSet\\Services\\VBoxGuest",
        L"SYSTEM\\CurrentControlSet\\Services\\VBoxMouse",
        L"SYSTEM\\CurrentControlSet\\Services\\VBoxSF",
        L"SYSTEM\\CurrentControlSet\\Services\\VBoxWddm",
        L"SYSTEM\\CurrentControlSet\\Services\\vmci",
        L"SYSTEM\\CurrentControlSet\\Services\\vmhgfs",
        L"SYSTEM\\CurrentControlSet\\Services\\VMTools",
        L"SYSTEM\\CurrentControlSet\\Services\\vmrawdsk",
        L"SYSTEM\\CurrentControlSet\\Services\\vmusbmouse",
        L"SYSTEM\\CurrentControlSet\\Services\\vm3dservice",
        L"SYSTEM\\CurrentControlSet\\Services\\vmmouse",
        L"SYSTEM\\CurrentControlSet\\Services\\xenvdb",
        L"SYSTEM\\CurrentControlSet\\Services\\xenbus",
    };

    for (const wchar_t* s : services)
    {
        if (RegKeyExists(HKEY_LOCAL_MACHINE, s))
            return true;
    }

    return false;
}

static bool HasVmDriversOnDisk()
{
    static const wchar_t* files[] =
    {
        L"C:\\Windows\\System32\\drivers\\VBoxGuest.sys",
        L"C:\\Windows\\System32\\drivers\\VBoxMouse.sys",
        L"C:\\Windows\\System32\\drivers\\VBoxSF.sys",
        L"C:\\Windows\\System32\\drivers\\vmmouse.sys",
        L"C:\\Windows\\System32\\drivers\\vmhgfs.sys",
        L"C:\\Windows\\System32\\drivers\\vmrawdsk.sys",
        L"C:\\Windows\\System32\\drivers\\vmusbmouse.sys",
        L"C:\\Windows\\System32\\VBoxService.exe",
        L"C:\\Windows\\System32\\VBoxTray.exe",
    };

    for (const wchar_t* f : files)
    {
        DWORD attr = GetFileAttributesW(f);

        if (attr != INVALID_FILE_ATTRIBUTES && !(attr & FILE_ATTRIBUTE_DIRECTORY))
            return true;
    }

    return false;
}

static bool HasVmProcesses()
{
    static const wchar_t* names[] =
    {
        L"VBoxService.exe",
        L"VBoxTray.exe",
        L"vmtoolsd.exe",
        L"vmwaretray.exe",
        L"vmwareuser.exe",
        L"VGAuthService.exe",
        L"qemu-ga.exe",
        L"xenservice.exe",
    };

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snap == INVALID_HANDLE_VALUE)
        return false;

    PROCESSENTRY32W pe{};
    pe.dwSize = sizeof(pe);

    bool found = false;

    if (Process32First(snap, &pe))
    {
        do
        {
            for (const wchar_t* n : names)
            {
                if (_wcsicmp(pe.szExeFile, n) == 0)
                {
                    found = true;
                    break;
                }
            }

            if (found)
                break;
        } while (Process32Next(snap, &pe));
    }

    CloseHandle(snap);
    return found;
}

// ---------------------------------------------------------------- public

bool IsVirtualMachine(std::wstring* reason)
{
    auto setReason = [reason](const wchar_t* text)
    {
        if (reason)
            *reason = text;
    };

    if (reason)
        reason->clear();

    CpuidResult cpu = QueryCpuid();

    // A hypervisor bit alone is not proof: VBS / Credential Guard / Hyper-V
    // on bare metal sets it too. Strong vendors are decisive.
    if (cpu.hypervisorPresent)
    {
        const std::string vendor = cpu.vendor;

        if (vendor == "VMwareVMware") { setReason(L"cpuid: VMware"); return true; }
        if (vendor == "VBoxVBoxVBox") { setReason(L"cpuid: VirtualBox"); return true; }
        if (vendor == "KVMKVMKVM")    { setReason(L"cpuid: KVM"); return true; }
        if (vendor == "XenVMMXenVMM") { setReason(L"cpuid: Xen"); return true; }
        if (vendor.rfind("TCGTCGTCG", 0) == 0) { setReason(L"cpuid: QEMU TCG"); return true; }
        if (vendor == "bhyve bhyve")  { setReason(L"cpuid: bhyve"); return true; }
        if (vendor == "VServerVServer") { setReason(L"cpuid: Virtual Server"); return true; }
    }

    const std::string smbios = ToLowerAscii(ReadSmbiosStrings());

    if (ContainsAny(smbios, { "virtualbox", "vmware", "qemu", "kvm", "bochs", "parallels", "xen", "virtual machine", "hyper-v" }))
    {
        setReason(L"smbios: virtual vendor string");
        return true;
    }

    if (cpu.hypervisorPresent && cpu.vendor == "Microsoft Hv")
    {
        // Hyper-V guest: confirmed by SMBIOS containing "virtual machine"
        // (handled above) or by Hyper-V integration services.
        if (RegKeyExists(HKEY_LOCAL_MACHINE, L"SYSTEM\\CurrentControlSet\\Services\\vmicheartbeat"))
        {
            setReason(L"hyperv: integration services");
            return true;
        }

        if (HasVirtualMac())
        {
            setReason(L"hyperv: virtual MAC");
            return true;
        }
    }

    std::string diskEnum = ToLowerAscii(ReadDiskEnumValue());

    if (ContainsAny(diskEnum, { "vbox", "vmware", "qemu", "xen", "virtio" }))
    {
        setReason(L"disk enum: virtual controller");
        return true;
    }

    if (HasVmServices())
    {
        setReason(L"registry: guest service installed");
        return true;
    }

    if (HasVmDriversOnDisk())
    {
        setReason(L"driver: guest driver present");
        return true;
    }

    if (HasVmProcesses())
    {
        setReason(L"process: guest helper running");
        return true;
    }

    if (HasVirtualMac())
    {
        setReason(L"network: virtual MAC OUI");
        return true;
    }

    return false;
}

#ifndef ARDRV_ALLOW_VM_DEFAULT
#define ARDRV_ALLOW_VM_DEFAULT 0
#endif

bool VmOverrideRequested(int argc, wchar_t** argv)
{
#if ARDRV_ALLOW_VM_DEFAULT
    // Built for lab/testing images that are themselves virtual machines.
    (void)argc;
    (void)argv;
    return true;
#endif

    for (int i = 1; i < argc; ++i)
    {
        if (argv[i] && (_wcsicmp(argv[i], L"--allow-vm") == 0 ||
                        _wcsicmp(argv[i], L"-allow-vm") == 0))
        {
            return true;
        }
    }

    wchar_t env[64]{};
    DWORD n = GetEnvironmentVariableW(L"ARDRV_ALLOW_VM", env, _countof(env));

    if (n > 0 && n < _countof(env) && (env[0] == L'1' || env[0] == L'y' || env[0] == L'Y'))
        return true;

    HKEY h = nullptr;

    if (RegOpenKeyExW(HKEY_LOCAL_MACHINE, L"SOFTWARE\\ArDrvLoader", 0, KEY_READ, &h) == ERROR_SUCCESS)
    {
        DWORD value = 0;
        DWORD size = sizeof(value);
        DWORD type = 0;

        bool enabled =
            RegQueryValueExW(h, L"AllowVm", nullptr, &type, (LPBYTE)&value, &size) == ERROR_SUCCESS &&
            type == REG_DWORD &&
            value != 0;

        RegCloseKey(h);

        if (enabled)
            return true;
    }

    return false;
}