#include <stdlib.h>
#include <fstream>
#include <iostream>
#include <string>
#include <sstream>
#include <iomanip>
#include <stdio.h>
#include <vector>
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include <pe_lib/pe_bliss.h>


//TODO clean all this crap up.
//using pointer tmp_image to image and -> to access member functions

class PE{
  private:

    struct PeSection{
      std::string name;
      int offset;
      int size;
    };

    std::vector<PeSection> PE_Sections;
    std::ifstream pe_file;
    pe_bliss::pe_base *image; //<< pointer to an pe_base object

  public:
    PE(std::string file_path);
    ~PE();
    int get_file_offset_entry_point();
    void get_basic_pe_info();


};


PE::PE(std::string file_path){
  using namespace pe_bliss;

  pe_file.open(file_path.c_str(), std::ifstream::in | std::ifstream::binary);
  if(!pe_file){
    std::cout << "PE Cannot open " << file_path << std::endl;
  }

  try{
    pe_bliss::pe_base new_image(pe_factory::create_pe(pe_file)); //creation of object
    image = &new_image; //pointer to object is set
  }
  catch(const pe_exception& e){
    std::cout << "Error: " << e.what() << std::endl;
  }
}

PE::~PE(){
  //dtor
}

int PE::get_file_offset_entry_point(){
  int FileOffSet;
  int AddressOfEntryPoint = image->get_ep(); //segfault here :(
  const pe_bliss::section_list sections = image->get_image_sections();

  for(pe_bliss::section_list::const_iterator it = sections.begin(); it != sections.end(); ++it){
    const pe_bliss::section& s = *it;
    if((AddressOfEntryPoint >= s.get_virtual_address()) &&
       AddressOfEntryPoint <= s.get_virtual_address() + s.get_size_of_raw_data()){
      int FileOffSet = image->rva_to_file_offset(s.get_virtual_address()) + (AddressOfEntryPoint - s.get_virtual_address());
      std::cout << "OEP FileOffSet: " << std::hex << (FileOffSet) << " in section: " << s.get_name() << std::endl;
      return(FileOffSet);
    }
  }
  return FileOffSet;
}

using namespace pe_bliss;

void PE::get_basic_pe_info(){

  std::cout << "PE file type: " << (image->get_pe_type() == pe_type_32 ? "PE32 (PE)" : "PE64 (PE+)") << std::endl;
}

/*

    try
    {

        pe_base image(pe_factory::create_pe(pe_file));
        std::cout << "PE file type: " << (image.get_pe_type() == pe_type_32 ? "PE32 (PE)" : "PE64 (PE+)") << std::endl;

        std::cout << "Calculated checksum: "<< std::hex << std::showbase << calculate_checksum(pe_file) << std::endl;
        std::cout << "Stored checksum: " << image.get_checksum() << std::endl;

        std::cout << "Characteristics: " << image.get_characteristics() << std::endl;

        std::cout << "Entry point: " << image.get_ep() << std::endl;

        std::cout << "File alignment: " << image.get_file_alignment() << std::endl;
        std::cout << "Section alignment: " << image.get_section_alignment() << std::endl;

        std::cout << "Image base: " << image.get_image_base_64() << std::endl;

        std::cout << "Subsystem: " << image.get_subsystem() << std::endl;
        std::cout << "Is console: " << (image.is_console() ? "YES" : "NO") << std::endl;
        std::cout << "Is windows GUI: " << (image.is_gui() ? "YES" : "NO") << std::endl;

        std::cout << "Has bound import: " << (image.has_bound_import() ? "YES" : "NO") << std::endl;
        std::cout << "Has config: " << (image.has_config() ? "YES" : "NO") << std::endl;
        std::cout << "Has debug: " << (image.has_debug() ? "YES" : "NO") << std::endl;
        std::cout << "Has delay import: " << (image.has_delay_import() ? "YES" : "NO") << std::endl;
        std::cout << "Has exception directory: " << (image.has_exception_directory() ? "YES" : "NO") << std::endl;
        std::cout << "Has exports: " << (image.has_exports() ? "YES" : "NO") << std::endl;
        std::cout << "Has imports: " << (image.has_imports() ? "YES" : "NO") << std::endl;
        std::cout << "Has reloc: " << (image.has_reloc() ? "YES" : "NO") << std::endl;
        std::cout << "Has resources: " << (image.has_resources() ? "YES" : "NO") << std::endl;
        std::cout << "Has security: " << (image.has_security() ? "YES" : "NO") << std::endl;
        std::cout << "Has tls: " << (image.has_tls() ? "YES" : "NO") << std::endl;
        std::cout << "Is .NET: " << (image.is_dotnet() ? "YES" : "NO") << std::endl;


        int AddressOfEntryPoint = image.get_ep();
        const section_list sections = image.get_image_sections();
        PE_Sections.clear();

    for(section_list::const_iterator it = sections.begin(); it != sections.end(); ++it){
            const section& s = *it; //Секция
            PE_Section tmp_Section;
            tmp_Section.name = s.get_name();
            tmp_Section.offset = s.get_pointer_to_raw_data();
            tmp_Section.size = s.get_size_of_raw_data();
            PE_Sections.push_back(tmp_Section);

            cout << "Section Name: " << s.get_name() << endl;
      cout << "Section Offset: " << tmp_Section.offset << endl;
      cout << "Section Size: " << tmp_Section.size << endl;

            if((AddressOfEntryPoint >= s.get_virtual_address()) && AddressOfEntryPoint <= s.get_virtual_address() + s.get_size_of_raw_data()){
        int FileOffSet = image.rva_to_file_offset(s.get_virtual_address()) + (AddressOfEntryPoint - s.get_virtual_address());
        cout << "OEP FileOffSet: " << std::hex << (FileOffSet) << " in section: " << s.get_name() << endl;
        return(FileOffSet);
        //cout << " Decimal: " << std::dec << FileOffSet << endl;
        //break;
            }
        }
    }
    catch(const pe_exception& e)
    {
        std::cout << "Error: " << e.what() << std::endl;
        return 0;
    }
}

*/