pe重定位及修正

#include "stdafx.h"


#include   
#include   
  
class MemBlock {  
public:  
    MemBlock() : m_buff(NULL) {}  
  
    bool alloc(size_t size) {  
        m_buff = VirtualAlloc(m_buff, size,  MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);   
        if(m_buff == NULL)   {    
            printf("VirtualAlloc error: %d\n",GetLastError());    
            return false;  
        }    
        return true;  
    }  
  
    ~MemBlock() {  
        if (!m_buff) {  
            return;  
        }  
        if(!VirtualFree(m_buff, 0, MEM_RELEASE))  {    
            printf("VirutalFree error: %d",GetLastError());    
        }   
    }  
  
    void load(size_t offset, void const* buff, size_t size) {  
        memcpy((char*)m_buff + offset, buff, size);  
    }  
  
    void* get() { return m_buff; }  
  
private:  
    void* m_buff;  
};  
  
class PEFile {  
public:  
    PEFile(char const* path) {  
        HANDLE fileHandle= CreateFile(path,GENERIC_READ,FILE_SHARE_READ,NULL,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL,NULL);  
        if(INVALID_HANDLE_VALUE == fileHandle){  
            std::cout << "打开文件映射对象失败！" << std::endl;  
            return;  
        }  
        HANDLE mapHandle=    CreateFileMapping(fileHandle,NULL,PAGE_READONLY,0,0,NULL);  
        if(mapHandle==NULL){  
            std::cout << "打开文件映射对象失败！" << std::endl;;  
            return;  
        }  
        m_BaseAddress =(LPBYTE)MapViewOfFile(mapHandle,FILE_MAP_READ,0,0,0);  
  
        m_dosHead=(PIMAGE_DOS_HEADER)m_BaseAddress;  
  
        m_ntHead=(PIMAGE_NT_HEADERS)(m_BaseAddress+m_dosHead->e_lfanew);  
  
        if(m_dosHead->e_magic==IMAGE_DOS_SIGNATURE&&m_ntHead->Signature != IMAGE_NT_SIGNATURE){  
            std::cout << "打开文件映射对象失败！" << std::endl;;  
            return;  
        }  
  
        m_sectionBase = IMAGE_FIRST_SECTION(m_ntHead);


        m_imgSize = 0;
        m_codeVA = 0;
        PIMAGE_SECTION_HEADER sectAddr = m_sectionBase;  
        for (WORD n = m_ntHead->FileHeader.NumberOfSections; n  > 0; --n) {  
            //模拟内存对齐机制    
            DWORD dwBlockCount  = (sectAddr->SizeOfRawData + m_ntHead->OptionalHeader.SectionAlignment - 1) /  
                m_ntHead->OptionalHeader.SectionAlignment;    
  
            DWORD dwBeginVA = sectAddr->VirtualAddress;    
            DWORD dwEndVA = dwBeginVA + dwBlockCount * m_ntHead->OptionalHeader.SectionAlignment;   
            if (m_imgSize < dwEndVA) {  
                m_imgSize = dwEndVA;
            }
            if (strcmp((char*)sectAddr->Name, ".text") == 0) {  
                m_codeVA = sectAddr->VirtualAddress;  
            }
            sectAddr++;
        }
    }
    void loadImage() {  
        m_image.alloc(m_imgSize);  


        loadSection(m_image);


        importFun((DWORD)m_image.get());


        relocAddr((DWORD)m_image.get());


        printf("image base address %p\n", m_image.get());


        void (*main_fun)();
        main_fun = (void (*)())((DWORD)m_image.get() + m_ntHead->OptionalHeader.AddressOfEntryPoint);
        main_fun();  
    }


    void loadSection(MemBlock& image) {
        PIMAGE_SECTION_HEADER sectAddr = m_sectionBase;  
        WORD n = m_ntHead->FileHeader.NumberOfSections;  
        while  (n  > 0) {  
            //模拟内存对齐机制    
            DWORD dwBlockCount  = (sectAddr->SizeOfRawData + m_ntHead->OptionalHeader.SectionAlignment - 1) /  
                m_ntHead->OptionalHeader.SectionAlignment;    
  
            DWORD dwBeginVA = sectAddr->VirtualAddress;    
            DWORD dwEndVA = sectAddr->VirtualAddress + dwBlockCount * m_ntHead->OptionalHeader.SectionAlignment;   
  
            image.load(sectAddr->VirtualAddress, (void*)(m_BaseAddress + sectAddr->PointerToRawData), sectAddr->SizeOfRawData);


            sectAddr++;  
            n--;  
        }
    }


    void importFun(DWORD imageLoadAddr) {
        PIMAGE_IMPORT_DESCRIPTOR impTable = (PIMAGE_IMPORT_DESCRIPTOR)(imageLoadAddr + m_ntHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);  
        while ( impTable->TimeDateStamp != 0 || impTable->Name != 0) {  
            PCHAR ptr = (PCHAR)((DWORD)m_image.get()+ impTable->Name); 
            std::cout << "handle import dll " << ptr << std::endl;
            HMODULE handle = GetModuleHandle(ptr);  
            if (handle == NULL)  
                handle = LoadLibrary(ptr);  
            if (handle == NULL) {  
                std::cout << "A required .DLL file, " << ptr << ", was not found.\n";
                return;
            }  
            if (impTable->TimeDateStamp != -1)  
            {  
                impTable->ForwarderChain  = (DWORD)handle;  
                impTable->TimeDateStamp = 0xCDC31337;  


                PIMAGE_THUNK_DATA funAddr = (PIMAGE_THUNK_DATA)(imageLoadAddr + (DWORD)impTable->FirstThunk);
                PIMAGE_THUNK_DATA impFunThunk    ;
                if (impTable->Characteristics == 0)  
                    impFunThunk  = (PIMAGE_THUNK_DATA)(imageLoadAddr + (DWORD)impTable->FirstThunk);
                else   
                    impFunThunk = (PIMAGE_THUNK_DATA)(imageLoadAddr + impTable->OriginalFirstThunk);  
  
                while ( impFunThunk->u1.AddressOfData != NULL) {  
                    DWORD dw = 0;
                    if (IMAGE_SNAP_BY_ORDINAL(impFunThunk    ->u1.Ordinal)) {  
                        LPCSTR fnAddr = (LPCSTR)IMAGE_ORDINAL(impFunThunk->u1.Ordinal);
                        dw = (DWORD)GetProcAddress(handle, fnAddr); 
                        printf("\t[0x%p] %d => 0x%08x\n", &(funAddr->u1.Function), IMAGE_ORDINAL(impFunThunk->u1.Ordinal), dw);
                    } else {  
                        PIMAGE_IMPORT_BY_NAME pName = (PIMAGE_IMPORT_BY_NAME)(impFunThunk->u1.AddressOfData + imageLoadAddr);  
                        dw = (DWORD)GetProcAddress(handle, pName->Name);
                        printf("\t[0x%p] %s => 0x%08x\n", &(funAddr->u1.Function), pName->Name, dw);
                    }  
                    if (dw == 0) {
                        std::cout << "not find import fun";
                        return;
                    }  
  
                    funAddr->u1.Function = dw;  
                    funAddr++;


                    impFunThunk ++;
                }  
         
            } else { // -1 = new style bound import   
                printf("New style bound import: %s\n", ptr);  
            }  
            impTable++;          
        } 
    }


    void relocAddr(DWORD imageLoadAddr) {
        IMAGE_DATA_DIRECTORY* pDataDir = &(m_ntHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC]);  
        if (pDataDir->VirtualAddress == 0) {  
            std::cout << "not exist data directory IMAGE_DIRECTORY_ENTRY_BASERELOC" << std::endl;;  
        }  
        printf("reloc entry va: %08x  size: %d\n",  pDataDir->VirtualAddress, pDataDir->Size);  
        DWORD va = pDataDir->VirtualAddress;  
        while (va < pDataDir->VirtualAddress + pDataDir->Size) {  
            PIMAGE_BASE_RELOCATION relocBase = (PIMAGE_BASE_RELOCATION)(m_BaseAddress + rva2foa(va));  
            PWORD   typeBase = (PWORD)(relocBase + 1);  
            int typeCount = (relocBase->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);  
            printf("reloc block va: %08x   size: %d  count: %d\n", pDataDir->VirtualAddress,  relocBase->SizeOfBlock, typeCount);  
            for (int i = 0; i < typeCount; ++i)  
            {  
                DWORD offset = typeBase[i] & 0x0fff;  
                WORD flag = (typeBase[i] >> 12) & 0x000f;  
                /* 
                IMAGE_REL_BASED_ABSOLUTE (0) 使块按照32位对齐，位置为0。 
                IMAGE_REL_BASED_HIGH (1) 高16位必须应用于偏移量所指高字16位。 
                IMAGE_REL_BASED_LOW (2) 低16位必须应用于偏移量所指低字16位。 
                IMAGE_REL_BASED_HIGHLOW (3) 全部32位应用于所有32位。. 
                IMAGE_REL_BASED_HIGHADJ (4) 需要32位，高16位位于偏移量，低16位位于下一个偏移量数组元素，组合为一个带符号数，加上32位的一个数，然后加上8000然后把高16位保存在偏移量的16位域内。 
                IMAGE_REL_BASED_MIPS_JMPADDR (5)        Unknown 
                IMAGE_REL_BASED_SECTION (6)        Unknown 
                IMAGE_REL_BASED_REL32 (7)        Unknown 
                */  
                if (flag == IMAGE_REL_BASED_HIGHLOW) {  
                    // 对需要重定位的数据进行修正  
                    // 需要进行重定位的地址 =  模块基址+重定位基址+每个数据表示偏移量  
                    PDWORD relocAddress = (PDWORD)(imageLoadAddr + relocBase->VirtualAddress + offset);  
                    DWORD oldAddr = *relocAddress;  
                    // 修正方法: 重定位之后的值 = 需要进行重定位的地址值 - IMAGE_OPTINAL_HEADER中的基址 + 实际基址  
                    *relocAddress = oldAddr - m_ntHead->OptionalHeader.ImageBase + imageLoadAddr;  
  
                    printf("\tadress: 0x%p   value: 0x%p => 0x%p \n", relocAddress,  oldAddr, *relocAddress);  
                }  
            }  
            va += relocBase->SizeOfBlock;  
        } 
    }


    size_t rva2foa(size_t rva)    
    {    
        PIMAGE_SECTION_HEADER sectAddr = m_sectionBase;  
        WORD n = m_ntHead->FileHeader.NumberOfSections;  
        while  (n  > 0) {  
            //模拟内存对齐机制    
            DWORD dwBlockCount  = (sectAddr->SizeOfRawData + m_ntHead->OptionalHeader.SectionAlignment - 1) /  
                m_ntHead->OptionalHeader.SectionAlignment;    
  
            DWORD dwBeginVA = sectAddr->VirtualAddress;    
            DWORD dwEndVA = sectAddr->VirtualAddress + dwBlockCount * m_ntHead->OptionalHeader.SectionAlignment;    
            //如果stRVA在某个区段中    
            if (rva >= dwBeginVA && rva < dwEndVA)   {    
                return rva - dwBeginVA + sectAddr->PointerToRawData;    
            }    
            else if (rva < dwBeginVA) {//在文件头中直接返回    
                return rva;    
            }   
            sectAddr++;  
            n--;  
        }   
        std::cout << "not find rva " <<  rva << std::endl;;  
        return 0;    
    } 






    void analyseReloc() {  
        IMAGE_DATA_DIRECTORY* pDataDir = &(m_ntHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC]);  
        if (pDataDir->VirtualAddress == 0) {  
            std::cout << "not exist data directory IMAGE_DIRECTORY_ENTRY_BASERELOC" << std::endl;;  
        }  
        printf("reloc entry va: %08x  size: %d\n",  pDataDir->VirtualAddress, pDataDir->Size);  
        DWORD va = pDataDir->VirtualAddress;  
        while (va < pDataDir->VirtualAddress + pDataDir->Size) {  
            PIMAGE_BASE_RELOCATION relocBase = (PIMAGE_BASE_RELOCATION)(m_BaseAddress + rva2foa(va));  
            PWORD   typeBase = (PWORD)(relocBase + 1);  
            int typeCount = (relocBase->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);  
            printf("reloc block va: %08x   size: %d  count: %d\n", pDataDir->VirtualAddress,  relocBase->SizeOfBlock, typeCount);  
            for (int i = 0; i < typeCount; ++i)  
            {  
                DWORD offset = typeBase[i] & 0x0fff;  
                offset += rva2foa(relocBase->VirtualAddress);  
                WORD flag = (typeBase[i] >> 12) & 0x000f;  
                static char* flag_str[] = {  
                    "REL_BASED_ABSOLUTE",  
                    "REL_BASED_HIGH",  
                    "REL_BASED_LOW",  
                    "REL_BASED_HIGHLOW",  
                    "REL_BASED_HIGHADJ",  
                    "REL_BASED_MACHINE_SPECIFIC_5",  
                    "REL_BASED_RESERVED",  
                    "REL_BASED_MACHINE_SPECIFIC_7",  
                    "REL_BASED_MACHINE_SPECIFIC_8",  
                    "REL_BASED_MACHINE_SPECIFIC_9",  
                    "REL_BASED_DIR64",  
                    "UNKOWN",  
                    "UNKOWN",  
                    "UNKOWN",  
                    "UNKOWN",  
                    "UNKOWN"  
                };  
                PDWORD addr = (PDWORD)(m_BaseAddress + offset);  
                printf("\trva: %08x   foa: %08x  addr: %08x  %s\n", offset,  rva2foa(offset), *addr, flag_str[flag]);  
            }  
            va += relocBase->SizeOfBlock;  
        }  
    }  
  
    void analyseImport() {  
        IMAGE_DATA_DIRECTORY* pDataDir = m_ntHead->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT;  
        if (pDataDir->VirtualAddress == 0) {  
            std::cout << "not exist data directory IMAGE_DIRECTORY_ENTRY_IMPORT";  
        }  
    }  
  
private:  
    LPBYTE m_BaseAddress;  
    PIMAGE_DOS_HEADER m_dosHead;  
    PIMAGE_NT_HEADERS m_ntHead;  
    PIMAGE_SECTION_HEADER m_sectionBase;  
    MemBlock m_image;  


    size_t m_imgSize;
    size_t m_codeVA;
};  






int _tmain(int argc, _TCHAR* argv[])
{
    PEFile pe(argv[1]);    
    pe.loadImage() ;
	return 0;
}