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

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



class AdminImpersonate : public ModuleBase {
    inline static decltype(&IsUserAnAdmin) TrueIsUserAnAdmin;
    inline static decltype(&CheckTokenMembership) TrueCheckTokenMembership;
    inline static decltype(&NtOpenFile) TrueNtOpenFile;
    inline static decltype(&NtCreateFile) TrueNtCreateFile;
    inline static decltype(&NtAccessCheck) TrueNtAccessCheck;
    inline static decltype(&NtAccessCheckAndAuditAlarm) TrueNtAccessCheckAndAuditAlarm;
    inline static decltype(&NtPrivilegeCheck) TrueNtPrivilegeCheck;

    static
    BOOL __stdcall DetourIsUserAnAdmin() {
        return TRUE;
    }

    static
    BOOL APIENTRY DetourCheckTokenMembership(_In_opt_ HANDLE TokenHandle,
                                             _In_ PSID SidToCheck, _Out_ PBOOL IsMember) {
        // fetch and allocate the local admin structure
        static SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
        static PSID LocalAdministratorsGroup = NULL;
        AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID,
            DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &LocalAdministratorsGroup);

        // get the real value of the function - return if failure
        BOOL bRealResult = TrueCheckTokenMembership(TokenHandle, SidToCheck, IsMember);
        if (bRealResult == 0) return bRealResult;

        // check if the local admin group is being requested
        if (EqualSid(SidToCheck, LocalAdministratorsGroup)) {
            // unconditionally say this user is running as an admin
            *IsMember = TRUE;
        }

        return bRealResult;
    }

    static
    NTSTATUS NTAPI DetourNtOpenFile(OUT PHANDLE FileHandle,
                                    IN ACCESS_MASK DesiredAccess, IN POBJECT_ATTRIBUTES ObjectAttributes, OUT PIO_STATUS_BLOCK IoStatusBlock,
                                    IN ULONG ShareAccess, IN ULONG OpenOptions) {
        DWORD iStatus = TrueNtOpenFile(FileHandle, DesiredAccess, ObjectAttributes,
            IoStatusBlock, ShareAccess, OpenOptions);

        if (iStatus == STATUS_SHARING_VIOLATION || iStatus == STATUS_ACCESS_DENIED) {
            return STATUS_NO_SUCH_FILE;
        }

        return iStatus;
    }

    static
    NTSTATUS NTAPI DetourNtCreateFile(OUT PHANDLE FileHandle, IN ACCESS_MASK DesiredAccess,
                                      IN POBJECT_ATTRIBUTES ObjectAttributes, OUT PIO_STATUS_BLOCK IoStatusBlock,
                                      IN PLARGE_INTEGER AllocationSize OPTIONAL, IN ULONG FileAttributes, IN ULONG ShareAccess,
                                      IN ULONG CreateDisposition, IN ULONG CreateOptions, IN PVOID EaBuffer OPTIONAL, IN ULONG EaLength) {
        NTSTATUS iStatus = TrueNtCreateFile(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, AllocationSize,
            FileAttributes, ShareAccess, CreateDisposition, CreateOptions, EaBuffer, EaLength);

        if (iStatus == STATUS_SHARING_VIOLATION || iStatus == STATUS_ACCESS_DENIED) {
            return STATUS_DISK_FULL;
        }

        return iStatus;
    }

    static
    NTSTATUS NTAPI DetourNtAccessCheck(IN PSECURITY_DESCRIPTOR SecurityDescriptor, IN HANDLE Handle, IN ACCESS_MASK AccessMask,
                                       IN PGENERIC_MAPPING pGenericMapping, IN PPRIVILEGE_SET PrivilegeSet, PULONG Unk1, PULONG Unk2, OUT NTSTATUS* Result) {
        TrueNtAccessCheck(SecurityDescriptor, Handle, AccessMask, pGenericMapping, PrivilegeSet, Unk1, Unk2, Result);
        *Result = STATUS_SUCCESS;
        return STATUS_SUCCESS;
    }
    static
    NTSTATUS NTAPI DetourNtAccessCheckAndAuditAlarm(PUNICODE_STRING A, HANDLE B, PUNICODE_STRING C, PUNICODE_STRING D, PSECURITY_DESCRIPTOR E, ACCESS_MASK F,
                                                    PGENERIC_MAPPING G, BOOLEAN H, PACCESS_MASK I , PBOOLEAN J, PBOOLEAN K) {
        auto res = MessageBoxA(nullptr, "Unsupported NT function called: NtAccessCheckAndAuditAlarm" , "Policy Tool Warning", MB_CANCELTRYCONTINUE);
        switch (res) {
        case IDCANCEL: return STATUS_NOT_IMPLEMENTED;
        case IDTRYAGAIN: return STATUS_SUCCESS;
        case IDCONTINUE: return TrueNtAccessCheckAndAuditAlarm(A, B, C, D, E, F, G, H, I, J, K);
        default: abort();
        }
    }

    static
    NTSTATUS NTAPI DetourNtPrivilegeCheck(IN HANDLE, PRIVILEGE_SET, OUT PBOOLEAN Result) {
        *Result = TRUE;
        return STATUS_SUCCESS;
    }

public:
    AdminImpersonate() {
        TrueIsUserAnAdmin = IsUserAnAdmin;
        TrueCheckTokenMembership = CheckTokenMembership;
        TrueNtOpenFile = reinterpret_cast<decltype(&NtOpenFile)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtOpenFile"));
        TrueNtCreateFile = reinterpret_cast<decltype(&NtCreateFile)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtCreateFile"));
        TrueNtAccessCheck = reinterpret_cast<decltype(&NtAccessCheck)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtAccessCheck"));
        TrueNtAccessCheckAndAuditAlarm = reinterpret_cast<decltype(&NtAccessCheckAndAuditAlarm)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtAccessCheckAndAuditAlarm"));
        TrueNtPrivilegeCheck = reinterpret_cast<decltype(&NtPrivilegeCheck)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtPrivilegeCheck"));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtOpenFile), reinterpret_cast<void*>(DetourNtOpenFile));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtCreateFile), reinterpret_cast<void*>(DetourNtCreateFile));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueIsUserAnAdmin), reinterpret_cast<void*>(DetourIsUserAnAdmin));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(DetourCheckTokenMembership));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtAccessCheck), reinterpret_cast<void*>(DetourNtAccessCheck));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtAccessCheckAndAuditAlarm), reinterpret_cast<void*>(DetourNtAccessCheckAndAuditAlarm));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtPrivilegeCheck), reinterpret_cast<void*>(DetourNtPrivilegeCheck));
    }
    ~AdminImpersonate() {
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtOpenFile), reinterpret_cast<void*>(DetourNtOpenFile));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtCreateFile), reinterpret_cast<void*>(DetourNtCreateFile));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueIsUserAnAdmin), reinterpret_cast<void*>(DetourIsUserAnAdmin));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(DetourCheckTokenMembership));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtAccessCheck), reinterpret_cast<void*>(DetourNtAccessCheck));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtAccessCheckAndAuditAlarm), reinterpret_cast<void*>(DetourNtAccessCheckAndAuditAlarm));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtPrivilegeCheck), reinterpret_cast<void*>(DetourNtPrivilegeCheck));
    }
};