adding tcp command server to ban ips

src/main.rs

@@ -1,8 +1,11 @@
 pub mod iptables_save;
 pub mod login_attempt;
+pub mod tpc_command_server;
 
+use crate::tpc_command_server::start_tcp_command_server;
 use linemux::MuxedLines;
 use login_attempt::LoginAttempt;
+use std::thread;
 use std::{collections::HashMap, thread::sleep, time::Duration};
 
 #[tokio::main]
@@ -17,12 +20,15 @@ async fn main() -> std::io::Result<()> {
         "starting iptables-save, run every {} seconds",
         seconds_iptables.as_secs()
     );
-    tokio::spawn(async move {
-        loop {
-            sleep(seconds_iptables);
-            iptables_save::save_iptables();
-            println!("saved iptables rules");
-        }
+
+    thread::spawn(move || loop {
+        sleep(seconds_iptables);
+        iptables_save::save_iptables();
+        println!("saved iptables rules");
+    });
+
+    thread::spawn(|| {
+        start_tcp_command_server();
     });
 
     println!("listening to changes over /host_ssh/auth.log");

src/tpc_command_server.rs

@@ -0,0 +1,91 @@
+use std::io::{BufRead, BufReader, Write};
+use std::net::{TcpListener, TcpStream};
+use std::thread;
+
+pub fn start_tcp_command_server() {
+    let listener = TcpListener::bind("127.0.0.1:9632").unwrap();
+    println!("listening on port 9632 for tcp commands");
+
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                thread::spawn(move || handle_client(&stream));
+            }
+            Err(e) => {
+                eprintln!("err with tcp conn: {}", e);
+            }
+        }
+    }
+}
+
+fn handle_client(mut stream: &TcpStream) {
+    let reader = BufReader::new(stream);
+
+    for line in reader.lines() {
+        let buffer = match line {
+            Ok(data) => data,
+            Err(_) => return,
+        };
+
+        if buffer.is_empty() {
+            break;
+        }
+
+        let mut parts = buffer.trim().split_whitespace();
+        if let Some(command) = parts.next() {
+            let arguments: Vec<&str> = parts.collect();
+            let response = handle_command(command, arguments);
+            stream.write_all(response.as_bytes()).unwrap();
+        } else {
+            stream.write_all("invalid command".as_bytes()).unwrap();
+        }
+    }
+}
+
+fn handle_command(command: &str, arguments: Vec<&str>) -> String {
+    match command {
+        "ban" => {
+            if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) {
+                let iptables = iptables::new(false).unwrap();
+                let _ = iptables.append_unique(
+                    "filter",
+                    "INPUT",
+                    &format!("-s {} -p tcp --dport {} -j ACCEPT", ip, port),
+                );
+
+                let _ = iptables.append_unique(
+                    "filter",
+                    "INPUT",
+                    &format!("-p tcp --dport {} -j DROP", port),
+                );
+
+                format!("banned port {}, only {} allowed", port, ip)
+            } else {
+                "missing args for ban: ip and port".to_string()
+            }
+        }
+        "unban" => {
+            if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) {
+                let iptables = iptables::new(false).unwrap();
+                let _ = iptables.delete(
+                    "filter",
+                    "INPUT",
+                    &format!("-s {} -p tcp --dport {} -j ACCEPT", ip, port),
+                );
+
+                let _ = iptables.delete(
+                    "filter",
+                    "INPUT",
+                    &format!("-p tcp --dport {} -j DROP", port),
+                );
+
+                format!("unbanned port {}, used for {}", port, ip)
+            } else {
+                "missing args for unban: ip and port".to_string()
+            }
+        }
+        _ => {
+            format!("unknown command: {}", command)
+        }
+    }
+}