/*

This file is part of HadesMem.

Copyright (C) 2011 Joshua Boyce (a.k.a. RaptorFactor).

<http://www.raptorfactor.com/> <raptorfactor@raptorfactor.com>



HadesMem is free software: you can redistribute it and/or modify

it under the terms of the GNU General Public License as published by

the Free Software Foundation, either version 3 of the License, or

(at your option) any later version.



HadesMem is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

GNU General Public License for more details.



You should have received a copy of the GNU General Public License

along with HadesMem.  If not, see <http://www.gnu.org/licenses/>.

*/



#pragma once



// Windows

#include <Windows.h>



// Hades

#include <HadesMemory/Fwd.hpp>

#include <HadesMemory/Error.hpp>

#include <HadesMemory/MemoryMgr.hpp>

#include <HadesMemory/PeLib/PeFile.hpp>



namespace Hades

{

  namespace Memory

  {

    // PE file NT headers

    class NtHeaders

    {

    public:

      // NT headers error class

      class Error : public virtual HadesMemError

      { };



      // Data directory entries

      enum DataDir

      {

        DataDir_Export, 

        DataDir_Import, 

        DataDir_Resource, 

        DataDir_Exception, 

        DataDir_Security, 

        DataDir_BaseReloc, 

        DataDir_Debug, 

        DataDir_Architecture, 

        DataDir_GlobalPTR, 

        DataDir_TLS, 

        DataDir_LoadConfig, 

        DataDir_BoundImport, 

        DataDir_IAT, 

        DataDir_DelayImport, 

        DataDir_COMDescriptor

      };



      // Constructor

      explicit NtHeaders(PeFile const& MyPeFile);



      // Get base of NT headers

      PVOID GetBase() const;



      // Whether signature is valid

      bool IsSignatureValid() const;



      // Ensure signature is valid

      void EnsureSignatureValid() const;



      // Get signature

      DWORD GetSignature() const;



      // Set signature

      void SetSignature(DWORD Signature) const;



      // Get machine

      WORD GetMachine() const;



      // Set machine

      void SetMachine(WORD Machine) const;



      // Get number of sections

      WORD GetNumberOfSections() const;



      // Set number of sections

      void SetNumberOfSections(WORD NumberOfSections) const;



      // Get time date stamp

      DWORD GetTimeDateStamp() const;



      // Set time date stamp

      void SetTimeDateStamp(DWORD TimeDateStamp) const;



      // Get pointer to symbol table

      DWORD GetPointerToSymbolTable() const;



      // Set pointer to symbol table

      void SetPointerToSymbolTable(DWORD PointerToSymbolTable) const;



      // Get number of symbols

      DWORD GetNumberOfSymbols() const;



      // Set number of symbols

      void SetNumberOfSymbols(DWORD NumberOfSymbols) const;



      // Get size of optional header

      WORD GetSizeOfOptionalHeader() const;



      // Set size of optional header

      void SetSizeOfOptionalHeader(WORD SizeOfOptionalHeader) const;



      // Get characteristics

      WORD GetCharacteristics() const;



      // Set characteristics

      void SetCharacteristics(WORD Characteristics) const;



      // Get magic

      WORD GetMagic() const;



      // Set magic

      void SetMagic(WORD Magic) const;



      // Get major linker version

      BYTE GetMajorLinkerVersion() const;



      // Set major linker version

      void SetMajorLinkerVersion(BYTE MajorLinkerVersion) const;



      // Get minor linker version

      BYTE GetMinorLinkerVersion() const;



      // Set major linker version

      void SetMinorLinkerVersion(BYTE MinorLinkerVersion) const;



      // Get minor linker version

      DWORD GetSizeOfCode() const;



      // Set major linker version

      void SetSizeOfCode(DWORD SizeOfCode) const;



      // Get minor linker version

      DWORD GetSizeOfInitializedData() const;



      // Set major linker version

      void SetSizeOfInitializedData(DWORD SizeOfInitializedData) const;



      // Get minor linker version

      DWORD GetSizeOfUninitializedData() const;



      // Set major linker version

      void SetSizeOfUninitializedData(DWORD SizeOfUninitializedData) const;



      // Get minor linker version

      DWORD GetAddressOfEntryPoint() const;



      // Set major linker version

      void SetAddressOfEntryPoint(DWORD AddressOfEntryPoint) const;



      // Get base of code

      DWORD GetBaseOfCode() const;



      // Set base of code

      void SetBaseOfCode(DWORD BaseOfCode) const;



#if defined(_M_IX86) 

      // Get base of data

      DWORD GetBaseOfData() const;



      // Set base of data

      void SetBaseOfData(DWORD BaseOfData) const;

#endif



      // Get base of code

      ULONG_PTR GetImageBase() const;



      // Set base of code

      void SetImageBase(ULONG_PTR ImageBase) const;



      // Get base of code

      DWORD GetSectionAlignment() const;



      // Set base of code

      void SetSectionAlignment(DWORD SectionAlignment) const;



      // Get base of code

      DWORD GetFileAlignment() const;



      // Set base of code

      void SetFileAlignment(DWORD FileAlignment) const;



      // Get base of code

      WORD GetMajorOperatingSystemVersion() const;



      // Set base of code

      void SetMajorOperatingSystemVersion(

        WORD MajorOperatingSystemVersion) const;



      // Get base of code

      WORD GetMinorOperatingSystemVersion() const;



      // Set base of code

      void SetMinorOperatingSystemVersion(

        WORD MinorOperatingSystemVersion) const;



      // Get base of code

      WORD GetMajorImageVersion() const;



      // Set base of code

      void SetMajorImageVersion(WORD MajorImageVersion) const;



      // Get base of code

      WORD GetMinorImageVersion() const;



      // Set base of code

      void SetMinorImageVersion(WORD MinorImageVersion) const;



      // Get base of code

      WORD GetMajorSubsystemVersion() const;



      // Set base of code

      void SetMajorSubsystemVersion(WORD MajorSubsystemVersion) const;



      // Get base of code

      WORD GetMinorSubsystemVersion() const;



      // Set base of code

      void SetMinorSubsystemVersion(WORD MinorSubsystemVersion) const;



      // Get base of code

      DWORD GetWin32VersionValue() const;



      // Set base of code

      void SetWin32VersionValue(DWORD Win32VersionValue) const;



      // Get size of image

      DWORD GetSizeOfImage() const;

      

      // Set size of image

      void SetSizeOfImage(DWORD SizeOfImage) const;



      // Get base of code

      DWORD GetSizeOfHeaders() const;



      // Set base of code

      void SetSizeOfHeaders(DWORD SizeOfHeaders) const;



      // Get base of code

      DWORD GetCheckSum() const;



      // Set base of code

      void SetCheckSum(DWORD CheckSum) const;



      // Get base of code

      WORD GetSubsystem() const;



      // Set base of code

      void SetSubsystem(WORD Subsystem) const;



      // Get base of code

      WORD GetDllCharacteristics() const;



      // Set base of code

      void SetDllCharacteristics(WORD DllCharacteristics) const;



      // Get base of code

      ULONG_PTR GetSizeOfStackReserve() const;



      // Set base of code

      void SetSizeOfStackReserve(ULONG_PTR SizeOfStackReserve) const;



      // Get base of code

      ULONG_PTR GetSizeOfStackCommit() const;



      // Set base of code

      void SetSizeOfStackCommit(ULONG_PTR SizeOfStackCommit) const;



      // Get base of code

      ULONG_PTR GetSizeOfHeapReserve() const;



      // Set base of code

      void SetSizeOfHeapReserve(ULONG_PTR SizeOfHeapReserve) const;



      // Get base of code

      ULONG_PTR GetSizeOfHeapCommit() const;



      // Set base of code

      void SetSizeOfHeapCommit(ULONG_PTR SizeOfHeapCommit) const;



      // Get base of code

      DWORD GetLoaderFlags() const;



      // Set base of code

      void SetLoaderFlags(DWORD LoaderFlags) const;



      // Get base of code

      DWORD GetNumberOfRvaAndSizes() const;



      // Set base of code

      void SetNumberOfRvaAndSizes(DWORD NumberOfRvaAndSizes) const;



      // Get base of code

      DWORD GetDataDirectoryVirtualAddress(DataDir MyDataDir) const;



      // Set base of code

      void SetDataDirectoryVirtualAddress(DataDir MyDataDir, 

        DWORD DataDirectoryVirtualAddress) const;



      // Get base of code

      DWORD GetDataDirectorySize(DataDir MyDataDir) const;



      // Set base of code

      void SetDataDirectorySize(DataDir MyDataDir, 

        DWORD DataDirectorySize) const;



      // Get raw NT headers

      IMAGE_NT_HEADERS GetHeadersRaw() const;



    private:

      // PE file

      PeFile m_PeFile;



      // Memory instance

      MemoryMgr m_Memory;



      // Base address

      PBYTE m_pBase;

    };

  }

}