Started security.

2 files changed, 59 insertions, 14 deletions

diff --git a/server/src/image/pe.h b/server/src/image/pe.h
index 23a0e4f..1025f13 100644
--- a/server/src/image/pe.h
+++ b/server/src/image/pe.h
@@ -18,6 +18,7 @@ template <bool x64 = false>
 class image {
   win::image_t<x64> *m_image;
   std::vector<char> m_buffer;
+  std::string m_name;
 
   std::unordered_map<std::string, std::vector<import_t>> m_imports;
   std::vector<section_t> m_sections;
@@ -25,7 +26,7 @@ class image {
 
  public:
   image() = default;
-  image(const std::string_view name) : m_image{nullptr} {
+  image(const std::string_view name) : m_image{nullptr}, m_name{name} {
     if (!io::read_file(name, m_buffer)) {
       io::logger->error("failed to load image {}.", name);
       return;
@@ -41,15 +42,17 @@ class image {
     parse_imports();
   }
 
-  void reload(const std::string_view name) {
-    io::read_file(name, m_buffer);
+  void reload() {
+    io::read_file(m_name, m_buffer);
     if (m_buffer.empty()) {
-      io::logger->error("failed to reload image {}.", name);
+      io::logger->error("failed to reload image {}.", m_name);
       return;
     }
 
     m_image = reinterpret_cast<win::image_t<x64> *>(m_buffer.data());
     load();
+
+    io::logger->info("reloaded {}.", m_name);
   }
 
   void parse_sections() {
@@ -103,16 +106,20 @@ class image {
         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;
+        if (thunk->is_ordinal) {
+          io::logger->error("found import by ordinal in module {}, {}.",
+                            mod_name, m_name);
+          continue;
+        }
 
-          std::transform(mod_name.begin(), mod_name.end(), mod_name.begin(),
-                         ::tolower);
+        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));
-        }
+        std::transform(mod_name.begin(), mod_name.end(), mod_name.begin(),
+                       ::tolower);
+
+        m_imports[mod_name].emplace_back(std::move(data));
       }
     }
   }
@@ -124,6 +131,10 @@ class image {
     out.resize(nt->optional_header.size_image);
 
     for (auto &sec : m_sections) {
+      if(sec.name == ".reloc" || sec.name == ".rsrc" || sec.name == ".idata") {
+        continue;
+      }
+
       std::memcpy(&out[sec.va], &m_buffer[sec.rva], sec.size);
     }
   }
diff --git a/server/src/main.cpp b/server/src/main.cpp
index 2a50ecd..80dbe62 100644
--- a/server/src/main.cpp
+++ b/server/src/main.cpp
@@ -18,12 +18,23 @@ int main(int argc, char* argv[]) {
 
   client_server.start();
 
+  uint16_t ver;
+  for(int i = 0; i < version.size(); ++i) {
+    if (i % 2) {
+      continue;
+    }
+
+    ver += static_cast<uint8_t>(version[i]) << 5;
+  }
+  io::logger->info("client version {}.", ver);
+
   client_server.connect_event.add([&](tcp::client& client) {
     auto ip = client.get_ip();
 
     io::logger->info("{} connected.", ip);
 
     client.gen_session();
+
     client.write(tcp::packet_t(version, tcp::packet_type::write, client(),
                                tcp::packet_id::session));
 
@@ -163,11 +174,11 @@ int main(int argc, char* argv[]) {
           }
 
           json["result"] = tcp::client_response::login_success;
-          json["games"]["notepad"] = {{"version", "0.1"},
+          json["games"]["notepad"] = {{"version", 1},
                                       {"id", 0},
                                       {"process", "notepad++.exe"},
                                       {"x64", false}};
-          json["games"]["sublime text"] = {{"version", "0.1"},
+          json["games"]["sublime text"] = {{"version", 1},
                                            {"id", 1},
                                            {"process", "sublime_text.exe"},
                                            {"x64", true}};
@@ -382,6 +393,29 @@ int main(int argc, char* argv[]) {
     io::logger->info("{} timed out.", client.get_ip());
   });
 
+
+  commands cmds;
+  cmds.add("reload", [&]() {
+    for(auto&[key, image] : client_server.images) {
+      image.reload();
+    }
+
+    for(auto&[key, image] : client_server.images64) {
+      image.reload();
+    }
+  });
+
+  std::thread t1{[&](tcp::server &srv) {
+    while (srv) {
+      std::string cmd;
+      getline(std::cin, cmd);
+      if(!cmds.parse_input(cmd)) {
+        io::logger->warn("invalid command.");
+      }
+    };
+  }, std::ref(client_server)};
+  t1.detach();
+
   std::thread t{tcp::server::monitor, std::ref(client_server)};
   t.join();
 }