#include        "defs.h"

void    peRebaseDll(__in void *lpBuffer, __in ULONG_PTR pNewDllBase){
        PIMAGE_DOS_HEADER       pmz;
        PPEHEADER32             pe32;
        PSECTION_HEADER         psection;
        PIMAGE_BASE_RELOCATION  p_reloc;    
        ULONG_PTR               delta, apply_reloc;
        ULONG                   dwRelocSize, dwRelocChunkSize;
        PUSHORT                 preloc;
        DWORD                   dwOldProt, dwOldProtNext;
        
        
        pmz = (PIMAGE_DOS_HEADER)lpBuffer;
        pe32= (PPEHEADER32)((ULONG_PTR)lpBuffer + pmz->e_lfanew);
        
        if (pe32->pe_reloc == 0) return;
                
        p_reloc = (PIMAGE_BASE_RELOCATION)(pe32->pe_reloc + (ULONG_PTR)lpBuffer);

        dwRelocSize = pe32->pe_relocsize;
        delta = (ULONG_PTR)pNewDllBase - pe32->pe_imagebase;
        
        while (dwRelocSize){
                preloc      = (PUSHORT)((ULONG_PTR)p_reloc + sizeof(IMAGE_BASE_RELOCATION));
                dwRelocChunkSize = p_reloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION);
                apply_reloc = (ULONG_PTR)lpBuffer + p_reloc->VirtualAddress;
                ntVirtualProtect((PVOID)((ULONG_PTR)lpBuffer + p_reloc->VirtualAddress), 0x1000, PAGE_EXECUTE_READWRITE, &dwOldProt);
                ntVirtualProtect((PVOID)((ULONG_PTR)lpBuffer + p_reloc->VirtualAddress + 0x1000), 0x1000, PAGE_EXECUTE_READWRITE, &dwOldProtNext);
                while (dwRelocChunkSize){
                        if (((*preloc & 0xF000) >> 12) == 0x3)
                                *(PULONG_PTR)(apply_reloc + (*preloc & 0xFFF)) += delta;
                        
                        preloc++;
                        dwRelocChunkSize -= sizeof(USHORT);        
                }  
                ntVirtualProtect((PVOID)((ULONG_PTR)lpBuffer + p_reloc->VirtualAddress), 0x1000, dwOldProt, &dwOldProt);              
                ntVirtualProtect((PVOID)((ULONG_PTR)lpBuffer + p_reloc->VirtualAddress + 0x1000), 0x1000, dwOldProtNext, &dwOldProtNext);   
                dwRelocSize -= p_reloc->SizeOfBlock;
                p_reloc     = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)p_reloc + p_reloc->SizeOfBlock);
        }      
}


void *peRemapNtdllProcess(__in HANDLE hProcess, __in void *lpNewBase){
        NTSTATUS status;
        UNICODE_STRING us;
        OBJECT_ATTRIBUTES oa;
        HANDLE  hSection;
        PVOID   pntdll;
        SIZE_T  size;
        LARGE_INTEGER SectionOffset;  
        HANDLE  hFile;
        IO_STATUS_BLOCK IoStatus;
        
        RtlInitUnicodeString(&us, L"\\KnownDlls\\ntdll.dll");
        InitializeObjectAttributes(&oa, &us, OBJ_CASE_INSENSITIVE, NULL, NULL);
        
        status = NtOpenSection(&hSection,
                               SECTION_QUERY | SECTION_MAP_WRITE | SECTION_MAP_READ | SECTION_MAP_EXECUTE,
                               &oa);
        if (status != STATUS_SUCCESS){
                RtlInitUnicodeString(&us, g_wsNtdllNtPath); 
                InitializeObjectAttributes(&oa, &us, OBJ_CASE_INSENSITIVE, NULL, NULL);
                status = NtCreateFile(&hFile,
                                      FILE_EXECUTE | SYNCHRONIZE,
                                      &oa,
                                      &IoStatus,
                                      NULL,
                                      0,
                                      FILE_SHARE_READ,
                                      FILE_OPEN,
                                      FILE_SYNCHRONOUS_IO_NONALERT | FILE_NON_DIRECTORY_FILE,
                                      NULL,
                                      0);
                if (status != STATUS_SUCCESS) return NULL;
                
                status = NtCreateSection(&hSection,
                                         SECTION_QUERY | SECTION_MAP_WRITE | SECTION_MAP_READ | SECTION_MAP_EXECUTE,
                                         NULL,
                                         NULL,
                                         PAGE_EXECUTE,
                                         SEC_IMAGE,
                                         hFile);
                NtClose(hFile);
                if (status != STATUS_SUCCESS) return NULL;                            
        }
                
        pntdll = (PVOID)lpNewBase; //NULL;
        SectionOffset.QuadPart = 0;
        size = 0;
        status = NtMapViewOfSection(hSection,
                                   hProcess, 
                                   &pntdll,
                                   0,
                                   0,
                                   0,
                                   &size,
                                   ViewShare,
                                   0,
                                   PAGE_READWRITE);
        NtClose(hSection);
        
        if (status != STATUS_SUCCESS && status != STATUS_IMAGE_NOT_AT_BASE){
                return NULL;        
        }
        return lpNewBase;  
}        


void *peRemapNtdll(__in void *lpNewBase){      
        return peRemapNtdllProcess(CURRENT_PROCESS, lpNewBase); 
}