1    :  // Copyright 2011 Google Inc.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  
  15    :  #ifndef SYZYGY_PE_PE_FILE_WRITER_H_
  16    :  #define SYZYGY_PE_PE_FILE_WRITER_H_
  17    :  
  18    :  #include "base/file_path.h"
  19    :  #include "syzygy/block_graph/block_graph.h"
  20    :  #include "syzygy/core/address_space.h"
  21    :  #include "syzygy/pe/decomposer.h"
  22    :  #include "syzygy/pe/image_layout.h"
  23    :  #include "syzygy/pe/pe_file_parser.h"
  24    :  
  25  m :  namespace pe {
  26    :  
  27    :  // Given an address space and header information, writes a BlockGraph out
  28    :  // to a PE image file.
  29  m :  class PEFileWriter {
  30  m :   public:
  31  m :    typedef block_graph::BlockGraph BlockGraph;
  32  m :    typedef core::AbsoluteAddress AbsoluteAddress;
  33  m :    typedef core::FileOffsetAddress FileOffsetAddress;
  34  m :    typedef core::RelativeAddress RelativeAddress;
  35    :  
  36    :    // @param image_layout the image layout to write.
  37  m :    explicit PEFileWriter(const ImageLayout& image_layout);
  38    :  
  39    :    // Writes the image to path.
  40  m :    bool WriteImage(const FilePath& path);
  41    :  
  42    :    // Updates the checksum for the image @p path.
  43  m :    static bool PEFileWriter::UpdateFileChecksum(const FilePath& path);
  44    :  
  45  m :   protected:
  46    :    // Validates the DOS header and the NT headers in the image.
  47    :    // On success, sets the nt_headers_ pointer.
  48  m :    bool ValidateHeaders();
  49  m :    bool InitializeSectionFileAddressSpace();
  50  m :    bool WriteBlocks(FILE* file);
  51  m :    bool WriteOneBlock(AbsoluteAddress image_base,
  52  m :                       const BlockGraph::Block* block,
  53  m :                       FILE* file);
  54    :  
  55    :    // Maps from the relative offset to the start of a section to
  56    :    // the file offset for the start of that same section.
  57  m :    typedef core::AddressSpace<RelativeAddress, size_t, FileOffsetAddress>
  58  m :        SectionFileAddressSpace;
  59  m :    SectionFileAddressSpace section_file_offsets_;
  60    :  
  61    :    // Maps from section virtual address range to section index.
  62  m :    typedef core::AddressSpace<RelativeAddress, size_t, size_t>
  63  m :        SectionAddressSpace;
  64  m :    SectionAddressSpace sections_;
  65    :  
  66    :    // Our image layout as provided to the constructor.
  67  m :    const ImageLayout& image_layout_;
  68    :  
  69    :    // Refers to the nt headers from the image during WriteImage.
  70  m :    const IMAGE_NT_HEADERS* nt_headers_;
  71    :  
  72  m :   private:
  73  m :    DISALLOW_COPY_AND_ASSIGN(PEFileWriter);
  74  m :  };
  75    :  
  76  m :  }  // namespace pe
  77    :  
  78    :  #endif  // SYZYGY_PE_PE_FILE_WRITER_H_