#include "sandbox.h"

std::string getDllNameFromApiSetMap(const std::string& apiSet) {
    const std::wstring wApiSet(apiSet.begin(), apiSet.end());

    // 获取系统版本信息
    using RtlGetVersionFunc = LONG(__stdcall*)(PRTL_OSVERSIONINFOW);
    const auto pRtlGetVersion = reinterpret_cast<RtlGetVersionFunc>(
        GetProcAddress(LoadLibraryA("ntdll.dll"), "RtlGetVersion"));

    RTL_OSVERSIONINFOEXW verInfo{};
    verInfo.dwOSVersionInfoSize = sizeof(verInfo);
    pRtlGetVersion(reinterpret_cast<PRTL_OSVERSIONINFOW>(&verInfo));

    const ULONG verShort = (verInfo.dwMajorVersion << 8) |
        (verInfo.dwMinorVersion << 4) |
        verInfo.wServicePackMajor;

    if (verShort >= static_cast<ULONG>(WinVer::kWin10)) {
        const auto apiSetMap = reinterpret_cast<API_SET_NAMESPACE_ARRAY_10*>(
            reinterpret_cast<X64PEB*>(__readgsqword(0x60))->ApiSetMap);
        const auto apiSetMapAsNumber = reinterpret_cast<ULONG_PTR>(apiSetMap);
        auto nsEntry = reinterpret_cast<PAPI_SET_NAMESPACE_ENTRY_10>(
            apiSetMap->Start + apiSetMapAsNumber);

        // 遍历API集合查找匹配项
        for (ULONG i = 0; i < apiSetMap->Count; i++) {
            UNICODE_STRING nameString{}, valueString{};
            nameString.MaximumLength = static_cast<USHORT>(nsEntry->NameLength);
            nameString.Length = static_cast<USHORT>(nsEntry->NameLength);
            nameString.Buffer = reinterpret_cast<PWCHAR>(apiSetMapAsNumber +
                nsEntry->NameOffset);

            const std::wstring name(nameString.Buffer,
                nameString.Length / sizeof(WCHAR));
            const std::wstring fullName = name + L".dll";

            if (_wcsicmp(wApiSet.c_str(), fullName.c_str()) == 0) {
                if (nsEntry->ValueCount == 0) {
                    return "";
                }

                const auto valueEntry =
                    reinterpret_cast<PAPI_SET_VALUE_ENTRY_10>(
                        apiSetMapAsNumber + nsEntry->ValueOffset);
                valueString.Buffer = reinterpret_cast<PWCHAR>(
                    apiSetMapAsNumber + valueEntry->ValueOffset);
                valueString.MaximumLength =
                    static_cast<USHORT>(valueEntry->ValueLength);
                valueString.Length =
                    static_cast<USHORT>(valueEntry->ValueLength);

                const std::wstring value(valueString.Buffer,
                    valueString.Length / sizeof(WCHAR));
                return { value.begin(), value.end() };
            }
            ++nsEntry;
        }
    }
    else {
        // 不支持Windows 10以下版本
        throw std::runtime_error("Unsupported Windows version");
    }
    return "";
}