utils/zenutils/source/shared/pe.h

   1 /* zenutils - Utilities for working with creative firmwares.
   2  * Copyright 2007 (c) Rasmus Ry <rasmus.ry{at}gmail.com>
   3  *
   4  * This program is free software; you can redistribute it and/or modify
   5  * it under the terms of the GNU General Public License as published by
   6  * the Free Software Foundation; either version 2 of the License, or
   7  * (at your option) any later version.
   8  *
   9  * This program is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12  * GNU General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License
  15  * along with this program; if not, write to the Free Software
  16  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  17  */
  18
  19 #ifndef SHARED_PE_H_INCLUDED
  20 #define SHARED_PE_H_INCLUDED
  21
  22 #include <string>
  23 #include <pelib/PeLib.h>
  24 #include <utils.h>
  25
  26 namespace shared {
  27     struct section_info
  28     {
  29         word index;
  30         dword virtual_address;
  31         dword virtual_size;
  32         dword raw_address;
  33         dword raw_size;
  34         dword characteristics;
  35     }; //struct section_info
  36
  37     class pe_file
  38     {
  39     public:
  40         pe_file(PeLib::PeFile* pef = NULL);
  41         ~pe_file();
  42
  43         bool is_valid() const;
  44         bool read(const std::string& filename);
  45         bool find_section(const std::string& name, section_info& info) const;
  46         bool add_section(const std::string& name, const bytes& buffer, section_info& info);
  47         dword get_image_base() const;
  48         dword pa_to_va(PeLib::dword pa) const;
  49
  50     protected:
  51         template <int _Bits>
  52         static bool find_section(const PeLib::PeFileT<_Bits>* pef,
  53                           const std::string& name, section_info& info);
  54         template <int _Bits>
  55         static bool add_section(PeLib::PeFileT<_Bits>* pef,
  56                           const std::string& name, const bytes& buffer,
  57                           section_info& info);
  58
  59     private:
  60         PeLib::PeFile* _pef;
  61     }; //class pe_file
  62
  63
  64     template <int _Bits>
  65     bool pe_file::find_section(const PeLib::PeFileT<_Bits>* pef,
  66                       const std::string& name, section_info& info)
  67     {
  68         for (PeLib::word i = 0; i < pef->peHeader().getNumberOfSections(); i++)
  69         {
  70             if (pef->peHeader().getSectionName(i) == name)
  71             {
  72                 info.index = i;
  73                 info.virtual_address = pef->peHeader().getVirtualAddress(i);
  74                 info.virtual_size = pef->peHeader().getVirtualSize(i);
  75                 info.raw_address = pef->peHeader().getPointerToRawData(i);
  76                 info.raw_size = pef->peHeader().getSizeOfRawData(i);
  77                 info.characteristics = pef->peHeader().getCharacteristics(i);
  78                 return true;
  79             }
  80         }
  81         return false;
  82     }
  83
  84     template <int _Bits>
  85     bool pe_file::add_section(PeLib::PeFileT<_Bits>* pef,
  86                       const std::string& name, const bytes& buffer,
  87                       section_info& info)
  88     {
  89         using namespace PeLib;
  90
  91         // Check if the last section has the same name as the one being added.
  92         PeLib::word secnum = pef->peHeader().getNumberOfSections();
  93         if (pef->peHeader().getSectionName(secnum-1) == name)
  94         {
  95             // If it is, we change the attributes of the existing section.
  96             secnum = secnum - 1;
  97             pef->peHeader().setSizeOfRawData(secnum,
  98                 alignOffset(buffer.size(),
  99                             pef->peHeader().getFileAlignment()));
 100             pef->peHeader().setVirtualSize(secnum,
 101                 alignOffset(buffer.size(),
 102                             pef->peHeader().getSectionAlignment()));
 103             PeLib::dword chars = pef->peHeader().getCharacteristics(secnum-1);
 104             pef->peHeader().setCharacteristics(secnum,
 105                 chars | PELIB_IMAGE_SCN_MEM_WRITE | PELIB_IMAGE_SCN_MEM_READ);
 106         }
 107         else
 108         {
 109             // Otherwise we add a new section.
 110             if (pef->peHeader().addSection(name, buffer.size()) != NO_ERROR)
 111             {
 112                 return false;
 113             }
 114             pef->peHeader().makeValid(pef->mzHeader().getAddressOfPeHeader());
 115             pef->peHeader().write(pef->getFileName(), pef->mzHeader().getAddressOfPeHeader());
 116         }
 117
 118         // Save the section headers to the file.
 119         if (pef->peHeader().writeSections(pef->getFileName()) != NO_ERROR)
 120         {
 121             return false;
 122         }
 123
 124         // Save the section data to the file.
 125         if (pef->peHeader().writeSectionData(pef->getFileName(), secnum, buffer) != NO_ERROR)
 126         {
 127             return false;
 128         }
 129
 130         // Fill out the section information.
 131         info.index = secnum;
 132         info.virtual_address = pef->peHeader().getVirtualAddress(secnum);
 133         info.virtual_size = pef->peHeader().getVirtualSize(secnum);
 134         info.raw_address = pef->peHeader().getPointerToRawData(secnum);
 135         info.raw_size = pef->peHeader().getSizeOfRawData(secnum);
 136         info.characteristics = pef->peHeader().getCharacteristics(secnum);
 137
 138         return true;
 139     }
 140 }; //namespace shared
 141
 142 #endif //SHARED_PE_H_INCLUDED