Added blacklist implementation on server.

Overall code cleanup and optimization.

10 files changed, 150 insertions, 47 deletions

diff --git a/server/src/client/blacklist.h b/server/src/client/blacklist.h
new file mode 100644
index 0000000..ea9e261
--- /dev/null
+++ b/server/src/client/blacklist.h
@@ -0,0 +1,54 @@
+#pragma once
+
+namespace tcp {
+
+struct blacklist_data {
+	std::string ip;
+	std::string hwid;
+};
+
+class blacklist {
+
+nlohmann::json m_data;
+std::string m_name;
+
+public:
+	void init(const std::string_view file = "blacklist") {
+		m_name = file;
+
+		std::string data;
+		if(!io::read_file(file, data))
+			return;
+
+		if(!nlohmann::json::accept(data)) {
+			io::logger->error("blacklist file isnt valid json.");
+			return;
+		}
+
+		m_data = nlohmann::json::parse(data);
+	}
+
+	void add(const blacklist_data &data) {
+		m_data["ips"].emplace_back(data.ip);
+		m_data["hwids"].emplace_back(data.hwid);
+
+		save();
+	}
+
+	void save() {
+		std::ofstream o(m_name, std::ios::trunc);
+		o << std::setw(4) << m_data;
+		o.close();
+	}
+
+	bool find(const std::string &key) {
+		for(auto &item : m_data["ips"]) {
+			if(item.get<std::string>() == key) {
+				return true;
+			}
+		}
+		return false;
+	}
+};
+
+}; 
\ No newline at end of file
diff --git a/server/src/client/client.h b/server/src/client/client.h
index 141ac67..a3d558d 100644
--- a/server/src/client/client.h
+++ b/server/src/client/client.h
@@ -24,8 +24,10 @@ class client {
 
   void cleanup() {
     close(m_socket);
-    SSL_shutdown(m_ssl);
-    SSL_free(m_ssl);
+    if (m_ssl) {
+      SSL_shutdown(m_ssl);
+      SSL_free(m_ssl);
+    }
 
     m_socket = -1;
   }
@@ -63,8 +65,9 @@ class client {
 
   int& get_socket() { return m_socket; }
   auto& get_ip() { return m_ip; }
-  auto& get_session() { return m_session_id; }
 
-  operator bool() const { return m_ssl && m_socket > 0; }
+  operator bool() const { return m_socket > 0; }
+  auto &operator()() { return m_session_id; }
+
 };
 };  // namespace tcp
\ No newline at end of file
diff --git a/server/src/image/pe.h b/server/src/image/pe.h
index 3bbe772..7243262 100644
--- a/server/src/image/pe.h
+++ b/server/src/image/pe.h
@@ -29,9 +29,7 @@ class image {
 
  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);
+    if (!io::read_file(name, m_buffer)) {
       return;
     }
 
diff --git a/server/src/include.h b/server/src/include.h
index 9fe4e09..9dd3255 100644
--- a/server/src/include.h
+++ b/server/src/include.h
@@ -22,6 +22,7 @@
 #include <utility>
 #include <atomic>
 #include <mutex>
+#include <iomanip>
 
 #include <arpa/inet.h>
 #include <netdb.h>
diff --git a/server/src/main.cpp b/server/src/main.cpp
index e7e0253..6628712 100644
--- a/server/src/main.cpp
+++ b/server/src/main.cpp
@@ -4,64 +4,84 @@
 #include "server/server.h"
 #include "image/pe.h"
 
-constexpr std::string_view client_version{"0.1.0"};
+constexpr std::string_view version{"0.1.0"};
 
 int main(int argc, char* argv[]) {
   io::init(false);
 
-  //pe::image image("test.dll");
+  // pe::image image("test.dll");
 
   tcp::server client_server("6666");
 
   client_server.start();
 
+  // client_server.bl().add({"127.0.0.1", "qsmdlkjqsdmlkqmdsl"});
+
   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(client_version, tcp::packet_type::write,
-                               client.get_session(),
+    client.write(tcp::packet_t(version, tcp::packet_type::write, client(),
                                tcp::packet_action::session));
-
-    io::logger->info("{} connected", ip);
   });
 
   client_server.disconnect_event.add([&](tcp::client& client) {
     client.cleanup();
 
-    io::logger->info("{} disconnected", client.get_ip());
+    auto it = std::find_if(
+        client_server().begin(), client_server().end(),
+        [&](tcp::client& c) { return c.get_socket() == client.get_socket(); });
+    
+    client_server().erase(it);
+
+    io::logger->info("{} disconnected.", client.get_ip());
   });
 
-  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();
-    auto message = packet();
-    auto action = packet.act;
+  client_server.receive_event.add(
+      [&](tcp::packet_t& packet, tcp::client& client) {
+        auto session = client();
+        auto packet_session = packet.session_id;
+        auto ip = client.get_ip();
+        auto message = packet();
+        auto action = packet.act;
 
-    if (!packet) {
-      io::logger->info("{} sent invalid packet", ip);
-      return;
-    }
+        if (!packet) {
+          io::logger->info("{} sent invalid packet.", ip);
 
-    if (packet_session != session) {
-      io::logger->info("{} sent wrong session id", ip);
-      return;
-    }
+          //client_server.disconnect_event.call(client);
+          return;
+        }
 
-    io::logger->info("{} : {}", packet_session, message);
+        if (packet_session != session) {
+          io::logger->info("{} sent wrong session id.", ip);
+          io::logger->info("{} : {}", packet_session, session);
 
-    if(action == tcp::packet_action::hwid) {
-      client.hwid = message;
+          //client_server.disconnect_event.call(client);
+          return;
+        }
 
-      io::logger->info("got hwid from {} : {}", ip, message);
-    }
+        io::logger->info("{} : {}", packet_session, message);
 
-    //client.write(tcp::packet_t(message, tcp::packet_type::write,
-                       //client.get_session()));
-  });
+        if (action == tcp::packet_action::hwid) {
+          client.hwid = message;
+
+          io::logger->info("got hwid from {} : {}", ip, message);
+        }
+
+        client.write(tcp::packet_t(message, tcp::packet_type::write, session));
+      });
 
   client_server.timeout_event.add([&](tcp::client& client) {
     client.cleanup();
+
+    auto it = std::find_if(
+        client_server().begin(), client_server().end(),
+        [&](tcp::client& c) { return c.get_socket() == client.get_socket(); });
+
+    client_server().erase(it);
+
     io::logger->info("{} timed out.", client.get_ip());
   });
 
diff --git a/server/src/server/enc.cpp b/server/src/server/enc.cpp
index 6a7baed..97e1c29 100644
--- a/server/src/server/enc.cpp
+++ b/server/src/server/enc.cpp
@@ -3,8 +3,9 @@
 
 namespace enc {
 
+std::random_device r;
+
 void encrypt_message(std::string &str) {
-  std::random_device r;
   std::default_random_engine e1(r());
   std::uniform_int_distribution<int> gen(0, 255);
 
diff --git a/server/src/server/server.cpp b/server/src/server/server.cpp
index b7f7606..2d4cafe 100644
--- a/server/src/server/server.cpp
+++ b/server/src/server/server.cpp
@@ -3,6 +3,8 @@
 #include "server.h"
 
 void tcp::server::start() {
+  m_blacklist.init();
+
   io::logger->info("starting server on port {}...", m_port.data());
 
   ssl ctx("ssl/server.crt", "ssl/server.key", "ssl/rootCA.crt");
@@ -57,13 +59,7 @@ tcp::select_status tcp::server::peek() {
 
   int maxfd = m_socket;
 
-  for (int i = 0; i < client_stack.size(); i++) {
-    auto c = client_stack[i];
-    if(!c) {
-      client_stack.erase(client_stack.begin() + i);
-      continue;
-    }
-
+  for (auto& c : client_stack) {
     const int s = c.get_socket();
     FD_SET(s, &m_server_set);
 
@@ -105,6 +101,16 @@ void tcp::server::accept_client() {
       cli.cleanup();
       return;
     }
+
+    //m_blacklist.add({"127.0.0.1", "ahahaahhahaha"});
+
+    // check for blacklist ip entry
+    if (m_blacklist.find(ip)) {
+      io::logger->info("{} is blacklisted, dropping...", ip);
+      cli.cleanup();
+      return;
+    }
+
     // check for an existing connection
     /*auto it = std::find_if(client_stack.begin(), client_stack.end(),
                            [&](client& c) { return c.get_ip() == ip; });
@@ -113,9 +119,11 @@ void tcp::server::accept_client() {
       cli.cleanup();
       return;
     }*/
+    
     cli.reset();
 
     connect_event.call(cli);
+
     client_stack.emplace_back(std::move(cli));
   }
 }
diff --git a/server/src/server/server.h b/server/src/server/server.h
index 1249784..ede0683 100644
--- a/server/src/server/server.h
+++ b/server/src/server/server.h
@@ -1,6 +1,7 @@
 #pragma once
 #include "../client/client.h"
 #include "../util/events.h"
+#include "../client/blacklist.h"
 #include "ssl.h"
 
 namespace tcp {
@@ -18,6 +19,8 @@ class server {
   std::atomic<bool> m_active;
 
   std::vector<tcp::client> client_stack;
+
+  blacklist m_blacklist;
  public:
   event<client&> connect_event;
   event<packet_t&, client&> receive_event;
@@ -37,6 +40,8 @@ class server {
   operator bool() const { return m_active; }
   auto &operator()() { return client_stack; }
 
+  auto &bl() { return m_blacklist; }
+
   static void monitor(server& srv) {
     while (srv) {
       auto ret = srv.peek();
diff --git a/server/src/util/io.cpp b/server/src/util/io.cpp
index aff8a53..d23cdda 100644
--- a/server/src/util/io.cpp
+++ b/server/src/util/io.cpp
@@ -22,11 +22,11 @@ void io::init(const bool& to_file) {
   spdlog::flush_every(std::chrono::seconds(1));
 }
 
-void io::read_file(const std::string_view name, std::vector<char>& out) {
+bool io::read_file(const std::string_view name, std::vector<char>& out) {
   std::ifstream file(name.data());
   if (!file.good()) {
     io::logger->error("failed to load {}.", name.data());
-    return;
+    return false;
   }
 
   file.unsetf(std::ios::skipws);
@@ -40,4 +40,16 @@ void io::read_file(const std::string_view name, std::vector<char>& out) {
   file.read(out.data(), size);
 
   file.close();
+
+  return true;
+}
+
+bool io::read_file(const std::string_view name, std::string& out) {
+  std::vector<char> vec;
+  if(!read_file(name, vec))
+    return false;
+
+  out.assign(vec.begin(), vec.end());
+
+  return true;
 }
diff --git a/server/src/util/io.h b/server/src/util/io.h
index 4681fbf..6a333bd 100644
--- a/server/src/util/io.h
+++ b/server/src/util/io.h
@@ -4,5 +4,6 @@ namespace io {
 extern std::shared_ptr<spdlog::logger> logger;
 
 void init(const bool& to_file);
-void read_file(const std::string_view name, std::vector<char>& out);
+bool read_file(const std::string_view name, std::vector<char>& out);
+bool read_file(const std::string_view name, std::string &out);
 };  // namespace io