

#include "PE_class.h"
#include <fstream>
#include <string>

using namespace std;








DWORD PE_file::Get_Size()
{
	return PE_size;
}

char* PE_file::GetBuffer()
{
	return buf;
}








PE_file::PE_file(char* name,DWORD extra_size)
{
	
	init =0;
	ifstream f(name,ios::binary);
	if(f==0)
	{
		printf("Can open input file\n");

		return;
	}
	f.seekg(0,ios::end);
	int size=f.tellg();
	f.seekg(0,ios::beg);
	buf=(char*)malloc(size+extra_size);
	PE_size=size+extra_size;
	int i=0;
	char ch;
	while(i<size)
	{
		f.get(ch);
		buf[i]=ch;
		i++;
	}
	f.close();
	if(buf[0] != 'M' || buf[1] != 'Z')
	{
		printf("input file isn`t correct PE file\n");
		return;
	}
	dos=(IMAGE_DOS_HEADER*)buf;
	PE_offset=dos->e_lfanew;
	fhead=(IMAGE_FILE_HEADER*)(buf+PE_offset+4);
	int size_opt=fhead->SizeOfOptionalHeader;
	opt_head=(IMAGE_OPTIONAL_HEADER32*)(buf+PE_offset+sizeof(IMAGE_FILE_HEADER)+4);
	section_offset=PE_offset+sizeof(IMAGE_FILE_HEADER)+ size_opt;
	init=true;
	cur_pos=0;
}

void PE_file::open(char* name,DWORD extra_size)
{
	
	init =0;
	ifstream f(name,ios::binary);
	if(f==0)
	{
		printf("Can open input file\n");

		return;
	}
	f.seekg(0,ios::end);
	int size=f.tellg();
	f.seekg(0,ios::beg);
	buf=(char*)malloc(size+extra_size);
	PE_size=size+extra_size;
	int i=0;
	char ch;
	while(i<size)
	{
		f.get(ch);
		buf[i]=ch;
		i++;
	}
	f.close();
	if(buf[0] != 'M' || buf[1] != 'Z')
	{
		printf("input file isn`t correct PE file\n");
		return;
	}
	dos=(IMAGE_DOS_HEADER*)buf;
	PE_offset=dos->e_lfanew;
	fhead=(IMAGE_FILE_HEADER*)(buf+PE_offset+4);
	int size_opt=fhead->SizeOfOptionalHeader;
	opt_head=(IMAGE_OPTIONAL_HEADER32*)(buf+PE_offset+sizeof(IMAGE_FILE_HEADER)+4);
	section_offset=PE_offset+sizeof(IMAGE_FILE_HEADER)+ size_opt;
	init=true;
	cur_pos=0;
}

PE_file::PE_file(HMODULE h)
{
	
	init =0;
	buf=(char*)h;
	/*if(buf[0] != 'M' || buf[1] != 'Z')
	{
	printf("input file isn`t correct PE file\n");
	return;
	}*/
	dos=(IMAGE_DOS_HEADER*)buf;
	PE_offset=dos->e_lfanew;
	fhead=(IMAGE_FILE_HEADER*)(buf+PE_offset+4);
	int size_opt=fhead->SizeOfOptionalHeader;
	opt_head=(IMAGE_OPTIONAL_HEADER32*)(buf+PE_offset+sizeof(IMAGE_FILE_HEADER)+4);
	section_offset=PE_offset+sizeof(IMAGE_FILE_HEADER)+ size_opt;
	init=true;
	cur_pos=0;
}

PE_file::PE_file()
{
	
}

void PE_file::open(HMODULE h,DWORD size)
{
	
init =0;
	buf=(char*)h;
	dos=(IMAGE_DOS_HEADER*)buf;
	PE_offset=dos->e_lfanew;
	fhead=(IMAGE_FILE_HEADER*)(buf+PE_offset+4);
	int size_opt=fhead->SizeOfOptionalHeader;
	opt_head=(IMAGE_OPTIONAL_HEADER32*)(buf+PE_offset+sizeof(IMAGE_FILE_HEADER)+4);
	section_offset=PE_offset+sizeof(IMAGE_FILE_HEADER)+ size_opt;
	init=true;
	cur_pos=0;
	PE_size=size;
	
}




DWORD PE_file::RVA_to_RAW(DWORD RVA)
{
	DWORD RAW=-1;
	int i=0;
	while(i<fhead->NumberOfSections)
	{
		sect_head=(IMAGE_SECTION_HEADER*)(buf+section_offset+4+(sizeof(IMAGE_SECTION_HEADER)*i));
		if(sect_head->VirtualAddress<=RVA && sect_head->VirtualAddress+sect_head->Misc.VirtualSize > RVA)
		{
			RAW=RVA-sect_head->VirtualAddress+sect_head->PointerToRawData;
		}
		i++;
	}
	return RAW;
}


DWORD PE_file::RAW_to_RVA(DWORD RAW)
{
	DWORD RVA=-1;
	int i=0;
	while(i<fhead->NumberOfSections)
	{
		sect_head=(IMAGE_SECTION_HEADER*)(buf+section_offset+4+(sizeof(IMAGE_SECTION_HEADER)*i));
		if(sect_head->PointerToRelocations<=RAW && sect_head->PointerToRawData+sect_head->SizeOfRawData > RAW)
		{
			RVA = RAW + sect_head->VirtualAddress-sect_head->PointerToRawData;
		}
		i++;
	}
	return RVA;
}

DWORD PE_file::VA_to_RVA(DWORD VA)
{
	return VA-opt_head->ImageBase;
}

DWORD PE_file::RVA_to_VA(DWORD RVA)
{
	return RVA+opt_head->ImageBase;
}



DWORD PE_file::Get_EP(void)
{
	return opt_head->AddressOfEntryPoint;
}

DWORD PE_file::Get_IB(void)
{
	return opt_head->ImageBase;
}

DWORD PE_file::Get_Section_Count(void)
{
	return fhead->NumberOfSections;
}

IMAGE_DOS_HEADER* PE_file::GetDosHeader(void)
{
	return dos;
}

IMAGE_FILE_HEADER* PE_file::GetFileHeader(void)
{
	return fhead;
}

IMAGE_OPTIONAL_HEADER32* PE_file::GetOptionalHeader(void)
{
	return opt_head;
}

IMAGE_SECTION_HEADER* PE_file::GetSectHeader(DWORD n)
{
	if(n>=fhead->NumberOfSections)
	{
		return 0;
	}
	return (IMAGE_SECTION_HEADER*)(buf+section_offset+4+(sizeof(IMAGE_SECTION_HEADER)*n));
}



