Overwrite fix

.gitignore

@@ -1,4 +1,5 @@
 cmake-build-debug/
 build/
 .vs
-.idea
+.idea
+*.exe

.travis.yml

@@ -8,6 +8,7 @@ addons:
     packages:
       - gcc-6
       - g++-6
+      - libboost-all-dev
 notifications:
  email:
   on_success: never

CMakeLists.txt

@@ -1,7 +1,12 @@
-cmake_minimum_required(VERSION 3.6)
-project(pe_ep_intercept)
-
-set(KEYSTONE_BUILD_STATIC_RUNTIME on)
-include_directories(include include/keystone/include)
-add_subdirectory(include/keystone)
-add_subdirectory(src)
+cmake_minimum_required(VERSION 3.6)
+project(pe_ep_intercept)
+
+# Keystone
+set(KEYSTONE_BUILD_STATIC_RUNTIME on)
+include_directories(include include/keystone/include)
+
+file(GLOB COMMON_SOURCE_FILES common/*)
+
+# Sources
+add_subdirectory(include/keystone)
+add_subdirectory(src)

src/PePatch.cpp → common/Assembler.cpp

@@ -1,12 +1,8 @@
 #include <keystone/include/keystone/keystone.h>
-#include "PePatch.hpp"
+#include "Assembler.hpp"
 
-namespace PeEpIntercept {
-    PePatch::PePatch(std::string &path) : PeFile(path) {
-
-    }
-
-    std::vector<char> PePatch::Assemble(const std::string &assembly) {
+namespace Interceptor {
+    std::vector<char> Assembler::assemble(Architecture type, const std::string &assembly) {
         std::vector<char> instructions;
 
         if (assembly.empty()) {
@@ -29,35 +25,34 @@ namespace PeEpIntercept {
         ks_mode instruct_mode;
 
         switch (type) {
-            case PeArch::x86:
+            case Architecture::x86:
                 instruct_mode = KS_MODE_32;
                 break;
-            case PeArch::x64:
+            case Architecture::x64:
                 instruct_mode = KS_MODE_64;
                 break;
             default:
-                throw std::runtime_error("executable type not supported");
+                throw std::runtime_error("Executable type not supported");
         }
 
         if (ks_open(KS_ARCH_X86, instruct_mode, &ks) != KS_ERR_OK) {
-            throw std::runtime_error("failed to open keystone");
+            throw std::runtime_error("Failed to open keystone");
         }
 
-        std::unique_ptr<ks_engine[],
-                decltype(ks_deleter)> ks_ptr(ks, ks_deleter);
+        std::unique_ptr<ks_engine[], decltype(ks_deleter)> ks_ptr(ks, ks_deleter);
 
         if (ks_asm(ks, assembly.c_str(), 0, &encode, &size, &count) != KS_ERR_OK) {
-            throw std::runtime_error("failed to assemble instructions");
+            throw std::runtime_error("Failed to assemble instructions");
         }
 
         std::unique_ptr<unsigned char[],
                 decltype(code_deleter)> encode_ptr(encode, code_deleter);
 
-        if (size > 0xffffffff) {
-            throw std::runtime_error("exceeded max section size");
+        if (size > sizeof(uint32_t) * 0xff) {
+            throw std::runtime_error("Exceeded max section size");
         }
 
-        for (size_t i = 0; i < size; i++) {
+        for (auto i = 0; i < size; i++) {
             auto encoded = static_cast<char>(encode[i]);
             instructions.push_back(encoded);
         }

common/Assembler.hpp

@@ -0,0 +1,15 @@
+#ifndef PE_EP_INTERCEPT_ASSEMBLER_HPP
+#define PE_EP_INTERCEPT_ASSEMBLER_HPP
+
+#include "PeFile.hpp"
+
+namespace Interceptor {
+    class Assembler {
+    protected:
+    public:
+        std::vector<char> assemble(Architecture type, const std::string &assembly);
+    };
+}
+
+
+#endif //PE_EP_INTERCEPT_ASSEMBLER_HPP

src/PeAssembly.cpp → common/Assembly.cpp

@@ -1,8 +1,8 @@
-#include "PeAssembly.hpp"
+#include "Assembly.hpp"
 
-namespace PeEpIntercept {
-    std::string EntryRedirectAssemblyX64(uint32_t oep) {
-        std::string address = std::to_string(oep);
+namespace Interceptor {
+    std::string entryRedirectAssemblyX64(uint32_t oep) {
+        auto address = std::to_string(oep);
 
         return "push rbp;"
                        "mov rbp, rsp;"
@@ -44,8 +44,8 @@ namespace PeEpIntercept {
                        "ret;";
     }
 
-    std::string EntryRedirectAssemblyX86(uint32_t oep) {
-        std::string address = std::to_string(oep);
+    std::string entryRedirectAssemblyX86(uint32_t oep) {
+        auto address = std::to_string(oep);
 
         return "push ebp;"
                        "mov ebp, esp;"
@@ -86,4 +86,4 @@ namespace PeEpIntercept {
                        // patch in ExitProcess
                        "ret;";
     }
-} // namespace PeEpIntercept
+} // namespace Interceptor

common/Assembly.hpp

@@ -0,0 +1,19 @@
+#ifndef PE_EP_INTERCEPT_ASSEMBLY_HPP
+#define PE_EP_INTERCEPT_ASSEMBLY_HPP
+
+#include <string>
+
+namespace Interceptor {
+    enum class Architecture {
+        x86,
+        x64,
+        unknown
+    };
+
+    std::string entryRedirectAssemblyX64(uint32_t oep);
+
+    std::string entryRedirectAssemblyX86(uint32_t oep);
+}
+
+
+#endif //PE_EP_INTERCEPT_ASSEMBLY_HPP

common/Editor.cpp

@@ -0,0 +1,43 @@
+#include <iostream>
+#include "Editor.hpp"
+#include "Assembler.hpp"
+#include "PePatchX64.hpp"
+
+namespace Interceptor {
+    PeFile Editor::original(std::fstream &file) {
+        try {
+            return PeFile(file);
+        } catch (std::runtime_error &ex) {
+            auto message = std::string("Could not parse the file. Reason: \"") + ex.what() + "\".";
+            throw std::runtime_error(message);
+        }
+    }
+
+    std::tuple<PeFile, bool> Editor::edit(
+            std::fstream &file,
+            const std::string &section
+    ) {
+        std::unique_ptr<Interceptor::PePatchX64> patcher;
+        auto exe_file = original(file);
+
+        if (exe_file.hasSection(section)) {
+            std::cout << "The executable already has the section \"" << section << "\"." << std::endl;
+            return std::make_tuple(exe_file, false);
+        }
+
+        Assembler assembler;
+        auto arch = exe_file.getPeArch();
+
+        switch (arch) {
+            case Architecture::x64:
+                patcher = std::make_unique<Interceptor::PePatchX64>(assembler, exe_file);
+                break;
+            default:
+                std::cout << "The executable has an unsupported architecture." << std::endl;
+                return std::make_tuple(exe_file, false);
+        }
+
+        auto patched_file = patcher->patch(section);
+        return std::make_tuple(patched_file, true);
+    }
+}

common/Editor.hpp

@@ -0,0 +1,25 @@
+#ifndef PE_EP_INTERCEPT_EDITOR_HPP
+#define PE_EP_INTERCEPT_EDITOR_HPP
+
+#include "PeFile.hpp"
+
+namespace Interceptor {
+    class Editor {
+    private:
+        static PeFile original(std::fstream &file);
+    public:
+        /**
+         *
+         * @param file The stream contents will be used to instantiate a PeFile object
+         * @param section The new section will have this string as its name
+         * @return Contains a success. Returns original file on failure, patched on success
+         */
+        static std::tuple<PeFile, bool> edit(
+                std::fstream &file,
+                const std::string &section
+        );
+    };
+}
+
+
+#endif //PE_EP_INTERCEPT_EDITOR_HPP