#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;

    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;
    }

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"));
        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*>(IsUserAnAdmin));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(CheckTokenMembership));
    }
    ~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*>(IsUserAnAdmin));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(CheckTokenMembership));
    }
};