#include        "defs.h"

void    ntdllFix(__in PVOID ntdllbase){
        PIMAGE_DOS_HEADER       pmz;
        PPEHEADER32             pe32;
        PIMAGE_EXPORT_DIRECTORY pexport;
        DWORD                   dwOldProt;
        PUCHAR                  pbuff;
        ULONG                   index;
        DWORD                   *AddressOfFunctions;
        ULONG_PTR               KiIntSystemCall;
        BOOL                    b_win8;
        
        KiIntSystemCall = (ULONG_PTR)getprocaddress(ntdllbase, "KiIntSystemCall");
        
        pmz = (PIMAGE_DOS_HEADER)ntdllbase;
        pe32= (PPEHEADER32)((ULONG_PTR)ntdllbase + pmz->e_lfanew);
        if (pe32->pe_export == 0) return;
                
        pexport = (PIMAGE_EXPORT_DIRECTORY)((ULONG_PTR)ntdllbase + pe32->pe_export);        
        AddressOfFunctions      = (DWORD *)((ULONG_PTR)ntdllbase + pexport->AddressOfFunctions);
        for (index = 0; index < pexport->NumberOfFunctions; index++){
                pbuff = (PUCHAR)((ULONG_PTR)ntdllbase + AddressOfFunctions[index]);        
                if (pbuff[0] == 0xB8 && pbuff[5] == 0xE8){
                        //windows 8        
                        b_win8 = TRUE;
                }else if (pbuff[0] == 0xB8 && pbuff[5] == 0xBA && pbuff[10] == 0xFF){
                        //xp/vista/win7        
                        b_win8 = FALSE;
                }else{
                        //win2k doesn't use sysenter... so no need to do anything here...
                        continue;
                }
                
                pbuff += 5;
                ntVirtualProtect(pbuff, 7, PAGE_EXECUTE_READWRITE, &dwOldProt);
                
                pbuff[0] = 0xE8;
                *(DWORD *)&pbuff[1] = KiIntSystemCall - ((ULONG_PTR)pbuff+5);
                if (b_win8 == FALSE){
                        pbuff[5] = 0x90;
                        pbuff[6] = 0x90;
                }
                
                ntVirtualProtect(pbuff, 7, dwOldProt, &dwOldProt);
        }
        
}