PeFile.h

/*
* PeFile.h - Part of the PeLib library.
*
* Copyright (c) 2004 Sebastian Porst (webmaster@the-interweb.com)
* All rights reserved.
*
* This software is licensed under the zlib/libpng License.
* For more details see http://www.opensource.org/licenses/zlib-license.php
* or the license information file (license.htm) in the root directory 
* of PeLib.
*/

#ifndef PEFILE_H
#define PEFILE_H

#include "PeLibInc.h"
#include "MzHeader.h"
#include "PeHeader.h"
#include "ImportDirectory.h"
#include "ExportDirectory.h"
#include "BoundImportDirectory.h"
#include "ResourceDirectory.h"
#include "RelocationsDirectory.h"
#include "ComHeaderDirectory.h"
#include "IatDirectory.h"
#include "DebugDirectory.h"
#include "TlsDirectory.h"

namespace PeLib
{
        class PeFile32;
        class PeFile64;
        
        class PeFileVisitor
        {
                public:
                  virtual void callback(PeFile32 &file){}
                  virtual void callback(PeFile64 &file){}
                  virtual ~PeFileVisitor(){}
        };

        template<int>
        struct PeFile_Traits;

        template<>
        struct PeFile_Traits<32>
        {
                typedef PeHeader32 PeHeader32_64;
        };
        
        template<>
        struct PeFile_Traits<64>
        {
                typedef PeHeader64 PeHeader32_64;
        };
        
        class PeFile
        {
                protected:
                  std::string m_filename; 
                  MzHeader m_mzh; 
                  ExportDirectory m_expdir; 
                  BoundImportDirectory m_boundimpdir; 
                  ResourceDirectory m_resdir; 
                  RelocationsDirectory m_relocs; 
                  ComHeaderDirectory m_comdesc; 
                  IatDirectory m_iat; 
                  DebugDirectory m_debugdir;
                public:
                  virtual ~PeFile();
                  
                  virtual std::string getFileName() const = 0;
                  virtual void setFileName(std::string strFilename) = 0;
                  
                  virtual void visit(PeFileVisitor &v) = 0;
                  
                  virtual int readMzHeader() = 0;
                  virtual int readExportDirectory()  = 0;
                  virtual int readPeHeader()  = 0;
                  virtual int readImportDirectory()  = 0;
                  virtual int readBoundImportDirectory()  = 0;
                  virtual int readResourceDirectory()  = 0;
                  virtual int readRelocationsDirectory()  = 0;
                  virtual int readComHeaderDirectory()  = 0;
                  virtual int readIatDirectory()  = 0;
                  virtual int readDebugDirectory()  = 0;
                  virtual int readTlsDirectory()  = 0;
                  
                  virtual unsigned int getBits() const = 0;

                  const MzHeader& mzHeader() const;
                  MzHeader& mzHeader();

                  const ExportDirectory& expDir() const;
                  ExportDirectory& expDir();
                  
                  const BoundImportDirectory& boundImpDir() const;
                  BoundImportDirectory& boundImpDir();
                  
                  const ResourceDirectory& resDir() const;
                  ResourceDirectory& resDir();

                  const RelocationsDirectory& relocDir() const;
                  RelocationsDirectory& relocDir();

                  const ComHeaderDirectory& comDir() const;
                  ComHeaderDirectory& comDir();

                  const IatDirectory& iatDir() const;
                  IatDirectory& iatDir();
                  
                  const DebugDirectory& debugDir() const;
                  DebugDirectory& debugDir();
                  
        };
        
        template<int bits>
        class PeFileT : public PeFile
        {
                typedef typename PeFile_Traits<bits>::PeHeader32_64 PeHeader32_64;
                
                private:
                  PeHeader32_64 m_peh; 
                  ImportDirectory<bits> m_impdir; 
                  TlsDirectory<bits> m_tlsdir;
                
                public:
                  PeFileT();
                  
                  virtual ~PeFileT() {}
                  
                  explicit PeFileT(const std::string& strFilename);
                  
                  std::string getFileName() const;
                  void setFileName(std::string strFilename);

                  int readMzHeader() ;
                  int readExportDirectory() ;
                  int readPeHeader() ;
                  int readImportDirectory() ;
                  int readBoundImportDirectory() ;
                  int readResourceDirectory() ;
                  int readRelocationsDirectory() ;
                  int readComHeaderDirectory() ;
                  int readIatDirectory() ;
                  int readDebugDirectory() ;
                  int readTlsDirectory() ;
                  
                  unsigned int getBits() const
                  {
                          return bits;
                  }
                  
                  const PeHeader32_64& peHeader() const;
                  PeHeader32_64& peHeader();

                  const ImportDirectory<bits>& impDir() const;
                  ImportDirectory<bits>& impDir();
                  
                  const TlsDirectory<bits>& tlsDir() const;
                  TlsDirectory<bits>& tlsDir();
        };

        class PeFile32 : public PeFileT<32>
        {
                public:
                  PeFile32();

                  explicit PeFile32(const std::string& strFlename);
                  virtual void visit(PeFileVisitor &v) { v.callback( *this ); }
        };
        
        class PeFile64 : public PeFileT<64>
        {
                public:
                  PeFile64();

                  explicit PeFile64(const std::string& strFlename);
                  virtual void visit(PeFileVisitor &v) { v.callback( *this ); }
        };
        
        //typedef PeFileT<32> PeFile32;
        //typedef PeFileT<64> PeFile64;
        
        template<int bits>
        PeFileT<bits>::PeFileT(const std::string& strFilename)
        {
                m_filename = strFilename;
        }

        template<int bits>
        PeFileT<bits>::PeFileT()
        {
        }

        template<int bits>
        int PeFileT<bits>::readPeHeader() 
        {
                return peHeader().read(getFileName(), mzHeader().getAddressOfPeHeader());
        }

        
        template<int bits>
        int PeFileT<bits>::readImportDirectory() 
        {
                if (peHeader().getIddImportRva() && peHeader().getIddImportSize())
                {
                        return impDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddImportRva())), peHeader().getIddImportSize(), peHeader());
                }
                return 1;
        }

        template<int bits>
        const typename PeFile_Traits<bits>::PeHeader32_64& PeFileT<bits>::peHeader() const
        {
                return m_peh;
        }

        template<int bits>
        typename PeFile_Traits<bits>::PeHeader32_64& PeFileT<bits>::peHeader()
        {
                return m_peh;
        }

        template<int bits>
        const ImportDirectory<bits>& PeFileT<bits>::impDir() const
        {
                return m_impdir;
        }

        template<int bits>
        ImportDirectory<bits>& PeFileT<bits>::impDir()
        {
                return m_impdir;
        }

        template<int bits>
        const TlsDirectory<bits>& PeFileT<bits>::tlsDir() const
        {
                return m_tlsdir;
        }

        template<int bits>
        TlsDirectory<bits>& PeFileT<bits>::tlsDir()
        {
                return m_tlsdir;
        }

        template<int bits>
        std::string PeFileT<bits>::getFileName() const
        {
                return m_filename;
        }

        template<int bits>
        void PeFileT<bits>::setFileName(std::string strFilename)
        {
                m_filename = strFilename;
        }

        template<int bits>
        int PeFileT<bits>::readMzHeader() 
        {
                return mzHeader().read(getFileName());
        }
        
        template<int bits>
        int PeFileT<bits>::readExportDirectory() 
        {
                if (peHeader().getIddExportRva() && peHeader().getIddExportSize())
                {
                        return expDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddExportRva())), peHeader().getIddExportSize(), peHeader());
                }
                return 1;
        }

        
        template<int bits>
        int PeFileT<bits>::readBoundImportDirectory() 
        {
                if (peHeader().getIddBoundImportRva() && peHeader().getIddBoundImportSize())
                {
                        return boundImpDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddBoundImportRva())), peHeader().getIddBoundImportSize());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readResourceDirectory() 
        {
                if (peHeader().getIddResourceRva() && peHeader().getIddResourceSize())
                {
                        return resDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddResourceRva())), peHeader().getIddResourceSize(), peHeader().getIddResourceRva());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readRelocationsDirectory() 
        {
                if (peHeader().getIddBaseRelocRva() && peHeader().getIddBaseRelocSize())
                {
                        return relocDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddBaseRelocRva())), peHeader().getIddBaseRelocSize());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readComHeaderDirectory() 
        {
                if (peHeader().getIddComHeaderRva() && peHeader().getIddComHeaderSize())
                {
                        return comDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddComHeaderRva())), peHeader().getIddComHeaderSize());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readIatDirectory() 
        {
                if (peHeader().getIddIatRva() && peHeader().getIddIatSize())
                {
                        return iatDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddIatRva())), peHeader().getIddIatSize());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readDebugDirectory() 
        {
                if (peHeader().getIddDebugRva() && peHeader().getIddDebugSize())
                {
                        return debugDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddDebugRva())), peHeader().getIddDebugSize());
                }
                return 1;
        }

        template<int bits>
        int PeFileT<bits>::readTlsDirectory() 
        {
                if (peHeader().getIddTlsRva() && peHeader().getIddTlsSize())
                {
                        return tlsDir().read(getFileName(), static_cast<unsigned int>(peHeader().rvaToOffset(peHeader().getIddTlsRva())), peHeader().getIddTlsSize());
                }
                return 1;
        }

}

#endif

---