PEFile.hs

module Gadgets.PEFile where

import Data.PE.Parser
import Data.PE.Structures
import Data.PE.Utils 
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.Maybe
import System.Environment
import Data.Word
import System.IO.Unsafe
import Data.Binary
import Data.Binary.Get
import Data.Char
import Data.Bits
import Data.Array.Unboxed
import Data.List

type Filename = String
type Secname = String
type SectionMeta = (SectionTable, LBS.ByteString)

getsecandinfo :: Filename -> Secname -> IO ((Maybe SectionMeta, MachineType))
getsecandinfo fn sn = buildFile fn >>= \pefile -> return (getsection pefile sn, getmachinetype pefile)

getsec :: Filename -> Secname -> IO (Maybe SectionMeta)
getsec fn sn = buildFile fn >>= \pefile -> return $ getsection pefile sn

getsecs :: Filename -> [SectionMeta]
getsecs fn = unsafePerformIO (buildFile fn >>= \pefile -> return $ (sectionTables.peHeader) pefile)

getary fn = arrayrep $ getsecs fn

getdirs :: Filename -> [DirectoryEntry]
getdirs fn = unsafePerformIO (buildFile fn >>= \pefile -> return $ (dataDirectories.peHeader) pefile)

getsection :: PEFile -> Secname -> Maybe SectionMeta
getsection pefile secn = let sections = (sectionTables.peHeader) pefile in
                          find (\x -> secn == (sectionHeaderName $ fst x)) sections 

getmachinetype :: PEFile -> MachineType
getmachinetype pe = targetMachine $ coffHeader $ peHeader pe

showsections :: Filename -> IO ()
showsections filename = do
                           pefile <- buildFile filename
                           let sections = (sectionTables.peHeader) pefile
                           let coff = (coffHeader.peHeader) pefile
                           let std = (standardFields.peHeader) pefile
                           let showme = \x -> (sectionHeaderName $ fst x)
                           --putStr $ show datadirs
                           putStr $ show $ coff
                           putStr $ show $ std
                           putStr $ show $ map showme sections 
                           --putStr $ show $ (numberOfRVAandSizes.windowsSpecFields.peHeader) pefile
                           --putStr $ show pefile
                           return () 


--Import Table Parsing stuff.  This should eventually move to the PE library.

type ImportDirectory = [ImportDirectoryEntry]
type ImportLookupTable = [ImportLookupTableEntry]

data ImportDirectoryEntry = ID {
                                lookupTableRVA :: Word32,
                                timeStamp :: Word32,
                                forwarderChain :: Word32,
                                nameRVA :: Word32,
                                importAddressTableRVA :: Word32
                               } | IDNull deriving (Show,Eq)

data HintNameEntry = HNE { 
                           hint :: Word16,
                           name :: String
                         } deriving (Show, Eq)

data ImportLookupTableEntry = ILTOrd Word16 | ILTHint Word32 | ILTNull deriving (Show,Eq)


getImpDir :: Get ImportDirectory
getImpDir = do
              entry <- get
              case (entry) of
                   IDNull -> return [IDNull]
                   x      -> getImpDir >>= \y -> return (x : y)


getLT :: Get ImportLookupTable
getLT = do
           entry <- get
           case (entry) of
                ILTNull -> return [ILTNull]
                x       -> getLT >>= \y -> return (x : y)


instance Binary HintNameEntry where
      put (HNE h n) = let words = (map fromIntegral $ map ord n)::[Word8] in
                       do 
                        put h
                        put words
                        if (length words `mod` 2 == 0) 
                          then put (0x0::Word8)
                          else return ()
      get = do
               ordinal <- getWord16le
               astr <- getAStr
               if (length astr `mod` 2 == 0)
                 then getWord8 >>= \_ -> return (HNE ordinal astr)
                 else return (HNE ordinal astr)

instance Binary ImportDirectoryEntry where
    put (ID lut ts fc nrva iarva) = put lut >> put ts >> put fc >> put nrva >> put iarva
    put (IDNull) = put (0x0::Word32) >> put (0x0::Word32) >> put (0x0::Word32) >> put (0x0::Word32) >> put (0x0::Word32)
    get = do
            lut <- getWord32le
            ts <- getWord32le
            fc <- getWord32le
            nrva <- getWord32le
            iarva <- getWord32le
            case (lut + ts + fc + nrva + iarva) of
                    0 -> return IDNull
                    _ -> return (ID lut ts fc nrva iarva)

instance Binary ImportLookupTableEntry where
   put (ILTOrd ord ) = put (0x80::Word8) >> put ord >> put (0x00::Word8)
   put (ILTHint rva) = put (setBit rva 31)
   put ILTNull = put (0x0::Word32)
   get = do
            word <- getWord32le
            case (word) of
                0 -> return ILTNull
                _ -> case (testBit word 31) of
                          True  -> return $ ILTOrd $ fromIntegral word
                          False -> return $ ILTHint (clearBit word 31)
--More PE Data structure stuff

importInfo fn = importInfo' (getsecs fn) (getdirs fn)

importInfo' secns dirs = map infos ientries 
                where ary = arrayrep secns
                      ientries = delete IDNull $  buildImport ary dirs
                      lookups = (buildLookup ary) 
                      hnts = (buildHintName ary) 
                      infos = \x -> (getdllname ary x, map name $ map hnts $ delete ILTNull $ lookups x)

--Build the Import table.
buildImport ary dirs = runGet getImpDir bstr  
                where itaddr = virtualAddr (dirs !! 1)
                      bstr = grabAt (fromIntegral itaddr) ary
                      
buildLookup ary ientry = runGet getLT (grabAt (fromIntegral rva) ary)
                where rva = lookupTableRVA ientry

buildHintName ary ltentry = case (ltentry) of
                                  (ILTHint x) -> runGet hnte (grabAt (fromIntegral x) ary)
                                  (ILTNull) -> error "Null encountered"
                                  _ -> error "Not working with ords today"
                where hnte = get >>= \x -> return x::Get HintNameEntry

getdllname ary ientry = case (ientry) of
                             (IDNull) -> ""
                             _        -> runGet getAStr (grabAt (fromIntegral rva) ary) 
                where rva = nameRVA ientry
--Building an array to represent the file structure
sectoblist (secn, bytes) = let words = LBS.unpack bytes in
                           let indxs x = x : indxs (x+1) in
                            zip (indxs $ fromIntegral $ virtualAddress secn) words

arrayrep :: [SectionMeta] -> UArray Word32 Word8
arrayrep secn = array (0,maxaddr) words 
        where
          words = concat $ map sectoblist secn
          maxaddr = maximum $ map fst words

--Ask for an address to begin a new head for a bytestring to build from, simple enough.
{-
grabAt :: Word32 -> UArray Word32 Word8 -> LBS.ByteString
grabAt indx ary = LBS.pack $ elems newarray 
        where maxdx = maximum $ indices ary
              newarray = ixmap (0,maxdx-indx) (\i -> i - indx) ary --remap the array
-}
grabAt :: Int -> UArray Word32 Word8 -> LBS.ByteString
grabAt indx ary = LBS.pack $ drop (indx) $ elems ary