From adad9f3b3122f0bc76a4bccc30083136813a6c82 Mon Sep 17 00:00:00 2001
From: marcelb <marcelb@bitelex.ml>
Date: Sat, 6 Jan 2024 22:22:20 +0100
Subject: [PATCH] Add white list feature

---
 lib/ipban.hpp | 25 ++++++++++++++++++++++--
 src/ipban.cpp | 53 ++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/lib/ipban.hpp b/lib/ipban.hpp
index 09c3658..ff737c7 100644
--- a/lib/ipban.hpp
+++ b/lib/ipban.hpp
@@ -8,6 +8,7 @@
 #include <future>
 #include <mutex>
 #include <fstream>
+#include <algorithm>
 
 #include <time.h>
 #include <unistd.h>
@@ -44,13 +45,14 @@ struct _fail {
  * Posjeduje vlastiti DB mehanizam za zaštitu od nepovratnog ban-a
 */
 class ipban {
-    mutex io, f_io;
+    mutex io, f_io, wl_io;
     time_t ban_duration;
     uint fail_interval;
     uint fail_limit;
     string db_file;
     vector<_ban> banned;
     map<string, struct _fail> failed;
+    vector<string> white_list;
     future<void> unban_bot;
     bool run_unban_bot = true;
     // interface možda bude trebao za ban
@@ -100,7 +102,7 @@ class ipban {
      * ako se prekorači broj dozvoljenih grešaka u intervalu - adresa se banuje
     */
 
-    void fail(const string& ip);
+    bool fail(const string& ip);
 
     /**
      * Uklanja greške za prosljeđenu adresu
@@ -108,6 +110,25 @@ class ipban {
 
     bool unfail(const string& ip);
 
+    /**
+     * Dodaje proslijeđenu adresu u white listu
+    */
+
+    void add_white_list(const string& ip);
+
+    /**
+     * Dodaje proslijeđene adrese u white listu
+    */
+
+    void add_white_list(const vector<string>& ips);
+
+    /**
+     * Provjerava da li je prosljeđena adresa u white listi
+     * Ako je vraća true, ako ne false
+    */
+
+    bool is_in_white_list(const string& ip);
+
     /**
      * Destruktor, uklanja sve zabrane.
     */
diff --git a/src/ipban.cpp b/src/ipban.cpp
index 113f62d..41fa607 100644
--- a/src/ipban.cpp
+++ b/src/ipban.cpp
@@ -73,11 +73,14 @@ bool marcelb::ipban::update_db() {
 }
 
 bool marcelb::ipban::ban(const string& ip) {
-    bool status = ufw_ban(ip);
-    io.lock();
-    banned.push_back({ip, time(NULL)});
-    status = status && update_db();
-    io.unlock();
+    bool status = !is_in_white_list(ip);
+    if (status) {
+        status = ufw_ban(ip);
+        io.lock();
+        banned.push_back({ip, time(NULL)});
+        status = status && update_db();
+        io.unlock();
+    }
     return status;
 }
 
@@ -91,7 +94,7 @@ bool marcelb::ipban::unban(vector<_ban>::iterator ban_itr) {
 }
 
 bool marcelb::ipban::ufw_ban(const string& ip) {
-    string ufw_cmd = "sudo ufw deny from " + ip + " to any";
+    string ufw_cmd = "sudo ufw insert 1 deny from " + ip + " to any";
     try {
         string execute_res = exec(ufw_cmd);
         if (execute_res == "Rule added\n") {
@@ -118,15 +121,19 @@ bool marcelb::ipban::ufw_unban(const string& ip) {
 }
 
 
-void marcelb::ipban::fail(const string& ip) {
+bool marcelb::ipban::fail(const string& ip) {
     lock_guard<mutex> _io(f_io);
-    if (failed[ip].n_fails == 0) {
-        failed[ip].n_fails = 1;
-        failed[ip].first_fail = time(NULL);
-    } else if (++failed[ip].n_fails >= fail_limit) {
-        ban(ip);
-        failed.erase(ip);
+    bool status = !is_in_white_list(ip);
+    if (status) {
+        if (failed[ip].n_fails == 0) {
+            failed[ip].n_fails = 1;
+            failed[ip].first_fail = time(NULL);
+        } else if (++failed[ip].n_fails >= fail_limit) {
+            ban(ip);
+            failed.erase(ip);
+        }
     }
+    return status;
 }
 
 
@@ -136,6 +143,26 @@ bool marcelb::ipban::unfail(const string& ip) {
 }
 
 
+void marcelb::ipban::add_white_list(const string& ip) {
+    lock_guard<mutex> _io(wl_io);
+    if (find(white_list.begin(), white_list.end(), ip) == white_list.end()) {
+        white_list.push_back(ip);
+    }
+}
+
+
+void marcelb::ipban::add_white_list(const vector<string>& ips) {
+    for (auto ip : ips) {
+        add_white_list(ip);
+    }
+}
+
+
+bool marcelb::ipban::is_in_white_list(const string& ip) {
+    return find(white_list.begin(), white_list.end(), ip) != white_list.end();
+}
+
+
 static void marcelb::sleep_if(const uint& _time, const bool& _condition) {
     time_t start_time = time(NULL);
     do {