#include <iostream>
#include <Windows.h>
#include <Commdlg.h>
#include <String.h>
#include <winnt.h>
#include <imagehlp.h>
#include <vector>
#include <string>
#include <fstream>
#include <tchar.h>
#include <stdio.h>

using namespace std;

// Check if its 32bit or 64bit
WORD  fileType;

// Exported names
vector<string> names;

const vector<string> explode(const string& s, const char& c)
{
	string buff{ "" };
	vector<string> v;

	for (auto n : s)
	{
		if (n != c) buff += n; else
			if (n == c && buff != "") { v.push_back(buff); buff = ""; }
	}
	if (buff != "") v.push_back(buff);

	return v;
}

bool GetImageFileHeaders(string fileName, IMAGE_NT_HEADERS &headers)
{
	HANDLE fileHandle = CreateFile(
		fileName.c_str(),
		GENERIC_READ,
		FILE_SHARE_READ,
		nullptr,
		OPEN_EXISTING,
		FILE_ATTRIBUTE_NORMAL,
		0
	);
	if (fileHandle == INVALID_HANDLE_VALUE)
		return false;

	HANDLE imageHandle = CreateFileMapping(
		fileHandle,
		nullptr,
		PAGE_READONLY,
		0,
		0,
		nullptr
	);
	if (imageHandle == 0)
	{
		CloseHandle(fileHandle);
		return false;
	}

	void *imagePtr = MapViewOfFile(
		imageHandle,
		FILE_MAP_READ,
		0,
		0,
		0
	);
	if (imagePtr == nullptr)
	{
		CloseHandle(imageHandle);
		CloseHandle(fileHandle);
		return false;
	}

	PIMAGE_NT_HEADERS headersPtr = ImageNtHeader(imagePtr);
	if (headersPtr == nullptr)
	{
		UnmapViewOfFile(imagePtr);
		CloseHandle(imageHandle);
		CloseHandle(fileHandle);
		return false;
	}

	headers = *headersPtr;

	UnmapViewOfFile(imagePtr);
	CloseHandle(imageHandle);
	CloseHandle(fileHandle);

	return true;
}

void ListDLLFunctions(string sADllName, vector<string>& slListOfDllFunctions)
{
	DWORD *dNameRVAs(0);
	DWORD *dNameRVAs2(0);
	_IMAGE_EXPORT_DIRECTORY *ImageExportDirectory;
	unsigned long cDirSize;
	_LOADED_IMAGE LoadedImage;
	string sName;
	slListOfDllFunctions.clear();
	if (MapAndLoad(sADllName.c_str(), NULL, &LoadedImage, TRUE, TRUE))
	{



		ImageExportDirectory = (_IMAGE_EXPORT_DIRECTORY*)
			ImageDirectoryEntryToData(LoadedImage.MappedAddress,
				false, IMAGE_DIRECTORY_ENTRY_EXPORT, &cDirSize);

		if (ImageExportDirectory != NULL)
		{
			dNameRVAs = (DWORD *)ImageRvaToVa(LoadedImage.FileHeader,
				LoadedImage.MappedAddress,
				ImageExportDirectory->AddressOfNames, NULL);

			for (size_t i = 0; i < ImageExportDirectory->NumberOfNames; i++)
			{
				sName = (char *)ImageRvaToVa(LoadedImage.FileHeader,
					LoadedImage.MappedAddress,
					dNameRVAs[i], NULL);
				
				slListOfDllFunctions.push_back(sName);
			}

		}
		UnMapAndLoad(&LoadedImage);
	}
}

void GenerateDEF(string name, vector<string> names)
{
	std::fstream file;
	file.open(name + ".def", std::ios::out);
	file << "LIBRARY " << name << endl;
	file << "EXPORTS" << endl;

	// Loop them
	for (int i = 0; i < names.size(); i++)
	{
		if (fileType == IMAGE_FILE_MACHINE_AMD64) {
			file << "\t" << names[i] << "=PROXY_" << i << " @" << i + 1 << endl;  //no mangling needed; see below
		} else {
			file << "\t" << names[i] << "=_PROXY_" << i << "@0 @" << i + 1 << endl;  //mangled; see below
		}
		
	}

	file.close();
}

void GenerateMainCPP(string name, vector<string> names)
{
	std::fstream file;
	file.open(name + ".cpp", std::ios::out);
	file << "#include <windows.h>" << endl << endl;

	file << "HINSTANCE hLThis = 0;" << endl;
	file << "FARPROC p[" << names.size() << "];" << endl;
	file << "HINSTANCE hL = 0;" << endl << endl;


	file << "BOOL WINAPI DllMain(HINSTANCE hInst,DWORD reason,LPVOID)" << endl;
	file << "{" << endl;
	file << "\tif (reason == DLL_PROCESS_ATTACH)" << endl;
	file << "\t{" << endl;
	file << "\t\thLThis = hInst;" << endl;
	file << "\t\thL = LoadLibrary(L\".\\\\" << name << "_.dll\");" << endl;  //L prefix since visual studio defaults to wide chars nowadays
	file << "\t\tif(!hL) return false;" << endl;
	file << "\t}" << endl << endl;

	// Exports addresses
	for (int i = 0; i < names.size(); i++)
	{
		file << "\tp[" << i << "] = GetProcAddress(hL, \"" << names[i] << "\");" << endl;
	}


	file << "\tif (reason == DLL_PROCESS_DETACH)" << endl;
	file << "\t{" << endl;
	file << "\t\tFreeLibrary(hL);" << endl;
	file << "\t\treturn 1;" << endl;

	file << "\t}" << endl << endl;;
	file << "\treturn 1;" << endl;
	file << "}" << endl << endl;


	// Generate Exports
	file << "extern \"C\"" << endl << "{" << endl;
	
	if (fileType == IMAGE_FILE_MACHINE_AMD64)
	{
		file << "\tFARPROC PA = NULL;" << endl;
		file << "\tint RunASM();" << endl << endl;

		for (int i = 0; i < names.size(); i++)
		{
			file << "\tvoid " << "PROXY_" << i << "() { //" << names[i] << endl;  //mangled functions in x64 doesnt matter; functions wont demangle anyways
			file << "\t\tPA = p[" << i << "];" << endl;
			file << "\t\tRunASM();" << endl;
			file << "\t}" << endl;
		}
	} else {
		for (int i = 0; i < names.size(); i++)
		{
			file << "\t__declspec(naked) void __stdcall " << "PROXY_" << i << "() { //" << names[i] << endl;  //declspec naked stdcall translates to _PROXY_<i>@0 when mangled; and mangled proxy functions somehow means exports retain their name exactly as defined which is what we need
			file << "\t\t__asm" << endl << "\t\t {";
			file << "\t\t\tjmp p[" << i << " * 4]" << endl;
			file << "\t\t}" << endl;
			file << "\t}" << endl;
		}
	}
	
	file << "}" << endl;

	file.close();
}

void GenerateASM(string name)
{
	std::fstream file;
	file.open(name + "a.asm", std::ios::out);  //append suffix "a" to avoid .obj files clashing when compiling with msvc
	file << ".data" << endl;
	file << "extern PA : qword" << endl;
	file << ".code" << endl;
	file << "RunASM proc" << endl;
	file << "jmp qword ptr [PA]" << endl;
	file << "RunASM endp" << endl;
	file << "end" << endl;

	file.close();
}

int main()
{
	OPENFILENAME ofn;
	char szFile[100];

	// open a file name
	ZeroMemory(&ofn, sizeof(ofn));
	ofn.lStructSize = sizeof(ofn);
	ofn.hwndOwner = NULL;
	ofn.lpstrFile = szFile;
	ofn.lpstrFile[0] = '\0';
	ofn.nMaxFile = sizeof(szFile);
	ofn.lpstrFilter = "DLL\0*.dll\0";
	ofn.nFilterIndex = 1;
	ofn.lpstrFileTitle = NULL;
	ofn.nMaxFileTitle = 0;
	ofn.lpstrInitialDir = NULL;
	ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST | OFN_NOCHANGEDIR;

	cout << "ProxiFy - Copyright (C) Kristoffer Blasiak." << endl;
	cout << "Select the DLL you want to make a proxy for." << endl;

	if (!GetOpenFileName(&ofn))
	{
		MessageBox(NULL, "You have to choose a file.", "File not opened", MB_OK);
		return 0;
	}
		

	IMAGE_NT_HEADERS headers;
	if (GetImageFileHeaders(ofn.lpstrFile, headers))
	{
		fileType = headers.FileHeader.Machine;
	}

	if (fileType == IMAGE_FILE_MACHINE_AMD64)
		MessageBox(NULL, "64 bit file", "64 bit file", MB_OK);

	if (fileType == IMAGE_FILE_MACHINE_I386)
		MessageBox(NULL, "32 bit file", "32 bit file", MB_OK);

	// Get filename
	vector<std::string> fileNameV = explode(ofn.lpstrFile, '\\');
	std::string fileName = fileNameV[fileNameV.size() - 1];
	fileName = fileName.substr(0, fileName.size() - 4);

	// Get dll export names
	ListDLLFunctions(ofn.lpstrFile, names);

	
	// Create Def File
	GenerateDEF(fileName, names);
	GenerateMainCPP(fileName, names);

	if (fileType == IMAGE_FILE_MACHINE_AMD64)
		GenerateASM(fileName);


	return 0;
}