diff --git a/src/cli.rs b/src/cli.rs index 5cbf7cd..ce2f448 100644 --- a/src/cli.rs +++ b/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, }, } diff --git a/src/iptables_wrapper.rs b/src/iptables_wrapper.rs index f8974ec..fe72103 100644 --- a/src/iptables_wrapper.rs +++ b/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() + } +} diff --git a/src/main.rs b/src/main.rs index bf0aa39..f3e9f83 100644 --- a/src/main.rs +++ b/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), } }