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

enum IATReferenceType
{
    IAT_REFERENCE_PTR_JMP,
    IAT_REFERENCE_PTR_CALL,
    IAT_REFERENCE_DIRECT_JMP,
    IAT_REFERENCE_DIRECT_CALL,
    IAT_REFERENCE_DIRECT_MOV,
    IAT_REFERENCE_DIRECT_PUSH,
    IAT_REFERENCE_DIRECT_LEA
};

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


class IATReferenceScan
{
public:

    IATReferenceScan()
    {
        apiReader = 0;
        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;
    bool JunkByteAfterInstruction;
    ApiReader* apiReader;

    void startScan(DWORD_PTR imageBase, DWORD imageSize, DWORD_PTR iatAddress, DWORD iatSize);
    //void patchNewIatBaseMemory(DWORD_PTR newIatBaseAddress);
    //void patchNewIatBaseFile(DWORD_PTR newIatBaseAddress);

    void patchNewIat(DWORD_PTR stdImagebase, DWORD_PTR newIatBaseAddress, PeParser* peParser);
    void patchDirectJumpTable(DWORD_PTR imageBase, DWORD directImportsJumpTableRVA, PeParser* peParser, BYTE* jmpTableMemory, DWORD newIatBase);
    void patchDirectImportsMemory(bool junkByteAfterInstruction);
    int numberOfFoundDirectImports();
    int numberOfFoundUniqueDirectImports();
    int numberOfDirectImportApisNotInIat();
    int getSizeInBytesOfJumpTableInSection();
    //static FileLog directImportLog;
    void printDirectImportLog();
    void changeIatBaseOfDirectImports(DWORD newIatBaseAddressRVA);
    DWORD addAdditionalApisToList();
private:
    DWORD_PTR NewIatAddressRVA;

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


    DWORD_PTR* iatBackup;

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

    void scanMemoryPage(PVOID BaseAddress, SIZE_T RegionSize);
    void analyzeInstruction(_DInst* instruction);
    void findNormalIatReference(_DInst* instruction);
    void getIatEntryAddress(IATReference* ref);
    void findDirectIatReferenceCallJmp(_DInst* instruction);
    bool isAddressValidImageMemory(DWORD_PTR address);
    void patchReferenceInMemory(IATReference* ref);
    void patchReferenceInFile(IATReference* ref);
    void patchDirectImportInMemory(IATReference* iter);
    DWORD_PTR lookUpIatForPointer(DWORD_PTR addr);
    void findDirectIatReferenceMov(_DInst* instruction);
    void findDirectIatReferencePush(_DInst* instruction);
    void checkMemoryRangeAndAddToList(IATReference* ref, _DInst* instruction);
    void findDirectIatReferenceLea(_DInst* instruction);
    void patchDirectImportInDump32(int patchPreFixBytes, int instructionSize, DWORD patchBytes, BYTE* memory, DWORD memorySize, bool generateReloc, DWORD patchOffset, DWORD sectionRVA);
    void patchDirectImportInDump64(int patchPreFixBytes, int instructionSize, DWORD_PTR patchBytes, BYTE* memory, DWORD memorySize, bool generateReloc, DWORD patchOffset, DWORD sectionRVA);
    void patchDirectJumpTableEntry(DWORD_PTR targetIatPointer, DWORD_PTR stdImagebase, DWORD directImportsJumpTableRVA, PeParser* peParser, BYTE* jmpTableMemory, DWORD newIatBase);


};

/*
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]
*/