#pragma once
#include <vector>
#include "ProcessAccessHelp.h"

enum IATReferenceType {
	IAT_REFERENCE_PTR_JMP,
	IAT_REFERENCE_PTR_CALL,
	IAT_REFERENCE_DIRECT_JMP,
	IAT_REFERENCE_DIRECT_CALL
};

class IATReference
{
public:
	DWORD_PTR addressVA; //Address of reference
	DWORD_PTR targetPointer; //Place inside IAT
	DWORD_PTR targetAddressInIat; //WIN API?
	IATReferenceType type;
};


class IATReferenceScan
{
public:

	IATReferenceScan()
	{
		IatAddressVA = 0;
		IatSize = 0;
		ImageBase = 0;
		ImageSize = 0;
		iatBackup = 0;
		ScanForDirectImports = false;
		ScanForNormalImports = true;
	}

	~IATReferenceScan()
	{
		iatReferenceList.clear();
		iatDirectImportList.clear();

		if (iatBackup)
		{
			free(iatBackup);
		}
	}

	bool ScanForDirectImports;
	bool ScanForNormalImports;


	void startScan(DWORD_PTR imageBase, DWORD imageSize, DWORD_PTR iatAddress, DWORD iatSize);
	void patchNewIatBaseMemory(DWORD_PTR newIatBaseAddress);
	void patchNewIatBaseFile(DWORD_PTR newIatBaseAddress);
	void patchDirectImportsMemory();
	int numberOfFoundDirectImports();
private:
	DWORD_PTR NewIatAddressVA;
	DWORD_PTR NewIatAddressRVA;

	DWORD_PTR IatAddressVA;
	DWORD IatSize;
	DWORD_PTR ImageBase;
	DWORD ImageSize;

	DWORD_PTR standardImageBase;

	DWORD_PTR * iatBackup;

	std::vector<IATReference> iatReferenceList;
	std::vector<IATReference> iatDirectImportList;

	bool isPageExecutable( DWORD Protect );
	void scanMemoryPage( PVOID BaseAddress, SIZE_T RegionSize );
	void analyzeInstruction( _DInst * instruction );
	void findNormalIatReference( _DInst * instruction );
	void getIatEntryAddress( IATReference* ref );
	void findDirectIatReference( _DInst * instruction );
	bool isAddressValidMemory( DWORD_PTR address );
	void patchReferenceInMemory( IATReference * ref );
	void patchReferenceInFile( IATReference* ref );
	void patchDirectImportInMemory( IATReference * iter );
	DWORD_PTR lookUpIatForPointer( DWORD_PTR addr );
};

/*
PE64
----------
000000013FF82D87 FF15 137C0A00 CALL QWORD [RIP+0xA7C13]
Result: 000000014002A9A0

000000013F65C952 FF25 F8EA0B00 JMP QWORD [RIP+0xBEAF8]
Result: 000000013F71B450

PE32
----------
0120FFA5 FF15 8C6D2601 CALL DWORD [0x01266D8C]

0120FF52 FF25 D4722601 JMP DWORD [0x012672D4]
*/