/*
 *
 * File Name: 
 *
 *		main.c
 *
 * Summary:
 *
 *		This file was created to be included within a 'disassembler' project for PE 
 *		image files running on x86 and x86-compatible processors.
 *
 *		File contains functions forming a framework for loading a PE file through 
 *		memory mapping, verifying if it's a valid PE and coordinating the process
 *		of disassembling
 * 
 *
 * Copyright (C) 2004, Isaac Sigasa [isigasa@ananzi.co.za]
 * All Rights Reserved
 *
 *
 *  
 * Redistribution and use in source and binary forms, with or without modification, 
 * are permitted provided that the following conditions are met:
 *
 *		-	Redistributions of source code must retain the above copyright notice, 
 *			this list of conditions and the following disclaimer. 
 *
 *		-	Redistributions in binary form must reproduce the above copyright notice, 
 *			this list of conditions and the following disclaimer in the documentation 
 *			and/or other materials provided with the distribution. 
 *
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, 
 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
 * ARE DISCLAIMED. 
 * IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 *
 */

#include <windows.h>
#include <winnt.h>
#include <string.h>
#include <stdio.h>
#include "disasm.h"


#define EXTRACT_DISASM			0x00010


BOOL FileExists(const char *strFileName)
{
	WIN32_FIND_DATA FindData;
	HANDLE HFind = FindFirstFile(strFileName,&FindData);
	BOOL Found = (HFind == INVALID_HANDLE_VALUE)? FALSE : TRUE;
	FindClose(HFind);
	return Found;
};


typedef union _ImageBase
{
	LPVOID ImageBase;
	PIMAGE_DOS_HEADER pDOSHeader;
}ImageBase, *PImageBase;


typedef struct _PEBFFDescriptor
{
	ImageBase UImageBase;
	PIMAGE_NT_HEADERS pImageNTHeader;
	PIMAGE_SECTION_HEADER pSectionHeader;
}PEBFFDescriptor, *PPEBFFDescriptor;


PBYTE RVAtoVA(PEBFFDescriptor *BFF, DWORD nRVA)
{
	DWORD i;

	if(nRVA == 0)
		return NULL;
	for(i = 0; i < BFF->pImageNTHeader->FileHeader.NumberOfSections; i++)
	{
		DWORD startVA = BFF->pSectionHeader[i].VirtualAddress;
		DWORD endVA = startVA + BFF->pSectionHeader[i].SizeOfRawData;

		if((nRVA >= startVA)&&(nRVA < endVA))
		{
			return (PBYTE)BFF->UImageBase.ImageBase + BFF->pSectionHeader[i].PointerToRawData 
				          + nRVA - BFF->pSectionHeader[i].VirtualAddress;
		}
	}
	return NULL;
};


typedef struct _FileMapping
{
	HANDLE HFile;						
	HANDLE HFileMapping;	
	LPVOID pBaseAddress;
}FileMapping, *PFileMapping;


DWORD AllocateFileMapping(const char *strFileName, FileMapping* pFileMapping)
{
	ZeroMemory(pFileMapping, sizeof(*pFileMapping));
	pFileMapping->HFile = CreateFile(strFileName,GENERIC_READ,FILE_SHARE_READ|FILE_SHARE_WRITE,NULL,OPEN_EXISTING,0,0);
	if(pFileMapping->HFile == INVALID_HANDLE_VALUE)
		return GetLastError();		
	pFileMapping->HFileMapping = CreateFileMapping(pFileMapping->HFile,0,PAGE_READONLY,0,0,NULL);
	if(!pFileMapping->HFileMapping)
		return GetLastError();
	pFileMapping->pBaseAddress = MapViewOfFile(pFileMapping->HFileMapping,FILE_MAP_READ,0,0,0);
	if(!pFileMapping->pBaseAddress)
		return GetLastError();
	return ERROR_SUCCESS;
};


void DeallocateFileMapping(FileMapping* pFileMapping)
{
	if(pFileMapping->HFileMapping)
		UnmapViewOfFile(pFileMapping->pBaseAddress);
	if(pFileMapping->HFileMapping)
		CloseHandle(pFileMapping->HFileMapping);
	if(pFileMapping->HFile)
		CloseHandle(pFileMapping->HFile);
};


DWORD InitializePEBFF(FileMapping *pFileMapping, PEBFFDescriptor* pBFF)
{
	ZeroMemory(pBFF,sizeof(*pBFF));
	pBFF->UImageBase.ImageBase = pFileMapping->pBaseAddress; 
	if(IsBadReadPtr(pBFF->UImageBase.ImageBase,sizeof(pBFF->UImageBase.ImageBase)))
		return ERROR_BAD_EXE_FORMAT;
	
	/* Check if we got a valid DOS signature */
	if(pBFF->UImageBase.pDOSHeader->e_magic != IMAGE_DOS_SIGNATURE)
		return ERROR_BAD_EXE_FORMAT;
	
	pBFF->pImageNTHeader = (PIMAGE_NT_HEADERS32)((char*)pBFF->UImageBase.ImageBase + pBFF->UImageBase.pDOSHeader->e_lfanew);
	if(IsBadReadPtr(pBFF->pImageNTHeader,sizeof(*pBFF->pImageNTHeader)))
		return ERROR_BAD_EXE_FORMAT;
	/* Check if we got a valid PE signature */
	if(pBFF->pImageNTHeader->Signature != IMAGE_NT_SIGNATURE)
		return ERROR_BAD_EXE_FORMAT;
	if(pBFF->pImageNTHeader->FileHeader.NumberOfSections)
		pBFF->pSectionHeader = (PIMAGE_SECTION_HEADER)((char*)(&pBFF->pImageNTHeader->OptionalHeader) + pBFF->pImageNTHeader->FileHeader.SizeOfOptionalHeader);
	else
		pBFF->pSectionHeader = NULL;

	return ERROR_SUCCESS;
};


int main(int argc, char *argv[])
{
	DWORD ret;
	DWORD i;
	FileMapping lFileMapping;
	char ErrorMessage[256];
	PBYTE pStart;	
	PBYTE pEnd;
	PBYTE pLoadAddress;
	PIMAGE_SECTION_HEADER pSection;
	PEBFFDescriptor BFF;

	ZeroMemory(ErrorMessage,sizeof(ErrorMessage));
	printf(	"PE Image Dumper\nCopyright (C) 2004, Isaac Sigasa [isigasa@ananzi.co.za].\n"
			"All Rights Reserved.\n\n");
	
	if(argc != 2)
	{
		printf("\nUsage:\n\tPED Filename");
		return -1;
	}
	if(!FileExists(argv[1]))
	{
		printf("\nError: file %s does not exist",argv[1]);
		return -1;
	}

	printf("\nProcessing file %s\n",argv[argc-1]);
	ret = AllocateFileMapping(argv[argc-1],&lFileMapping);
	if(ret != ERROR_SUCCESS)
	{
		FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM,0,ret,0,ErrorMessage,sizeof(ErrorMessage),0);
		printf("\nError %d: Failed to initialize file mapping,'%s'",ret,ErrorMessage);
		DeallocateFileMapping(&lFileMapping);
		return 1;
	};

	ret = InitializePEBFF(&lFileMapping,&BFF);
	if(ret != ERROR_SUCCESS)
	{
		strcpy(ErrorMessage,"Bad EXE or DLL format");
		printf("\nError %d: Failed to initialize BE BFF, '%s'",ret,ErrorMessage);
		DeallocateFileMapping(&lFileMapping);
		return 1;
	};

	// For now we are dealing with EXEs or DLLs, let's check if that's the case
	if(!((BFF.pImageNTHeader->FileHeader.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE) ||
		(BFF.pImageNTHeader->FileHeader.Characteristics & IMAGE_FILE_DLL)))
	{
		printf("\nOnly dealing with either EXECUTABLE or DLL images");
		DeallocateFileMapping(&lFileMapping);
		return 1;
	}; 
	if(BFF.pImageNTHeader->FileHeader.Characteristics & IMAGE_FILE_DLL)
		printf("\nFile type: DLL\n");
	else
	{
		if(BFF.pImageNTHeader->FileHeader.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE)
			printf("\nFile type: EXECUTABLE\n");
	}
	
	/* scan sections - if a section is marked as executable, disassemble it */
	for(i = 0, pSection = BFF.pSectionHeader; i < BFF.pImageNTHeader->FileHeader.NumberOfSections; i++,pSection++)
	{
		if(pSection->Characteristics & IMAGE_SCN_MEM_EXECUTE)
		{
			pLoadAddress = (PBYTE)BFF.pImageNTHeader->OptionalHeader.ImageBase + pSection->VirtualAddress;
			pStart = RVAtoVA(&BFF,pSection->VirtualAddress);
			pEnd = pStart + pSection->Misc.VirtualSize;			
			Disassemble((char*)pLoadAddress,OpSize32,pStart,pEnd);
			printf("\n");
		}
	};
		
	DeallocateFileMapping(&lFileMapping);
	return 0;
}