#include "../main.h"
#include "ModuleBase.hpp"

#include <string>
#include <string_view>
#include <windows.h>
#include "../wine_winternl.h"
#include <ntstatus.h>
#include <detours.h>



class PolicyDisable : public ModuleBase {
    inline static decltype(&NtQueryValueKey) TrueNtQueryValueKey;

    static
    NTSTATUS WINAPI DetourNtQueryValueKey(_In_ HANDLE KeyHandle,
        _In_ PUNICODE_STRING ValueName, _In_ KEY_VALUE_INFORMATION_CLASS KeyValueInformationClass,
        _Out_opt_ PVOID KeyValueInformation, _In_ ULONG Length, _Out_ PULONG ResultLength) {
        // lookup the size for the key name so we can allocate space for it
        DWORD iKeyNameSize;
        if (NtQueryKey(KeyHandle, KeyNameInformation, NULL, 0, &iKeyNameSize) != STATUS_BUFFER_TOO_SMALL) {
            // should never happen
            return STATUS_INVALID_PARAMETER;
        }

        // allocate space for name and lookup
        NTSTATUS iStatus = -1;
        PKEY_NAME_INFORMATION pNameInfo = (PKEY_NAME_INFORMATION)malloc(iKeyNameSize);
        if (pNameInfo != NULL && NtQueryKey(KeyHandle, KeyNameInformation, pNameInfo,
            iKeyNameSize, &iKeyNameSize) == STATUS_SUCCESS) {
            // get wstring_view for quick comparisations
            std::wstring_view sKeyName{pNameInfo->Name, pNameInfo->NameLength/2};

            // return failure if key name contains "polic", "Polic" or "POLIC"
            if (sKeyName.find(L"polic") != sKeyName.npos ||
                sKeyName.find(L"Polic") != sKeyName.npos ||
                sKeyName.find(L"POLIC") != sKeyName.npos) {
                *log_out << "Denied policy access: " << std::string_view{reinterpret_cast<char*>(pNameInfo->Name), pNameInfo->NameLength} << "\r\n";
                free(pNameInfo);
                *ResultLength = 0;
                return STATUS_OBJECT_NAME_NOT_FOUND;
            }
            *log_out << "Allowed policy access: " << std::string_view{reinterpret_cast<char*>(pNameInfo->Name), pNameInfo->NameLength} << "\r\n";
        }

        // clean up
        free(pNameInfo);

        // execute real function and return its value
        return TrueNtQueryValueKey(KeyHandle, ValueName, KeyValueInformationClass, KeyValueInformation, Length, ResultLength);
    }

public:
    PolicyDisable() {
        TrueNtQueryValueKey = reinterpret_cast<decltype(&NtQueryValueKey)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtQueryValueKey"));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtQueryValueKey), reinterpret_cast<void*>(DetourNtQueryValueKey));
    }
    ~PolicyDisable() {
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtQueryValueKey), reinterpret_cast<void*>(DetourNtQueryValueKey));
    }
};