#include <tlhelp32.h>
#include <fstream>
#include <string>
#include <iostream>
#include <sstream>
#include <vector>
#include <dbghelp.h>
#include <delayimp.h>

using namespace System;
using namespace System::Collections::Generic;

typedef KeyValuePair<String^, bool> Function;
typedef List<Function> FunctionList;
typedef KeyValuePair<String^, FunctionList^> ModuleImports;
typedef Dictionary<String^, FunctionList^> ImportMap;

void ParseOldStyleDelayImports(ImportMap^ imports, ImgDelayDescr* importStruct, char* base, DWORD preferred)
{
	DWORD nameTableOffset = importStruct->rvaINT - preferred;
	DWORD dllNameOffset = importStruct->rvaDLLName - preferred;
	PIMAGE_THUNK_DATA pThunk = reinterpret_cast<PIMAGE_THUNK_DATA>(base + nameTableOffset);
	char* dllName = base + dllNameOffset;

	String^ thisDll = String(dllName).ToLower();
	FunctionList^ importsFromThisDll = gcnew FunctionList();
	imports->Add(thisDll, importsFromThisDll);

	while(pThunk->u1.AddressOfData)
	{
		String^ funcName = nullptr;
		// check whether this is imported by ordinal only
		if(pThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG)
		{
			// remove the ordinal flag to get the real ordinal number
			WORD ordinal = static_cast<WORD>(pThunk->u1.Ordinal & (~IMAGE_ORDINAL_FLAG));
			funcName = String::Concat(L"Ordinal ", ordinal.ToString());
		}
		else
		{
			// the function is exported by name, the address we have 
			// tells us where its IMAGE_IMPORT_BY_NAME struct is which will
			// give us the name of the function
			DWORD nameStructOffset = pThunk->u1.AddressOfData - preferred;
			PIMAGE_IMPORT_BY_NAME funcNameStruct = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(base + nameStructOffset);
			funcName = gcnew String(reinterpret_cast<PCHAR>(funcNameStruct->Name));
		}
		// add the function name to the vector for this dll
		importsFromThisDll->Add(Function(funcName, true));
		++pThunk; // move on to the next imported function
	}
}

void ParseDllImportEntry(ImportMap^ imports, char* modName, char* base, PIMAGE_THUNK_DATA functionList, bool isDelay)
{
	// Get the dll name and ensure its lowercase (to avoid having duplicates in the ExportMap)
	String^ thisDll = String(modName).ToLower();
	FunctionList^ importsFromThisDll = gcnew FunctionList();
	imports->Add(thisDll, importsFromThisDll);

	while(functionList->u1.AddressOfData)
	{
		String^ funcName = nullptr;
		// check whether this is imported by ordinal only
		if(functionList->u1.Ordinal & IMAGE_ORDINAL_FLAG)
		{
			// remove the ordinal flag to get the real ordinal number
			WORD ordinal = static_cast<WORD>(functionList->u1.Ordinal & (~IMAGE_ORDINAL_FLAG));
			funcName = String::Concat(L"Ordinal ", ordinal.ToString());
		}
		else
		{
			// the function is exported by name, the address we have 
			// tells us where its IMAGE_IMPORT_BY_NAME struct is which will
			// give us the name of the function
			PIMAGE_IMPORT_BY_NAME funcNameStruct = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(base + functionList->u1.AddressOfData);
			funcName = gcnew String(reinterpret_cast<PCHAR>(funcNameStruct->Name));
		}
		// add the function name to the vector for this dll
		importsFromThisDll->Add(Function(funcName, isDelay));
		++functionList; // move on to the next imported function
	}
}

void ParseImports(PVOID pImage, ImportMap^ imports)
{
	ULONG size = 0;
	PIMAGE_IMPORT_DESCRIPTOR importFiles = static_cast<PIMAGE_IMPORT_DESCRIPTOR>(ImageDirectoryEntryToData(pImage, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &size));
	if(importFiles)
	{
		char* base = static_cast<char*>(pImage);
		// The INT is an array of Import descriptor structures
		// ended by a struct that contains all zeros
		while(importFiles->OriginalFirstThunk != 0)
		{
			ParseDllImportEntry(imports, base + importFiles->Name, base, reinterpret_cast<PIMAGE_THUNK_DATA>(base + importFiles->OriginalFirstThunk), false);
			++importFiles; // move to the next image we import from
		}
	}
}

void ParseDelayLoad(PVOID pImage, ImportMap^ imports)
{
	ULONG size = 0;
	ImgDelayDescr* importFiles = static_cast<ImgDelayDescr*>(ImageDirectoryEntryToData(pImage, TRUE, IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT, &size));
	if(importFiles)
	{
		char* base = static_cast<char*>(pImage);
		// The INT is an array of Import descriptor structures
		// ended by a struct that contains all zeros
		while(importFiles->rvaIAT)
		{
			char* nameTable = 0;
			char* dllName = 0;
			if(importFiles->grAttrs & 1)
			{
				nameTable = base + importFiles->rvaINT;
				dllName = base + importFiles->rvaDLLName;
				ParseDllImportEntry(imports, dllName, base, reinterpret_cast<PIMAGE_THUNK_DATA>(nameTable), true);
			}
			else
			{
				PIMAGE_NT_HEADERS pHeaders = ImageNtHeader(pImage);
				DWORD preferred = pHeaders->OptionalHeader.ImageBase;
				// this can only be the case on 32-bit images, since only VC6
				// produces these types of import records
				ParseOldStyleDelayImports(imports, importFiles, base, preferred);
			}
			++importFiles; // move to the next image we import from
		}
	}
}

std::wstring StringToCppString(String^ string)
{
	cli::array<wchar_t>^ charArray = string->ToCharArray(0, string->Length);
	cli::pin_ptr<wchar_t> ptr = &charArray[0];
	return std::wstring(ptr);
}

LPVOID MapImageFile(LPCWSTR file)
{
	LPVOID lpMapping = NULL;
	HANDLE hFile = CreateFileW(file, GENERIC_READ, 0, NULL, OPEN_EXISTING, 0, NULL);
	if(hFile != INVALID_HANDLE_VALUE)
	{
		HANDLE hMapping = CreateFileMappingW(hFile, NULL, SEC_IMAGE | PAGE_READONLY, 0, 0, L"File");
		CloseHandle(hFile);
		if(hMapping != INVALID_HANDLE_VALUE)
		{
			lpMapping = MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);
			CloseHandle(hMapping);
		}
	}
	return lpMapping;
}

LPVOID MapImageFile(String^ file)
{
	return MapImageFile(StringToCppString(file).c_str());
}

ImportMap^ Retrieve(String^ fileName)
{
	ImportMap^ imports = gcnew ImportMap();
	LPVOID pImage = MapImageFile(fileName);
	if(pImage)
	{
		ParseImports(pImage, imports);
		ParseDelayLoad(pImage, imports);
		UnmapViewOfFile(pImage);
	}
	return imports;
}