module Data.PE.Structures where
import Data.Word
import Data.PE.Utils
import Data.ByteString.Lazy
import Numeric
import Data.Binary
import Data.Binary.Get
import Data.Char
--import System.Time

-- |The over-arching container.  Holds the headers and a list of binary sections
data PEFile = PEFile {
    peHeader :: PEHeader
} deriving Show

data PEObject = PEObj {
    peObjHeader :: PEObjectHeader
} deriving Show

-- |The Binary Section container.  Holds names and containers.
data BinSection = BinSection {
    secname :: String,
    binSection :: ByteString
} deriving Show

data PEObjectHeader = PEObjHdr {
  objcoffhdr :: COFFHeader,
  objsectionTables :: [(SectionTable,ByteString)]
} deriving Show

-- |The Header section, holds entries for each header in the PE File
data PEHeader = PEHeader {
        msdosHeader :: MSDOSHeader,
        peSignature :: PESignature,
        coffHeader :: COFFHeader,
        standardFields :: StandardFields,
        windowsSpecFields :: WindowsSpecFields,
        dataDirectories :: [DirectoryEntry],
        sectionTables :: [(SectionTable,ByteString)]
--	sectionbytes :: [ByteString]
} deriving Show

data MSDOSHeader = MSDOSHeader {
        signature :: Word16,
        lastsize :: Word16,
        pagesInFile :: Word16,
        relocations :: Word16,
        headerSizeInParagraph :: Word16,
        minExtraParagraphs :: Word16,
        maxExtraParagraphs :: Word16,
        ss :: Word16,
        sp :: Word16,
        checksum :: Word16,
        ip :: Word16,
        cs :: Word16,
        relocTableOffset :: Word16,
        overlayNumber :: Word16,
        --4 * short reserved spaces
        oemIdentifier :: Word16,
        oemInformation :: Word16,
        --10 * short reserved spaces
        offset :: Word32
}
instance Show MSDOSHeader where
        show header = "Signature: " ++ (show.encode $ signature header) ++ "\n"
                ++ "Last Page Size: " ++ (show $ lastsize header) ++ "\n"
                ++ "Number of Pages: " ++ (show $ pagesInFile header) ++ "\n"
                ++ "Relocations: " ++ (show $ relocations header) ++ "\n"
                ++ "Header Size in Paragraphs: " ++ (show $ headerSizeInParagraph header) ++ "\n"
                ++ "Min Extra Paragraphs: " ++ (show $ minExtraParagraphs header) ++ "\n"
                ++ "Max Extra Paragraphs: " ++ (show $ maxExtraParagraphs header) ++ "\n"
                ++ "Stack Segment: 0x" ++ (showHex (ss header) "") ++ "\n"
                ++ "Stack Pointer: 0x" ++ (showHex (sp header) "") ++ "\n"
                ++ "File checksum: " ++ (show $ checksum header) ++ "\n"
                ++ "Code Segment: " ++ (show $ cs header) ++ "\n"
                ++ "Instruction Pointer: " ++ (show $ ip header) ++ "\n"
                ++ "Relocation Offset: " ++ (show $ relocTableOffset header) ++ "\n"
                ++ "Overlay Number: " ++ (show $ overlayNumber header) ++ "\n"
                ++ "OEM Identifier: 0x" ++ (showHex (oemIdentifier header) "") ++ "\n"
                ++ "OEM Information: 0x" ++(showHex (oemInformation header) "") ++ "\n"
                ++ "PE Header Offset: " ++ (show $ offset header) ++ "\n"


data PESignature = PESignature {
        pesignature :: Word32 --0x00004550
}
instance Show PESignature where
        show (PESignature sig) = "PE-Signature: 0x" ++ (showHex sig "") ++ "\n"

data COFFHeader = COFFHeader {
        targetMachine :: MachineType,
        numberOfSections :: Word16,  --IMPORTANT
        timeDateStamp :: Word32,
        pointerToSymbolTable :: Word32, --0 for image
        numberOfSymbols :: Word32, --0 for image
        sizeofOptionalHeaders :: Word16,
        coffCharacteristics :: Word16
}
instance Show COFFHeader where
        show hdr = "Target Machine: " ++ (show $ targetMachine hdr) ++"\n"
                ++ "Number of Sections: " ++ (show (numberOfSections hdr)) ++"\n"
                ++ "Timestamp: " ++ (show $ timeDateStamp $ hdr) ++ "\n"
                ++ "Symbol Table Pointer: 0x" ++ (showHex (pointerToSymbolTable hdr) "") ++ "\n"
                ++ "Number of Symbols: " ++ (show $ numberOfSymbols hdr) ++ "\n"
                ++ "Size of Optional Headers: " ++ (show $ sizeofOptionalHeaders hdr) ++ "\n"
                ++ "COFF Characteristics: 0x" ++ (showHex (coffCharacteristics hdr) "") ++ "\n"

data StandardFields = StandardFields {
        standardSig :: Word16, -- Should be 0x10B or 0x20B if PE32+
        lnMajorVersion :: Word8,
        lnMinorVersion :: Word8,
        sizeOfCode :: Word32,
        sizeOfInitializedData :: Word32,
        sizeOfUninitData :: Word32,
        addressOfEntryPoint :: Word32,
        baseOfCode :: Word32,
        baseOfData :: Word32
} | SFPlus { standardSig :: Word16,
              lnMajorVersion :: Word8,
              lnMinorVersion :: Word8,
              sizeOfCode :: Word32,
              sizeOfInitializedData :: Word32,
              sizeOfUninitData :: Word32,
              addressOfEntryPoint :: Word32,
                    baseOfCode :: Word32 }

instance Show StandardFields where
        show sf = "Signature: 0x" ++ (showHex (standardSig sf) "") ++ "\n"
                ++ "Linker Major Version: " ++ (show $ lnMajorVersion sf) ++ "\n"
                ++ "Linker Minor Version: " ++ (show $ lnMinorVersion sf) ++ "\n"
                ++ "Size of Code: " ++ (show $ sizeOfCode sf) ++ "\n"
                ++ "Size of Initialized Data: " ++ (show $ sizeOfInitializedData sf) ++ "\n"
                ++ "Size of Un-initialized Data: " ++ (show $ sizeOfUninitData sf) ++ "\n"
                ++ "Entry Point Address: 0x" ++ (showHex (addressOfEntryPoint sf) "") ++ "\n"
                ++ "Code base Address: 0x" ++ (showHex (baseOfCode sf) "") ++ "\n"
                -- ++"Data base Address: 0x"++(showHex (baseOfData sf) "")++"\n"

data WindowsSpecFields = WindowsSpecFields {
        imageBase :: Word32,
        sectionAlignment :: Word32,
        fileAlignment :: Word32,
        majorOSVersion :: Word16,
        minorOSVersion :: Word16,
        majorImageVersion :: Word16,
        minorImageVersion :: Word16,
        majorSubSystemVersion :: Word16,
        minorSubSystemVersion :: Word16,
        win32VersionValue :: Word32,
        sizeOfImage :: Word32,
        sizeOfHeaders :: Word32,
        checkSum32 :: Word32,
        checkSum16 :: Word16,
        dllCharacteristics :: Word16,
        sizeOfStackReserve :: Word32,
        sizeOfStackCommit :: Word32,
        sizeOfHeapReserve :: Word32,
        sizeOfHeapCommit :: Word32,
        loaderFlags :: Word32,
        numberOfRVAandSizes :: Word32
} | WSFPlus { imgBase :: Word64,
            sectionAlignment :: Word32,
            fileAlignment :: Word32,
            majorOSVersion :: Word16,
            minorOSVersion :: Word16,
            majorImageVersion :: Word16,
            minorImageVersion :: Word16,
            majorSubSystemVersion :: Word16,
            minorSubSystemVersion :: Word16,
            win32VersionValue :: Word32,
            sizeOfImage :: Word32,
            sizeOfHeaders :: Word32,
            checkSum32 :: Word32,
            checkSum16 :: Word16,
            dllCharacteristics :: Word16,
            szOfStackReserve :: Word64,
            szOfStackCommit :: Word64,
            szOfHeapReserve :: Word64,
            szOfHeapCommit :: Word64,
            loaderFlags :: Word32,
            numberOfRVAandSizes :: Word32 }

--instance Binary WindowsSpecFields where
--  get = 

instance Show WindowsSpecFields where
        show hdr = "Image Base: 0x" ++ (showHex (imageBase hdr) "") ++ "\n"
                ++ "Section Alignment: 0x" ++ (showHex (sectionAlignment hdr) "") ++ "\n"
                ++ "File Alignment: 0x" ++ (showHex (fileAlignment hdr) "") ++ "\n"
                ++ "Major OS Version: " ++ (show $ majorOSVersion hdr) ++ "\n"
                ++ "Minor OS Version: " ++ (show $ minorOSVersion hdr) ++ "\n"
                ++ "Major Subsystem Version: " ++ (show $ majorSubSystemVersion hdr) ++ "\n"
                ++ "Minor Subsystem Version: " ++ (show $ minorSubSystemVersion hdr) ++ "\n"
                ++ "Win 32 Version Value: " ++ (show $ win32VersionValue hdr) ++ "\n"
                ++ "Size of Image: " ++ (show $ sizeOfImage hdr) ++ "\n"
                ++ "Size of Headers: " ++ (show $ sizeOfHeaders hdr) ++ "\n"
                ++ "Checksum 32: " ++ (show $ checkSum32 hdr) ++ "\n"
                ++ "Checksum 15: " ++ (show $ checkSum16 hdr) ++ "\n"
                ++ "DLL Characteristics: 0x" ++ (showHex (dllCharacteristics hdr) "") ++ "\n"
                ++ "Size of Stack Reserved: " ++ (show $ sizeOfStackReserve hdr) ++ "\n"
                ++ "Size of Stack Commit: " ++ (show $ sizeOfStackCommit hdr) ++ "\n"
                ++ "Size of Heap Reserved: " ++ (show $ sizeOfHeapReserve hdr) ++ "\n"
                ++ "Size of Heap Commit: " ++ (show $ sizeOfHeapCommit hdr) ++ "\n"
                ++ "Loader Flags: 0x" ++ (showHex (loaderFlags hdr) "") ++ "\n"
                ++ "RVA: " ++ (show $ numberOfRVAandSizes hdr) ++ "\n"


data DirectoryEntry = DirEntry {
  virtualAddr :: Word32,
  entrySize :: Word32
} deriving Show

instance Binary DirectoryEntry where
  get = do
          addr <- getWord32le
          size <- getWord32le
          let entry = DirEntry {virtualAddr=addr, entrySize=size}
          return $ entry
  put _ = error "Serialization of DirectoryEntry not supported."

data SectionTable = SectionTable {
        sectionHeaderName :: String,
        virtualSize :: Word32,
        virtualAddress :: Word32,
        sizeOfRawData :: Word32,
        pointerToRawData :: Word32,
        pointerToRelocations :: Word32,
        pointerToLineNumbers :: Word32,
        numberOfRelocations :: Word16,
        numberOfLineNumbers :: Word16,
        secCharacteristics :: Word32
} deriving Show

instance Binary SectionTable where
  get = do
           sectionHeaderName' <- getWord64le
           virtualSize' <- getWord32le
           virtualAddress' <- getWord32le
           sizeOfRawData' <- getWord32le
           pointerToRawData' <- getWord32le
           pointerToRelocations' <- getWord32le
           pointerToLineNumbers' <- getWord32le
           numberOfRelocations' <- getWord16le
           numberOfLineNumbers' <- getWord16le
           secCharacteristics' <- getWord32le
           let header = SectionTable { sectionHeaderName=(byte64String sectionHeaderName'), virtualSize=virtualSize',
                                       virtualAddress=virtualAddress', sizeOfRawData=sizeOfRawData',
                                       pointerToRawData=pointerToRawData', pointerToRelocations=pointerToRelocations',
                                       pointerToLineNumbers=pointerToLineNumbers', numberOfRelocations=numberOfRelocations',
                                       numberOfLineNumbers=numberOfLineNumbers', secCharacteristics=secCharacteristics'}
           return header

  put _ = error "SectionTable serialization not supported"




data MachineType = UNKNOWN | AM33 | AMD64 | ARM | ARMV7 | EBC | I386 | IA64 | M32R | MIPS16 | MIPSFPU | MIPSFPU16 |
                    PPC | PPCFP | R4000 | SH3 | SH3DSP | SH4 | SH5 | THUMB | WCE | INVALID deriving (Show)


instance Binary MachineType where
  get = do
          x <- getWord16le
          return $ mapMachine x
  put _ = error "Serialization of MachineType not supported"



mapMachine :: Word16 -> MachineType
mapMachine w = case w of
                   0x00 -> UNKNOWN
                   0x1d3 -> AM33
                   0x8664 -> AMD64
                   0x1c0 -> ARM
                   0x1c4 -> ARMV7
                   0xebc -> EBC
                   0x14c -> I386
                   0x200 -> IA64
                   0x9041 -> M32R
                   0x266 -> MIPS16
                   0x366 -> MIPSFPU
                   0x466 -> MIPSFPU16
                   0x1f0 -> PPC
                   0x1f1 -> PPCFP
                   0x166 -> R4000
                   0x1a2 -> SH3
                   0x1a3 -> SH3DSP
                   0x1a6 -> SH4
                   0x1a8 -> SH5
                   0x1c2 -> THUMB
                   0x169 -> WCE
                   _ -> error "Bad machine type."

getAStr :: Get String
getAStr = do
            x <- getWord8
            case (x == 0x0) of
                  True  -> return []
                  False -> getAStr >>= \xs -> return ((chr $ fromIntegral x) : xs)