#include"./pelib/source/PeFile.h"
#include <iostream>
#include <windows.h>
#include "recommom/Utils/utils.h"
using namespace std;
using namespace PeLib;

bool replaceString(string& src_str, size_t loc, size_t _size, string new_str) {
	bool b_ret = false;
	do {
		if(_size!=new_str.size())break;
		if(src_str.size()<_size+loc)break;

		src_str.erase(loc, _size);
		src_str.insert(loc, new_str);
		b_ret = true;
	} while (false);
	return b_ret;
}

int main_tes(int argc, char** argv) {
	PeLib::PeFile32 pef("E:\\svc\\test\\pelib\\pelib\\test.exe");
	string file_data = readFile(L"E:\\svc\\test\\pelib\\pelib\\test.exe");

	pef.readMzHeader();
	pef.readPeHeader();
	pef.readImportDirectory();
	PeLib::MzHeader mzh = pef.mzHeader();
	PeLib::PeHeader32 peh = pef.peHeader();
	//PeLib::P
	PeLib::ImportDirectory32 imp = pef.impDir();

	DWORD dw_add_size = 0x1000;
	string s_add_data = string(dw_add_size, 'A');
	dw_add_size = PeLib::alignOffset(dw_add_size, peh.getFileAlignment());
	s_add_data += string(dw_add_size - s_add_data.size(), 'B');

	//1.get exe stud and section data
	
	DWORD dw_pe_stud_size = mzh.getAddressOfPeHeader() -mzh.size();
	string pe_stud = file_data.substr(mzh.size(), dw_pe_stud_size);
	size_t size_header_align = peh.getSizeOfHeaders() - mzh.size() - peh.size() ;
	string s_header_align_add_data(size_header_align, '0');


	//2.modify e_lfanew / AddressOfEntryPoint / SizeOfHeaders / SizeOfImage
	mzh.setAddressOfPeHeader(mzh.getAddressOfPeHeader() + dw_add_size);
	peh.setAddressOfEntryPoint(peh.getAddressOfEntryPoint() + dw_add_size);
	peh.setSizeOfHeaders(peh.getSizeOfHeaders() + dw_add_size);
	peh.setSizeOfImage(peh.getSizeOfImage() + dw_add_size);

	for (size_t i = 0; i < peh.calcNumberOfRvaAndSizes(); i++){
		if(peh.m_inthHeader.dataDirectories[i].VirtualAddress==0)continue;
		peh.m_inthHeader.dataDirectories[i].VirtualAddress += dw_add_size;
		
	}
	for (size_t i = 0; i < peh.calcNumberOfSections(); i++) {
		if (peh.m_vIsh[i].VirtualAddress == 0)continue;
		peh.m_vIsh[i].VirtualAddress += dw_add_size;
		peh.m_vIsh[i].PointerToRawData += dw_add_size;

	}

	vector<PeLib::byte> mz_vec = {};
	vector<PeLib::byte> nt_vec = {};
	
	
	mzh.rebuild(mz_vec);
	peh.rebuild(nt_vec);

	string s_mz_header = string(mz_vec.begin(), mz_vec.end());
	string s_nt_header = string(nt_vec.begin(), nt_vec.end());
	

	string file_data_out = s_mz_header+pe_stud+s_add_data+s_nt_header+s_header_align_add_data;
	string mod_str = file_data.substr(peh.getSizeOfHeaders() - dw_add_size);
	file_data_out += mod_str;

	//=====================
	for (unsigned int i = 0; i < imp.m_vOldiid.size(); i++){
		if (imp.m_vOldiid[i].impdesc.OriginalFirstThunk == 0)continue;
		imp.m_vOldiid[i].impdesc.OriginalFirstThunk += dw_add_size;
		imp.m_vOldiid[i].impdesc.Name += dw_add_size;
		imp.m_vOldiid[i].impdesc.FirstThunk += dw_add_size;
	}
	vector<PeLib::byte> imp_vec = {};
	imp.rebuild(imp_vec, peh.m_inthHeader.dataDirectories[PELIB_IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress - dw_add_size);


	replaceString(file_data_out, peh.rvaToOffset(peh.getIddImportRva()), peh.getIddImportSize(), string(imp_vec.begin(), imp_vec.begin()+peh.getIddImportSize()));

	for (unsigned int i = 0; i < imp.getNumberOfFiles(PeLib::OLDDIR); i++){
		
		DWORD origin_offset = peh.rvaToOffset(imp.m_vOldiid[i].impdesc.OriginalFirstThunk);
		DWORD first_offset = peh.rvaToOffset(imp.m_vOldiid[i].impdesc.FirstThunk);
		DWORD count = 0;
		do {

			DWORD origin_data = 0;
			DWORD first_data = 0;
			memcpy(&origin_data, file_data_out.data() + origin_offset+count*sizeof(origin_data), sizeof(origin_data));
			memcpy(&first_data, file_data_out.data() + first_offset+count * sizeof(first_data), sizeof(first_data));
			if (origin_data == 0)break;
			origin_data += dw_add_size;
			first_data += dw_add_size;
			replaceString(file_data_out, origin_offset + count * sizeof(origin_data), sizeof(DWORD), string((char*)&origin_data, sizeof(DWORD)));
			replaceString(file_data_out, first_offset + count * sizeof(origin_data), sizeof(DWORD), string((char*)&first_data, sizeof(DWORD)));
			
			memcpy(&origin_data, file_data_out.data() +peh.rvaToOffset(origin_data) , sizeof(origin_data));
			if (origin_data == 0)break;
			++count;

		} while (1);

	}


	
	writeFile(L"E:\\svc\\test\\pelib\\pelib\\test_mod.exe", file_data_out);


	return 0;
}