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

#include <windows.h>
#include <shlobj.h>
#include <detours.h>



class AdminImpersonate : public ModuleBase {
    inline static decltype(&IsUserAnAdmin) TrueIsUserAnAdmin;
    inline static decltype(&CheckTokenMembership) TrueCheckTokenMembership;

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


public:
    AdminImpersonate() {
        TrueIsUserAnAdmin = IsUserAnAdmin;
        TrueCheckTokenMembership = CheckTokenMembership;
        DetourAttach(&reinterpret_cast<PVOID&>(TrueIsUserAnAdmin), reinterpret_cast<void*>(IsUserAnAdmin));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(CheckTokenMembership));
    }
    ~AdminImpersonate() {
        DetourDetach(&reinterpret_cast<PVOID&>(TrueIsUserAnAdmin), reinterpret_cast<void*>(IsUserAnAdmin));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueCheckTokenMembership), reinterpret_cast<void*>(CheckTokenMembership));
    }
};