Added client timeout.

Added json imports for pe images.
Added wrapper to support streaming strings.

7 files changed, 70 insertions, 22 deletions

diff --git a/server/src/client/client.cpp b/server/src/client/client.cpp
index 75b25f2..1163ab5 100644
--- a/server/src/client/client.cpp
+++ b/server/src/client/client.cpp
@@ -5,13 +5,13 @@
 bool tcp::client::init_ssl(SSL_CTX* server_ctx) {
   m_ssl = SSL_new(server_ctx);
   if (!m_ssl) {
-    io::logger->error("failed to create ssl on client {}.", m_ip);
+    io::logger->error("failed to create ssl on {}.", m_ip);
     return false;
   }
 
   int ret = SSL_set_fd(m_ssl, m_socket);
   if (ret <= 0) {
-    io::logger->error("failed to set descriptor on client {}.", m_ip);
+    io::logger->error("failed to set descriptor on {}.", m_ip);
     return false;
   }
 
@@ -19,7 +19,7 @@ bool tcp::client::init_ssl(SSL_CTX* server_ctx) {
 
   if (ret <= 0) {
     int err = SSL_get_error(m_ssl, ret);
-    io::logger->error("client {} failed to accept ssl, return code {}", m_ip,
+    io::logger->error("{} failed to accept ssl, return code {}", m_ip,
                       err);
     return false;
   }
diff --git a/server/src/client/client.h b/server/src/client/client.h
index 1dbb8d5..52e0c42 100644
--- a/server/src/client/client.h
+++ b/server/src/client/client.h
@@ -7,7 +7,7 @@ class client {
   int m_socket;
   SSL* m_ssl;
 
-  time_t m_time;
+  std::time_t m_time;
 
   std::string m_ip;
   std::string m_session_id;
@@ -26,20 +26,35 @@ class client {
     SSL_free(m_ssl);
   }
 
-  int write(const void* data, size_t size) {
-    return SSL_write(m_ssl, data, size);
-  }
+  void reset() { std::time(&m_time); }
+  bool timeout() { return std::difftime(std::time(nullptr), m_time) >= 30; }
 
   int write(const packet_t& packet) {
     if (!packet) return 0;
     return write(packet.message.data(), packet.message.size());
   }
 
+  int write(const void* data, size_t size) {
+    return SSL_write(m_ssl, data, size);
+  }
+
   int read(void* data, size_t size) { return SSL_read(m_ssl, data, size); }
 
-  int stream(std::vector<char>& data, float *dur = nullptr);
+  int stream(std::vector<char>& data, float* dur = nullptr);
   int read_stream(std::vector<char>& out);
 
+  int stream(std::string &str) {
+    std::vector<char> vec(str.begin(), str.end());
+    return stream(vec);
+  }
+
+  int read_stream(std::string &str) {
+    std::vector<char> out;
+    int ret = read_stream(out);
+    str.assign(out.begin(), out.end());
+    return ret;
+  }
+
   void gen_session();
 
   int get_socket() { return m_socket; }
diff --git a/server/src/image/pe.h b/server/src/image/pe.h
index 2f81616..3bbe772 100644
--- a/server/src/image/pe.h
+++ b/server/src/image/pe.h
@@ -124,6 +124,16 @@ class image {
   auto imports() const { return m_imports; }
   auto relocs() const { return m_relocs; }
   auto sections() const { return m_sections; }
+
+  std::string get_json_imports() {
+    nlohmann::json json;
+    for(auto &[mod, imports] : m_imports) {
+      for(auto &i : imports) {
+        json[mod].emplace_back(std::make_pair(i.name, i.rva));
+      }
+    }
+    return json.dump();
+  }
 };
 
 };  // namespace pe
\ No newline at end of file
diff --git a/server/src/main.cpp b/server/src/main.cpp
index f5b58be..138f733 100644
--- a/server/src/main.cpp
+++ b/server/src/main.cpp
@@ -9,6 +9,8 @@ constexpr std::string_view client_version{"0.1.0"};
 int main(int argc, char* argv[]) {
   io::init(false);
 
+  //pe::image image("test.dll");
+
   tcp::server client_server("6666");
 
   client_server.start();
@@ -41,7 +43,7 @@ int main(int argc, char* argv[]) {
     auto message = packet();
 
     if (!packet) {
-      io::logger->info("{} sent invalid packet, id {}", ip, packet.id);
+      io::logger->info("{} sent invalid packet", ip);
       return;
     }
 
@@ -52,9 +54,11 @@ int main(int argc, char* argv[]) {
 
     io::logger->info("{} : {}", packet_session, message);
 
-    tcp::packet_t resp(message, tcp::packet_type::write,
-                       client.get_session());
-    client.write(resp);
+    client.write(tcp::packet_t(message, tcp::packet_type::write,
+                       client.get_session()));
+
+    /*auto imports = image.get_json_imports();
+    client.stream(imports);*/
 
     /*std::vector<char> t;
     io::read_file("test.dll", t);
@@ -69,6 +73,13 @@ int main(int argc, char* argv[]) {
     
   });
 
+  client_server.timeout_event.add([&](tcp::client& client) {
+    client.write(tcp::packet_t("timedout", tcp::packet_type::write,
+                               client.get_session()));
+
+    io::logger->info("{} timed out.", client.get_ip());
+  });
+
   std::thread t{tcp::server::monitor, std::ref(client_server)};
   t.join();
 }
diff --git a/server/src/server/packet.h b/server/src/server/packet.h
index a385960..661b0df 100644
--- a/server/src/server/packet.h
+++ b/server/src/server/packet.h
@@ -5,21 +5,18 @@
 
 namespace tcp {
 constexpr size_t session_id_len = 10;
-constexpr size_t message_len = 1024;
+constexpr size_t message_len = 512;
 
 enum packet_type : int { write = 0, read };
 
 struct packet_t {
   std::string message;
   std::string session_id;
-  int id;
 
   packet_t() {}
   packet_t(const std::string_view msg, const packet_type& type,
            std::string_view session = "") {
     if (type == read) {
-      ++id;
-
       message = msg;
       enc::decrypt_message(message);
 
diff --git a/server/src/server/server.cpp b/server/src/server/server.cpp
index b0383bd..5996eb3 100644
--- a/server/src/server/server.cpp
+++ b/server/src/server/server.cpp
@@ -99,18 +99,18 @@ void tcp::server::accept_client() {
       cli.cleanup();
       return;
     }
-
     // check for an existing connection
-    auto it = std::find_if(client_stack.begin(), client_stack.end(),
+    /*auto it = std::find_if(client_stack.begin(), client_stack.end(),
                            [&](client& c) { return c.get_ip() == ip; });
     if (it != client_stack.end()) {
       io::logger->info("{} is already connected, dropping...", ip);
       cli.cleanup();
       return;
-    }
+    }*/
+    cli.reset();
 
     connect_event.call(cli);
-    client_stack.emplace_back(cli);
+    client_stack.emplace_back(std::move(cli));
   }
 }
 
@@ -125,6 +125,8 @@ void tcp::server::receive() {
 
     const int read = c.read(&buf[0], buf.size());
     if (read > 0) {
+      c.reset();
+
       std::string msg(buf.data(), read);
 
       tcp::packet_t packet(msg, tcp::packet_type::read);
@@ -136,6 +138,18 @@ void tcp::server::receive() {
   }
 }
 
+void tcp::server::check_timeout() {
+  auto it = std::find_if(client_stack.begin(), client_stack.end(),
+                         [&](client& c) { return c.timeout(); });
+
+  if (it != client_stack.end()) {
+    timeout_event.call(*it);
+
+    it->cleanup();
+    client_stack.erase(it);
+  }
+}
+
 void tcp::server::stop() {
   io::logger->info("stopping server on port {}.", m_port);
   close(m_socket);
diff --git a/server/src/server/server.h b/server/src/server/server.h
index 90c31f7..1249784 100644
--- a/server/src/server/server.h
+++ b/server/src/server/server.h
@@ -19,10 +19,10 @@ class server {
 
   std::vector<tcp::client> client_stack;
  public:
-
   event<client&> connect_event;
   event<packet_t&, client&> receive_event;
   event<client&> disconnect_event;
+  event<client&> timeout_event;
 
   server(const std::string_view port) : m_port{port}, m_active{false} {}
   ~server() = default;
@@ -32,6 +32,7 @@ class server {
   void accept_client();
   void receive();
   void stop();
+  void check_timeout();
 
   operator bool() const { return m_active; }
   auto &operator()() { return client_stack; }
@@ -43,7 +44,7 @@ class server {
         srv.accept_client();
         srv.receive();
       } else if (ret == select_status::standby) {
-        // check for timeout
+        srv.check_timeout();
       } else {
         break;
       }