improving docker ban and allow

src/cli.rs

@@ -15,7 +15,10 @@ pub enum Cli {
         iptables_save: Option<PathBuf>,
     },
     #[structopt(about = "List all banned ips")]
-    ListBannedIps,
+    ListBannedIps {
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
+    },
     #[structopt(about = "List all secured ports")]
     ListSecuredPorts {
         #[structopt(name = "Docker", short = "d", long = "docker")]

src/iptables_wrapper.rs

@@ -19,7 +19,9 @@ pub fn is_port_secured(port: u16, docker: bool) -> bool {
 
 pub fn list_secured_ports(docker: bool) -> Vec<u16> {
     let iptables = iptables::new(false).unwrap();
-    let rules = iptables.list("filter", &get_chain(docker));
+
+    let chain = get_chain(docker);
+    let rules = iptables.list("filter", &chain);
     if rules.is_err() {
         return vec![];
     }
@@ -33,9 +35,11 @@ pub fn list_secured_ports(docker: bool) -> Vec<u16> {
         .collect()
 }
 
-pub fn list_banned_ips() -> Vec<String> {
+pub fn list_banned_ips(docker: bool) -> Vec<String> {
     let iptables = iptables::new(false).unwrap();
-    let rules = iptables.list("filter", &get_chain(false));
+
+    let chain = get_chain(docker);
+    let rules = iptables.list("filter", &chain);
     if rules.is_err() {
         return vec![];
     }
@@ -44,7 +48,9 @@ pub fn list_banned_ips() -> Vec<String> {
     rules
         .unwrap()
         .iter()
-        .filter(|r| r.contains("-A INPUT") && r.contains("-j DROP") && r.contains("-s"))
+        .filter(|r| {
+            r.contains(&format!("-A {}", chain)) && r.contains("-j DROP") && r.contains("-s")
+        })
         .map(|r| extract_ip(&rgx, r).unwrap())
         .collect()
 }
@@ -57,7 +63,8 @@ pub fn map_secured_ports_allowed_ips(docker: bool) -> HashMap<u16, Vec<String>>
     }
 
     let iptables = iptables::new(false).unwrap();
-    let rules = iptables.list("filter", &get_chain(docker));
+    let chain = get_chain(docker);
+    let rules = iptables.list("filter", &chain);
     if rules.is_err() {
         return result;
     }
@@ -68,7 +75,7 @@ pub fn map_secured_ports_allowed_ips(docker: bool) -> HashMap<u16, Vec<String>>
         let ips = rules
             .iter()
             .filter(|r| {
-                r.contains("-A INPUT -s")
+                r.contains(&format!("-A {} -s", chain))
                     && r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port))
             })
             .map(|r| extract_ip(&rgx, r).unwrap())
@@ -79,8 +86,25 @@ pub fn map_secured_ports_allowed_ips(docker: bool) -> HashMap<u16, Vec<String>>
     result
 }
 
-fn secure_port_rule(port: u16) -> String {
+fn secure_port_rule(iptables: &IPTables, port: u16, docker: bool) -> Option<String> {
-    format!("-p tcp --dport {} -j DROP", port)
+    if docker {
+        let internal_ip = get_internal_ip_for_docker_port(iptables, port)?;
+        Some(format!(
+            "-d {} -p tcp --dport {} -j DROP",
+            internal_ip, port
+        ))
+    } else {
+        Some(format!("-p tcp --dport {} -j DROP", port))
+    }
+}
+
+fn get_internal_ip_for_docker_port(iptables: &IPTables, port: u16) -> Option<String> {
+    let rules = iptables.list("filter", "DOCKER").unwrap();
+    let rule = rules
+        .iter()
+        .find(|r| r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port)));
+
+    rule.map(|r| extract_ip(&get_regex_for_ip(), r).unwrap())
 }
 
 pub fn secure_port(
@@ -92,7 +116,12 @@ pub fn secure_port(
 
     let table = "filter";
     let chain = get_chain(docker);
-    let rule = secure_port_rule(port);
+    let rule = secure_port_rule(&iptables, port, docker);
+    if rule.is_none() {
+        return Err("Err gathering secure port rule".into());
+    }
+
+    let rule = rule.unwrap();
     if let Some(position) = position {
         insert_unique(&iptables, table, &chain, &rule, position)
     } else {
@@ -102,11 +131,31 @@ pub fn secure_port(
 
 pub fn unsecure_port(port: u16, docker: bool) -> Result<(), Box<dyn std::error::Error>> {
     let iptables = iptables::new(false).unwrap();
-    iptables.delete("filter", &get_chain(docker), &secure_port_rule(port))
+
+    let rule = secure_port_rule(&iptables, port, docker);
+    if rule.is_none() {
+        return Err("Err gathering secure port rule".into());
+    }
+
+    let rule = rule.unwrap();
+    iptables.delete("filter", &get_chain(docker), &rule)
 }
 
-fn allow_ip_for_port_rule(port: u16, ip: &str) -> String {
+fn allow_ip_for_port_rule(
-    format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip)
+    iptables: &IPTables,
+    port: u16,
+    ip: &str,
+    docker: bool,
+) -> Option<String> {
+    if docker {
+        let internal_ip = get_internal_ip_for_docker_port(iptables, port)?;
+        Some(format!(
+            "-p tcp --dport {} -s {} -d {} -j ACCEPT",
+            port, ip, internal_ip
+        ))
+    } else {
+        Some(format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip))
+    }
 }
 
 pub fn allow_ip_for_port(
@@ -119,7 +168,12 @@ pub fn allow_ip_for_port(
 
     let table = "filter";
     let chain = get_chain(docker);
-    let rule = allow_ip_for_port_rule(port, ip);
+    let rule = allow_ip_for_port_rule(&iptables, port, ip, docker);
+    if rule.is_none() {
+        return Err("Err gathering allow ip for port rule".into());
+    }
+
+    let rule = rule.unwrap();
     if let Some(position) = position {
         insert_unique(&iptables, table, &chain, &rule, position)
     } else {
@@ -133,11 +187,13 @@ pub fn remove_allow_ip_for_port(
     docker: bool,
 ) -> Result<(), Box<dyn std::error::Error>> {
     let iptables = iptables::new(false).unwrap();
-    iptables.delete(
+    let rule = allow_ip_for_port_rule(&iptables, port, ip, docker);
-        "filter",
+    if rule.is_none() {
-        &get_chain(docker),
+        return Err("Err gathering allow ip for port rule".into());
-        &allow_ip_for_port_rule(port, ip),
+    }
-    )
+
+    let rule = rule.unwrap();
+    iptables.delete("filter", &get_chain(docker), &rule)
 }
 
 fn get_chain(docker: bool) -> String {
@@ -199,9 +255,16 @@ mod tests {
     }
 
     #[test]
-    fn no_match_extrac_ip() {
+    fn no_match_extract_ip() {
         let regex = get_regex_for_ip();
         let input = "-A INPUT -j DROP";
         assert_eq!(extract_ip(&regex, input), None);
     }
+
+    #[test]
+    fn docker_extract_ip() {
+        let regex = get_regex_for_ip();
+        let input = "-A DOCKER -d 172.18.0.2/32 ! -i br-127d33df48a4 -o br-127d33df48a4 -p tcp -m tcp --dport 8078 -j ACCEPT";
+        assert_eq!(extract_ip(&regex, input), Some("172.18.0.2".to_string()));
+    }
 }

src/main.rs

@@ -10,9 +10,7 @@ use iptables_wrapper::{
 };
 use linemux::MuxedLines;
 use login_attempt::LoginAttempt;
-use std::path::PathBuf;
+use std::{collections::HashMap, path::PathBuf, thread::sleep, thread::spawn, time::Duration};
-use std::thread::spawn;
-use std::{collections::HashMap, thread::sleep, time::Duration};
 use structopt::StructOpt;
 
 #[tokio::main]
@@ -24,8 +22,11 @@ async fn main() {
         } => {
             let _ = start_ban_server(ssh_auth_log, iptables_save).await;
         }
-        Cli::ListBannedIps => {
+        Cli::ListBannedIps { docker } => {
-            println!("{}", serde_json::to_string(&list_banned_ips()).unwrap());
+            println!(
+                "{}",
+                serde_json::to_string(&list_banned_ips(docker)).unwrap()
+            );
         }
         Cli::ListSecuredPorts { docker } => {
             println!(