From f7c2f313f71695f4ccaeafedd3ffa83f69982388 Mon Sep 17 00:00:00 2001
From: midefos <midefos@gmail.com>
Date: Mon, 20 May 2024 23:54:15 +0200
Subject: [PATCH] better commands

---
 src/tpc_command_server.rs | 58 +++++++++++++++++++++++----------------
 1 file changed, 35 insertions(+), 23 deletions(-)

diff --git a/src/tpc_command_server.rs b/src/tpc_command_server.rs
index 8f87560..5232c84 100644
--- a/src/tpc_command_server.rs
+++ b/src/tpc_command_server.rs
@@ -19,9 +19,7 @@ pub fn start_tcp_command_server() {
 }
 
 fn handle_client(mut stream: &TcpStream) {
-    let reader = BufReader::new(stream);
-
-    for line in reader.lines() {
+    for line in BufReader::new(stream).lines() {
         let buffer = match line {
             Ok(data) => data,
             Err(_) => return,
@@ -44,27 +42,48 @@ fn handle_client(mut stream: &TcpStream) {
 
 fn handle_command(command: &str, arguments: Vec<&str>) -> String {
     match command {
-        "ban" => {
-            if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) {
+        "banport" => {
+            if let Some(port) = arguments.get(0) {
                 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)
+                format!("banned port {} for all ips", port)
             } else {
-                "missing args for ban: ip and port".to_string()
+                "missing args for banport: port".to_string()
             }
         }
-        "unban" => {
+        "unbanport" => {
+            if let Some(port) = arguments.get(0) {
+                let iptables = iptables::new(false).unwrap();
+                let _ = iptables.delete(
+                    "filter",
+                    "INPUT",
+                    &format!("-p tcp --dport {} -j DROP", port),
+                );
+
+                format!("unbanned port {}", port)
+            } else {
+                "missing args for unbanport: port".to_string()
+            }
+        }
+        "allowipport" => {
+            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),
+                );
+                format!("allowed {} to access {}", ip, port)
+            } else {
+                "missing args for allowipport: ip and port".to_string()
+            }
+        }
+        "removeipport" => {
             if let (Some(ip), Some(port)) = (arguments.get(0), arguments.get(1)) {
                 let iptables = iptables::new(false).unwrap();
                 let _ = iptables.delete(
@@ -72,16 +91,9 @@ fn handle_command(command: &str, arguments: Vec<&str>) -> String {
                     "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)
+                format!("rm {} access to {}", ip, port)
             } else {
-                "missing args for unban: ip and port".to_string()
+                "missing args for rmipport: ip and port".to_string()
             }
         }
         _ => {