Finished pe image implementation wrapper.

Moved some files to more appropriate folders.
Added linuxpe submodule.
Added linuxpe as a precompiled header.

8 files changed, 165 insertions, 10 deletions

diff --git a/server/src/image/pe.h b/server/src/image/pe.h
new file mode 100644
index 0000000..2f81616
--- /dev/null
+++ b/server/src/image/pe.h
@@ -0,0 +1,129 @@
+#pragma once
+
+// Had to put implementation inside a header file because im using templated
+// class to make it work in its separate source file, I had to do some weird
+// fuckery but I didn't want to do it
+
+namespace pe {
+
+struct import_t {
+  std::string name;
+  uint32_t rva;
+};
+
+struct section_t {
+  std::string name;
+  size_t size;
+  uint32_t rva;
+  uint32_t va;
+};
+
+template <bool x64 = false>
+class image {
+  win::image_t<x64> *m_image;
+  std::vector<char> m_buffer;
+
+  std::unordered_map<std::string, std::vector<import_t>> m_imports;
+  std::vector<section_t> m_sections;
+  std::vector<std::pair<uint32_t, win::reloc_entry_t>> m_relocs;
+
+ public:
+  image(const std::string_view name) : m_image{nullptr} {
+    io::read_file(name, m_buffer);
+    if (m_buffer.empty()) {
+      io::logger->error("failed to read image {}", name);
+      return;
+    }
+
+    m_image = reinterpret_cast<win::image_t<x64> *>(m_buffer.data());
+    load();
+  }
+
+  void load() {
+    parse_sections();
+    parse_relocs();
+    parse_imports();
+  }
+
+  void reload(const std::string_view name) {
+    io::read_file(name, m_buffer);
+    if (m_buffer.empty()) {
+      io::logger->error("failed to reload image {}", name);
+      return;
+    }
+
+    m_image = reinterpret_cast<win::image_t<x64> *>(m_buffer.data());
+    load();
+  }
+
+  void parse_sections() {
+    const auto nt = m_image->get_nt_headers();
+    const size_t n = nt->file_header.num_sections;
+
+    for (size_t i = 0; i < n; i++) {
+      auto section = nt->get_section(i);
+      m_sections.emplace_back(section_t{section->name, section->size_raw_data,
+                               section->ptr_raw_data,
+                               section->virtual_address});
+    }
+  };
+
+  void parse_relocs() {
+    const auto reloc_dir =
+        m_image->get_directory(win::directory_id::directory_entry_basereloc);
+    if (!reloc_dir) return;
+
+    const auto ptr = m_image->rva_to_ptr(reloc_dir->rva);
+    auto block = reinterpret_cast<win::reloc_block_t *>(ptr);
+
+    while (block->base_rva) {
+      for (size_t i = 0; i < block->num_entries(); ++i) {
+        auto entry = block->entries[i];
+
+        m_relocs.emplace_back(std::make_pair(block->base_rva, entry));
+      }
+      block = block->get_next();
+    }
+  }
+
+  void parse_imports() {
+    const auto import_dir =
+        m_image->get_directory(win::directory_id::directory_entry_import);
+    if (!import_dir) return;
+
+    const auto ptr = m_image->rva_to_ptr(import_dir->rva);
+    auto table = reinterpret_cast<win::import_directory_t *>(ptr);
+
+    for (uint32_t previous_name = 0; previous_name < table->rva_name;
+         previous_name = table->rva_name, ++table) {
+      auto name_ptr = m_image->rva_to_ptr(table->rva_name);
+      auto mod_name = std::string(reinterpret_cast<char *>(name_ptr));
+
+      auto thunk = reinterpret_cast<win::image_thunk_data_t<x64> *>(
+          m_image->rva_to_ptr(table->rva_original_first_thunk));
+
+      for (uint32_t index = 0; thunk->address;
+           index += sizeof(uint32_t), ++thunk) {
+        auto named_import = reinterpret_cast<win::image_named_import_t *>(
+            m_image->rva_to_ptr(thunk->address));
+
+        if (!thunk->is_ordinal) {
+          import_t data;
+          data.name = reinterpret_cast<const char *>(named_import->name);
+          data.rva = table->rva_first_thunk + index;
+
+          m_imports[mod_name].emplace_back(std::move(data));
+        }
+      }
+    }
+  }
+
+  const auto operator()() { return m_image; }
+  operator bool() const { return m_image != nullptr; }
+
+  auto imports() const { return m_imports; }
+  auto relocs() const { return m_relocs; }
+  auto sections() const { return m_sections; }
+};
+
+};  // namespace pe
\ No newline at end of file
diff --git a/server/src/include.h b/server/src/include.h
index 19e3972..9fe4e09 100644
--- a/server/src/include.h
+++ b/server/src/include.h
@@ -38,3 +38,7 @@
 #include <spdlog/sinks/stdout_color_sinks.h>
 
 #include <openssl/ssl.h>
+
+#include <linux-pe/linuxpe>
+
+#include "util/util.h"
\ No newline at end of file
diff --git a/server/src/main.cpp b/server/src/main.cpp
index b4cb3cb..7b96ff3 100644
--- a/server/src/main.cpp
+++ b/server/src/main.cpp
@@ -2,17 +2,27 @@
 #include "util/io.h"
 #include "util/commands.h"
 #include "server/server.h"
+#include "image/pe.h"
 
 constexpr std::string_view client_version{"0.1.0"};
 
 int main(int argc, char* argv[]) {
   io::init(false);
 
-  tcp::server server("6666");
+  pe::image test("out");
 
-  server.start();
+  for(auto&[mod, imports] : test.imports()) {
+    io::logger->info(mod);
+    for(auto &i : imports) {
+      io::logger->info("  {}->{:x}", i.name, i.rva);
+    }
+  }
+
+  tcp::server client_server("6666");
+
+  client_server.start();
 
-  server.connect_event.add([&](tcp::client& client) {
+  client_server.connect_event.add([&](tcp::client& client) {
     auto ip = client.get_ip();
     client.gen_session();
     client.write(tcp::packet_t(client_version,
@@ -21,19 +31,19 @@ int main(int argc, char* argv[]) {
     io::logger->info("{} connected", ip);
   });
 
-  server.disconnect_event.add([&](tcp::client& client) {
-    auto it = std::find_if(server.client_stack.begin(),
-                           server.client_stack.end(), [&](tcp::client& c) {
+  client_server.disconnect_event.add([&](tcp::client& client) {
+    auto it = std::find_if(client_server.client_stack.begin(),
+                           client_server.client_stack.end(), [&](tcp::client& c) {
                              return client.get_socket() == client.get_socket();
                            });
 
-    server.client_stack.erase(it);
+    client_server.client_stack.erase(it);
     client.cleanup();
 
     io::logger->info("{} disconnected", client.get_ip());
   });
 
-  server.receive_event.add([&](tcp::packet_t& packet, tcp::client& client) {
+  client_server.receive_event.add([&](tcp::packet_t& packet, tcp::client& client) {
     auto session = client.get_session();
     auto packet_session = packet.session_id;
     auto ip = client.get_ip();
@@ -68,6 +78,6 @@ int main(int argc, char* argv[]) {
     
   });
 
-  std::thread t{tcp::server::monitor, std::ref(server)};
+  std::thread t{tcp::server::monitor, std::ref(client_server)};
   t.join();
 }
diff --git a/server/src/util/enc.cpp b/server/src/server/enc.cpp
index 6a7baed..6a7baed 100644
--- a/server/src/util/enc.cpp
+++ b/server/src/server/enc.cpp
diff --git a/server/src/util/enc.h b/server/src/server/enc.h
index ae8d5a6..ae8d5a6 100644
--- a/server/src/util/enc.h
+++ b/server/src/server/enc.h
diff --git a/server/src/server/packet.h b/server/src/server/packet.h
index 57c6229..3c27c2d 100644
--- a/server/src/server/packet.h
+++ b/server/src/server/packet.h
@@ -1,5 +1,5 @@
 #pragma once
-#include "../util/enc.h"
+#include "enc.h"
 
 namespace tcp {
 constexpr size_t session_id_len = 10;
diff --git a/server/src/util/util.cpp b/server/src/util/util.cpp
new file mode 100644
index 0000000..799663f
--- /dev/null
+++ b/server/src/util/util.cpp
@@ -0,0 +1,7 @@
+#include "../include.h"
+#include "util.h"
+
+
+void util::to_lowercase(std::string &str) {
+	std::transform(str.begin(), str.end(), str.begin(), ::tolower);
+}
\ No newline at end of file
diff --git a/server/src/util/util.h b/server/src/util/util.h
new file mode 100644
index 0000000..4e5a0db
--- /dev/null
+++ b/server/src/util/util.h
@@ -0,0 +1,5 @@
+#pragma once
+
+namespace util {
+	void to_lowercase(std::string &str);
+};
\ No newline at end of file