#pragma once

#include <windows.h>
#include <vector>

class PeSection
{
public:
    WCHAR name[IMAGE_SIZEOF_SHORT_NAME + 1];
    DWORD_PTR virtualAddress;
    DWORD  virtualSize;
    DWORD  rawAddress;
    DWORD  rawSize;
    DWORD characteristics;

    bool isDumped;

    //highlight big virtual sizes -> anti-dump protection
    bool highlightVirtualSize()
    {
        return (virtualSize > 0x2000000);
    };
};

class PeFileSection
{
public:
    IMAGE_SECTION_HEADER sectionHeader;
    BYTE* data;
    DWORD dataSize;
    DWORD normalSize;

    PeFileSection()
    {
        ZeroMemory(&sectionHeader, sizeof(IMAGE_SECTION_HEADER));
        data = 0;
        dataSize = 0;
        normalSize = 0;
    }
};

class PeParser
{
public:
    PeParser(const WCHAR* file, bool readSectionHeaders = true);
    PeParser(const DWORD_PTR moduleBase, bool readSectionHeaders = true);
    PeParser(const DWORD_PTR iatVA, const DWORD_PTR FileMapVA, const HANDLE hFileMap, bool readSectionHeaders = true);

    ~PeParser();

    bool isValidPeFile();
    bool isPE64();
    bool isPE32();

    bool isTargetFileSamePeFormat();

    WORD getNumberOfSections();
    std::vector<PeFileSection> & getSectionHeaderList();

    bool hasExportDirectory();
    bool hasTLSDirectory();
    bool hasRelocationDirectory();
    bool hasOverlayData();

    DWORD getEntryPoint();

    bool getSectionNameUnicode(const int sectionIndex, WCHAR* output, const int outputLen);

    DWORD getSectionHeaderBasedFileSize();
    DWORD getSectionHeaderBasedSizeOfImage();

    bool readPeSectionsFromProcess();
    bool readPeSectionsFromFile();
    bool readPeSectionsFromMappedFile();
    bool savePeFileToDisk(const WCHAR* newFile);
    void removeDosStub();
    void alignAllSectionHeaders();
    void fixPeHeader();
    void setDefaultFileAlignment();
    bool dumpProcess(DWORD_PTR modBase, DWORD_PTR entryPoint, const WCHAR* dumpFilePath);
    bool dumpProcess(DWORD_PTR modBase, DWORD_PTR entryPoint, const WCHAR* dumpFilePath, std::vector<PeSection> & sectionList);

    void setEntryPointVa(DWORD_PTR entryPoint);
    void setEntryPointRva(DWORD entryPoint);

    static bool updatePeHeaderChecksum(const WCHAR* targetFile, DWORD fileSize);
    BYTE* getSectionMemoryByIndex(int index);
    DWORD getSectionMemorySizeByIndex(int index);
    int convertRVAToOffsetVectorIndex(DWORD_PTR dwRVA);
    DWORD_PTR convertOffsetToRVAVector(DWORD_PTR dwOffset);
    DWORD_PTR convertRVAToOffsetVector(DWORD_PTR dwRVA);
    DWORD_PTR convertRVAToOffsetRelative(DWORD_PTR dwRVA);
    DWORD getSectionAddressRVAByIndex(int index);

    PIMAGE_NT_HEADERS getCurrentNtHeader();
protected:
    PeParser();


    static const DWORD FileAlignmentConstant = 0x200;

    const WCHAR* filename;
    DWORD_PTR moduleBaseAddress;
    DWORD_PTR fileMapVA;

    /************************************************************************/
    /* PE FILE                                                              */
    /*                                                                      */
    /*  IMAGE_DOS_HEADER      64   0x40                                     */
    /*  IMAGE_NT_HEADERS32   248   0xF8                                     */
    /*  IMAGE_NT_HEADERS64   264  0x108                                     */
    /*  IMAGE_SECTION_HEADER  40   0x28                                     */
    /************************************************************************/

    PIMAGE_DOS_HEADER pDosHeader;
    BYTE* pDosStub;  //between dos header and section header
    DWORD dosStubSize;
    PIMAGE_NT_HEADERS32 pNTHeader32;
    PIMAGE_NT_HEADERS64 pNTHeader64;
    std::vector<PeFileSection> listPeSection;
    BYTE* overlayData;
    DWORD overlaySize;
    /************************************************************************/

    BYTE* fileMemory;
    BYTE* headerMemory;

    HANDLE hFile;
    DWORD fileSize;

    bool readPeHeaderFromFile(bool readSectionHeaders);
    bool readPeHeaderFromFileMapping(bool readSectionHeaders);
    bool readPeHeaderFromProcess(bool readSectionHeaders);

    bool hasDirectory(const int directoryIndex);
    bool getSectionHeaders();
    void getDosAndNtHeader(BYTE* memory, LONG size);
    DWORD calcCorrectPeHeaderSize(bool readSectionHeaders);
    DWORD getInitialHeaderReadSize(bool readSectionHeaders);
    bool openFileHandle();
    void closeFileHandle();
    void initClass();

    DWORD isMemoryNotNull(BYTE* data, int dataSize);
    bool openWriteFileHandle(const WCHAR* newFile);
    bool writeZeroMemoryToFile(HANDLE hFile, DWORD fileOffset, DWORD size);

    bool readPeSectionFromFile(DWORD readOffset, PeFileSection & peFileSection);
    bool readPeSectionFromProcess(DWORD_PTR readOffset, PeFileSection & peFileSection);
    bool readPeSectionFromFileMapping(DWORD readOffset, PeFileSection & peFileSection);

    bool readSectionFromProcess(const DWORD_PTR readOffset, PeFileSection & peFileSection);
    bool readSectionFromFile(const DWORD readOffset, PeFileSection & peFileSection);
    bool readSectionFrom(const DWORD_PTR readOffset, PeFileSection & peFileSection, const bool isProcess);


    DWORD_PTR getStandardImagebase();

    bool addNewLastSection(const CHAR* sectionName, DWORD sectionSize, BYTE* sectionData);
    DWORD alignValue(DWORD badValue, DWORD alignTo);

    void setNumberOfSections(WORD numberOfSections);

    void removeIatDirectory();
    bool getFileOverlay();
};