altering also docker chain

src/cli.rs

@@ -16,12 +16,17 @@ pub enum Arguments {
     BanPort {
         #[structopt(name = "Port to ban", short = "p", long = "port")]
         port: u16,
+
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
     },
 
     #[structopt(about = "Unban port")]
     UnbanPort {
         #[structopt(name = "Port to unban", short = "p", long = "port")]
         port: u16,
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
     },
 
     #[structopt(about = "Allow ip and port")]
@@ -31,6 +36,8 @@ pub enum Arguments {
 
         #[structopt(name = "Port to allow", short = "p", long = "port")]
         port: u16,
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
     },
 
     #[structopt(about = "Remove ip and port")]
@@ -40,5 +47,7 @@ pub enum Arguments {
 
         #[structopt(name = "Port to remove", short = "p", long = "port")]
         port: u16,
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
     },
 }

src/iptables_wrapper.rs

@@ -1,41 +1,50 @@
-pub fn ban_port(port: u16) {
+pub fn ban_port(port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
+
     let _ = iptables.append_unique(
         "filter",
-        "INPUT",
+        &get_chain(docker),
         &format!("-p tcp --dport {} -j DROP", port),
     );
 
     println!("banned port {}", port);
 }
 
-pub fn unban_port(port: u16) {
+pub fn unban_port(port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
     let _ = iptables.delete(
         "filter",
-        "INPUT",
+        &get_chain(docker),
         &format!("-p tcp --dport {} -j DROP", port),
     );
 
     println!("unbanned port {}", port);
 }
 
-pub fn allow_ip_port(ip: &str, port: u16) {
+pub fn allow_ip_port(ip: &str, port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
     let _ = iptables.append_unique(
         "filter",
-        "INPUT",
+        &get_chain(docker),
         &format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip),
     );
     println!("allowed {} to access {}", ip, port);
 }
 
-pub fn remove_ip_port(ip: &str, port: u16) {
+pub fn remove_ip_port(ip: &str, port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
     let _ = iptables.delete(
         "filter",
-        "INPUT",
+        &get_chain(docker),
         &format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip),
     );
     println!("removed access {} to {}", ip, port);
 }
+
+fn get_chain(docker: bool) -> String {
+    if docker {
+        "DOCKER-USER".to_string()
+    } else {
+        "INPUT".to_string()
+    }
+}

src/main.rs

@@ -22,10 +22,10 @@ async fn main() {
         } => {
             let _ = start_ban_server(ssh_auth_log, iptables_save).await;
         }
-        Arguments::BanPort { port } => ban_port(port),
-        Arguments::UnbanPort { port } => unban_port(port),
-        Arguments::AllowIpPort { ip, port } => allow_ip_port(&ip, port),
-        Arguments::RemoveIpPort { ip, port } => remove_ip_port(&ip, port),
+        Arguments::BanPort { port, docker } => ban_port(port, docker),
+        Arguments::UnbanPort { port, docker } => unban_port(port, docker),
+        Arguments::AllowIpPort { ip, port, docker } => allow_ip_port(&ip, port, docker),
+        Arguments::RemoveIpPort { ip, port, docker } => remove_ip_port(&ip, port, docker),
     }
 }