// 01Dasm.cpp : 定义控制台应用程序的入口点。
//

#include "stdafx.h"

using std::set ;
using std::vector ;

set<DWORD> setRelBaseAddr ; //放置重定位地址
vector<IMAGE_SECTION_HEADER> vetSetions ;//放置区段头信息，用于将RVA转换成RAWoffset
LPSTR lpMapFileBase = NULL ;//文件映射的基址

LONGLONG GetEip(DWORD& dwBase) ;//取得PE入口,参数用于保存PE中的imagebase。
DWORD GetBaseAddrAfterLoad(LPWSTR lpwszModuleName) ;//取得实际载入内存的基址
void AddRelBaseAddrToSet() ;//把需要重定位的地址放入集合中
INT AccessDeniedFilter(DWORD dwExceptionCode, DWORD dwBase, DWORD dwNew, LPDISASM lpDisasm) ;//SEH的过滤函数
int RAVToRawOffset(DWORD ImageBase, DWORD dwBase) ;

int _tmain(int argc, _TCHAR* argv[])
{
	DWORD dwNewBaseAddr = 0 ;
	DWORD dwDefaultBaseAddr = 0 ;
	AddRelBaseAddrToSet() ;//最初就遍历重定位表  放入集合
	DISASM stDasm = {} ;
	LONGLONG lEip = GetEip(dwDefaultBaseAddr) ;//得到eip的 rva
	_tprintf(_T("base addr: 0x%08x , new base addr: 0x%08x\n"),dwDefaultBaseAddr, dwNewBaseAddr = GetBaseAddrAfterLoad(L"01dasm.exe")) ;
	if (-1 == lEip)
	{
		return -1 ;
	}

	stDasm.EIP = lEip + dwNewBaseAddr;//加上新的基址就是内存偏移,不管是否与建议装载地址相同
	

	UINT nCount = 0 ;
	int nLen = -1 ;


	//_tsystem(_T("pause")) ;
	while(nCount ++ < 1500)
	{
		__try//进行SEH处理
		{
			nLen = Disasm(&stDasm) ;
			if (nLen != UNKNOWN_OPCODE)
			{
				_tprintf(_T("0x%08x:"),stDasm.EIP) ;
				puts(stDasm.CompleteInstr) ;
				stDasm.EIP += nLen ;
			}
			else
			{
				_tprintf(_T("DASM ERROR\n")) ;
				break;
			}
		}
		__except(AccessDeniedFilter(GetExceptionCode(), dwDefaultBaseAddr, dwNewBaseAddr, &stDasm))
		{
			_tprintf(_T("Unkown Error! \n")) ;
			break;
		}
	}
	_tsystem(_T("pause")) ;
	return 0;
}

LONGLONG GetEip(DWORD& dwBase)
{
	HANDLE hFile = INVALID_HANDLE_VALUE ;
	TCHAR tszFileName[MAX_PATH] = {0} ;
	IMAGE_DOS_HEADER iDosHeader = {0} ;
	IMAGE_NT_HEADERS iNTHeader = {0} ;
	DWORD dwBytesRead = 0 ;
	do 
	{
		GetModuleFileName(NULL, tszFileName, MAX_PATH) ;

		hFile = CreateFile(tszFileName,GENERIC_READ, FILE_SHARE_READ, 
			NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL) ;
		if (INVALID_HANDLE_VALUE == hFile)
		{
			break;
		}

		if(!ReadFile(hFile, &iDosHeader, sizeof(IMAGE_DOS_HEADER), &dwBytesRead, NULL))
		{
			break;
		}
		if (IMAGE_DOS_SIGNATURE != iDosHeader.e_magic)
		{
			break ;
		}

		SetFilePointer(hFile, iDosHeader.e_lfanew, NULL, FILE_BEGIN) ;

		if(!ReadFile(hFile, &iNTHeader, sizeof(IMAGE_NT_HEADERS), &dwBytesRead, NULL))
		{
			break;
		}
		if (IMAGE_NT_SIGNATURE != iNTHeader.Signature)
		{
			break;
		}
	} while (FALSE);

	if (hFile)
	{
		CloseHandle(hFile) ;
		dwBase = iNTHeader.OptionalHeader.ImageBase ;
		return (LONGLONG)(iNTHeader.OptionalHeader.AddressOfEntryPoint) ;
	}
	return -1 ;
}

DWORD GetBaseAddrAfterLoad(LPWSTR lpwszModuleName)
{
	DWORD dwPID = GetCurrentProcessId() ;

	HANDLE hTlHelp = INVALID_HANDLE_VALUE ;

	hTlHelp = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, dwPID) ;

	if (INVALID_HANDLE_VALUE == hTlHelp)
	{
		return 0 ;
	}

	MODULEENTRY32 me32 = {sizeof(MODULEENTRY32)} ;
	if (Module32First(hTlHelp, &me32))
	{
		do 
		{
			if (0 == lstrcmpiW(me32.szModule, lpwszModuleName))
			{
				CloseHandle(hTlHelp) ;
				return (DWORD)me32.modBaseAddr ;
			}
		} while (Module32Next(hTlHelp, &me32));
	}

	CloseHandle(hTlHelp) ;
	return 0 ;

}



void AddRelBaseAddrToSet()
{
	IMAGE_DOS_HEADER iDosHeader = {0} ;
	IMAGE_NT_HEADERS iNTHeader = {0} ;
	//
	TCHAR tszFileName[MAX_PATH + 1] = {0} ;
	if (!setRelBaseAddr.empty())
	{
		setRelBaseAddr.clear() ;
	}
	GetModuleFileName(NULL, tszFileName , MAX_PATH) ;
	HANDLE hFile ;
	HANDLE hFileMap ;

	do 
	{
		hFile = CreateFile(tszFileName,GENERIC_READ, FILE_SHARE_READ, 
			NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL) ;
		if (INVALID_HANDLE_VALUE == hFile)
		{
			break ;
		}

		hFileMap = CreateFileMapping(hFile,NULL,PAGE_READONLY,0,0,NULL) ;
		if (NULL == hFileMap)
		{
			break  ;
		}

		lpMapFileBase = (LPSTR)MapViewOfFile(hFileMap, FILE_MAP_READ, 0 , 0 , 0) ;

		if (NULL == lpMapFileBase)
		{
			break;
		}

		LPSTR lpTmpAddr = lpMapFileBase ;

		memcpy_s(&iDosHeader, sizeof(IMAGE_DOS_HEADER), lpTmpAddr, sizeof(IMAGE_DOS_HEADER)) ;
		lpTmpAddr += iDosHeader.e_lfanew ;
		memcpy_s(&iNTHeader, sizeof(IMAGE_NT_HEADERS), lpTmpAddr, sizeof(IMAGE_NT_HEADERS)) ;
		lpTmpAddr += sizeof(IMAGE_NT_HEADERS) ;

		vetSetions.clear() ;
		IMAGE_SECTION_HEADER iSectionHeader = {0} ;
		for(int i = 0 ; i < iNTHeader.FileHeader.NumberOfSections ; i ++)
		{
			ZeroMemory(&iSectionHeader, sizeof(IMAGE_SECTION_HEADER)) ;
			memcpy_s(&iSectionHeader,sizeof(IMAGE_SECTION_HEADER),lpTmpAddr,sizeof(IMAGE_SECTION_HEADER)) ;
			lpTmpAddr += sizeof(IMAGE_SECTION_HEADER) ;
			vetSetions.push_back(iSectionHeader) ;
		}

		DWORD dwSizeofBlock = 0 ;
		DWORD dwRelVirtrualAddr =  0; 

		dwRelVirtrualAddr = iNTHeader.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress ;
		lpTmpAddr = lpMapFileBase + RAVToRawOffset(dwRelVirtrualAddr + iNTHeader.OptionalHeader.ImageBase, iNTHeader.OptionalHeader.ImageBase) ;

		WORD wRelBlock = 0 ;
		DWORD dwAddrNeedRel = 0 ;
		while(dwRelVirtrualAddr = *((DWORD*)lpTmpAddr))
		{
			dwSizeofBlock = ( *((DWORD*)lpTmpAddr + 1)  - sizeof(IMAGE_BASE_RELOCATION)) / 2 ;
			lpTmpAddr += 8 ;
			for (DWORD i = 0 ; i < dwSizeofBlock ; i ++)
			{
				wRelBlock = *((WORD*)lpTmpAddr) ;
				lpTmpAddr += 2 ;
				if (0x03000 == (wRelBlock & 0x0f000))
				{
					dwAddrNeedRel = iNTHeader.OptionalHeader.ImageBase + dwRelVirtrualAddr + (DWORD)(wRelBlock & 0x0fff) ;
					setRelBaseAddr.insert(dwAddrNeedRel) ;
				}
			}
		}

	} while (FALSE);

	if (lpMapFileBase)
	{
		UnmapViewOfFile(lpMapFileBase) ;
	}

	if (hFileMap)
	{
		CloseHandle(hFileMap) ;
	}
	if (INVALID_HANDLE_VALUE != hFile)
	{
		CloseHandle(hFile) ;
	}

}

INT AccessDeniedFilter(DWORD dwExceptionCode, DWORD dwBase, DWORD dwNew, LPDISASM lpDisasm)
{
	if (EXCEPTION_ACCESS_VIOLATION != dwExceptionCode)
	{
		return EXCEPTION_EXECUTE_HANDLER ;
	}


	if (setRelBaseAddr.end() != setRelBaseAddr.find(lpDisasm->EIP) && (dwBase != dwNew))
	{
		lpDisasm->EIP = lpDisasm->EIP - dwBase + dwNew ;
		return EXCEPTION_CONTINUE_EXECUTION ;
	}

	return EXCEPTION_EXECUTE_HANDLER ;

}

int RAVToRawOffset(DWORD RVA, DWORD dwBaseAddr)
{
	for(vector<IMAGE_SECTION_HEADER>::iterator it = vetSetions.begin();
		it != vetSetions.end();it++)
	{
		if (RVA >= dwBaseAddr + it->VirtualAddress  && RVA <= dwBaseAddr + it->VirtualAddress + it->SizeOfRawData)
		{
			return it->PointerToRawData + RVA - it->VirtualAddress - dwBaseAddr ;
		}
	}
	return -1 ;
}