improving martillo-maldito to server to protect also nodes (port, ip, etc..)

src/main.rs

@@ -3,53 +3,72 @@ pub mod iptables_save;
 pub mod iptables_wrapper;
 pub mod login_attempt;
 
-use cli::Arguments;
+use cli::Cli;
 use iptables_wrapper::{
-    allow_ip_port, ban_port, is_port_secured, list_banned_ips, remove_ip_port, unban_port,
+    allow_ip_for_port, is_port_secured, list_banned_ips, list_secured_ports,
+    map_secured_ports_allowed_ips, remove_allow_ip_for_port, secure_port, unsecure_port,
 };
 use linemux::MuxedLines;
 use login_attempt::LoginAttempt;
 use std::path::PathBuf;
-use std::thread;
+use std::thread::spawn;
 use std::{collections::HashMap, thread::sleep, time::Duration};
 use structopt::StructOpt;
 
 #[tokio::main]
 async fn main() {
-    let opts = Arguments::from_args();
-    match opts {
-        Arguments::BanServer {
+    match Cli::from_args() {
+        Cli::BanServer {
             ssh_auth_log,
             iptables_save,
         } => {
             let _ = start_ban_server(ssh_auth_log, iptables_save).await;
         }
-        Arguments::ListBannedIps => {
-            for ip in list_banned_ips() {
-                println!("{}", ip);
-            }
+        Cli::ListBannedIps => {
+            println!("{}", serde_json::to_string(&list_banned_ips()).unwrap());
         }
-        Arguments::IsPortSecured { port, docker } => {
-            let is_secured = is_port_secured(port, docker);
-            println!("{}", is_secured);
+        Cli::ListSecuredPorts { docker } => {
+            println!(
+                "{}",
+                serde_json::to_string(&list_secured_ports(docker)).unwrap()
+            );
         }
-        Arguments::BanPort {
+        Cli::MapSecuredPortsAllowedIps { docker } => {
+            println!(
+                "{}",
+                serde_json::to_string(&map_secured_ports_allowed_ips(docker)).unwrap()
+            );
+        }
+        Cli::IsPortSecured { port, docker } => {
+            println!("{}", is_port_secured(port, docker));
+        }
+        Cli::SecurePort {
             port,
             docker,
             position,
-        } => ban_port(port, docker, position),
-        Arguments::UnbanPort { port, docker } => unban_port(port, docker),
-        Arguments::AllowIpPort {
+        } => secure_port(port, docker, position),
+        Cli::UnsecurePort { port, docker } => unsecure_port(port, docker),
+        Cli::AllowIpForPort {
             ip,
             port,
             docker,
             position,
-        } => allow_ip_port(&ip, port, docker, position),
-        Arguments::OnlyIpPort { ip, port, docker } => {
-            allow_ip_port(&ip, port, docker, Some(1));
-            ban_port(port, docker, Some(2));
+        } => allow_ip_for_port(&ip, port, docker, position),
+        Cli::OnlyIpForPort { ip, port, docker } => {
+            allow_ip_for_port(&ip, port, docker, Some(1));
+            secure_port(port, docker, Some(2));
+        }
+        Cli::RemoveAllowIpPort { ip, port, docker } => remove_allow_ip_for_port(&ip, port, docker),
+        Cli::SaveIPTables { iptables_save } => {
+            let path = if let Some(iptables_save) = iptables_save {
+                iptables_save
+            } else {
+                PathBuf::from("/etc/iptables/rules.v4")
+            };
+
+            iptables_save::save_iptables(&path);
+            println!("Saved IPTables to {}", path.display());
         }
-        Arguments::RemoveIpPort { ip, port, docker } => remove_ip_port(&ip, port, docker),
     }
 }
 
@@ -65,47 +84,51 @@ async fn start_ban_server(
     if let Some(iptables_save) = iptables_save {
         let seconds_iptables = Duration::from_secs(60);
         println!(
-            "starting iptables-save, save every {} seconds",
+            "Saving IPTables every {} seconds",
             seconds_iptables.as_secs()
         );
-        thread::spawn(move || loop {
+
+        spawn(move || loop {
             sleep(seconds_iptables);
             iptables_save::save_iptables(&iptables_save);
         });
     }
 
-    println!("listening to changes over {}", ssh_auth_log.display());
-    while let Ok(Some(line)) = lines.next_line().await {
-        if let Some(login_attempt) = LoginAttempt::capture(line.line()) {
-            println!(
-                "failed login attempt from {}@{}:{}",
-                login_attempt.user, login_attempt.ip, login_attempt.port
-            );
-            match login_attempts.get_mut(&login_attempt.ip) {
-                Some(count) => {
-                    *count += 1;
+    println!("Listeging to changer over file: {}", ssh_auth_log.display());
+    loop {
+        while let Ok(Some(line)) = lines.next_line().await {
+            if let Some(login_attempt) = LoginAttempt::capture(line.line()) {
+                println!(
+                    "Failed login attempt from {}@{}:{}",
+                    login_attempt.user, login_attempt.ip, login_attempt.port
+                );
 
-                    if *count == 3 {
-                        match iptables.append_unique(
-                            "filter",
-                            "INPUT",
-                            &format!("--source {} -j DROP", login_attempt.ip),
-                        ) {
-                            Ok(_) => {
-                                println!("{} banned", login_attempt.ip);
-                            }
-                            Err(_) => {
-                                println!("{} already banned", login_attempt.ip);
+                match login_attempts.get_mut(&login_attempt.ip) {
+                    Some(count) => {
+                        *count += 1;
+
+                        if *count == 3 {
+                            if iptables
+                                .append_unique(
+                                    "filter",
+                                    "INPUT",
+                                    &format!("--source {} -j DROP", login_attempt.ip),
+                                )
+                                .is_ok()
+                            {
+                                println!("IP {} banned", login_attempt.ip);
+                            } else {
+                                println!("IP {} already banned", login_attempt.ip);
                             }
+
+                            login_attempts.remove(&login_attempt.ip);
                         }
-                        login_attempts.remove(&login_attempt.ip);
                     }
-                }
-                None => {
-                    login_attempts.insert(login_attempt.ip.clone(), 1);
+                    None => {
+                        login_attempts.insert(login_attempt.ip.clone(), 1);
+                    }
                 }
             }
         }
     }
-    Ok(())
 }