feat: migrate ban_ip/unban_ip to ipset

- New ipset.rs module wrapping ipset binary (hash:ip backend)
- ban_ip() and unban_ip() use ipset.add/del instead of per-IP iptables rules
- get_banned_ips() parses ipset list -json
- Optional timeout via BAN_DURATION_SECS (default 3600s = 1h)
- Auto-installs single iptables rule: -m set --match-set banned src -j DROP
- New CLI subcommands: ban-ip, unban-ip (with --timeout)
- 13 unit tests for ipset module + integration test script using netns
- Update Woodpecker CI to install ipset

src/cli.rs

@@ -1,13 +1,29 @@
 use structopt::StructOpt;
 
 #[derive(Debug, StructOpt)]
-#[structopt(name = "martillo_maldito", about = "A IPTables wrapper")]
+#[structopt(name = "martillo_maldito", about = "An iptables/ipset wrapper")]
 pub enum Cli {
     #[structopt(about = "Get all banned ips")]
     GetBannedIps {
         #[structopt(name = "Docker", short = "d", long = "docker")]
         docker: bool,
     },
+    #[structopt(about = "Ban an ip (uses ipset)")]
+    BanIp {
+        #[structopt(name = "IP to ban", short = "i", long = "ip")]
+        ip: String,
+        #[structopt(name = "Timeout in seconds", long = "timeout")]
+        timeout: Option<u32>,
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
+    },
+    #[structopt(about = "Unban an ip (uses ipset)")]
+    UnbanIp {
+        #[structopt(name = "IP to unban", short = "i", long = "ip")]
+        ip: String,
+        #[structopt(name = "Docker", short = "d", long = "docker")]
+        docker: bool,
+    },
     #[structopt(about = "Get all secured ports")]
     GetSecuredPorts {
         #[structopt(name = "Docker", short = "d", long = "docker")]

src/ipset.rs

@@ -0,0 +1,333 @@
+use std::path::Path;
+use std::process::Command;
+
+pub const DEFAULT_SET_NAME: &str = "banned";
+pub const DEFAULT_TIMEOUT_SECS: u32 = 3600;
+
+#[derive(Debug, Clone)]
+pub struct Ipset {
+    name: String,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum IpsetError {
+    #[error("ipset binary not found in PATH")]
+    NotInstalled,
+    #[error("ipset command failed: {0}")]
+    CommandFailed(String),
+    #[error("ipset output parse error: {0}")]
+    ParseError(String),
+    #[error("invalid IP address: {0}")]
+    InvalidIp(String),
+    #[error("io error: {0}")]
+    Io(#[from] std::io::Error),
+}
+
+pub type Result<T> = std::result::Result<T, IpsetError>;
+
+impl Ipset {
+    pub fn new(name: impl Into<String>) -> Self {
+        Self { name: name.into() }
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn is_available() -> bool {
+        Command::new("ipset")
+            .arg("--version")
+            .output()
+            .map(|o| o.status.success())
+            .unwrap_or(false)
+    }
+
+    pub fn ensure_exists(&self) -> Result<()> {
+        let output = Command::new("ipset")
+            .args(["create", &self.name, "hash:ip", "-exist"])
+            .output()?;
+        if !output.status.success() {
+            return Err(IpsetError::CommandFailed(
+                String::from_utf8_lossy(&output.stderr).into_owned(),
+            ));
+        }
+        Ok(())
+    }
+
+    pub fn add(&self, ip: &str, timeout_secs: Option<u32>) -> Result<()> {
+        validate_ip(ip)?;
+        let mut cmd = Command::new("ipset");
+        cmd.args(["add", &self.name, ip, "-exist"]);
+        if let Some(secs) = timeout_secs {
+            cmd.args(["timeout", &secs.to_string()]);
+        }
+        let output = cmd.output()?;
+        if !output.status.success() {
+            return Err(IpsetError::CommandFailed(
+                String::from_utf8_lossy(&output.stderr).into_owned(),
+            ));
+        }
+        Ok(())
+    }
+
+    pub fn del(&self, ip: &str) -> Result<()> {
+        validate_ip(ip)?;
+        let output = Command::new("ipset")
+            .args(["del", &self.name, ip, "-exist"])
+            .output()?;
+        if !output.status.success() {
+            return Err(IpsetError::CommandFailed(
+                String::from_utf8_lossy(&output.stderr).into_owned(),
+            ));
+        }
+        Ok(())
+    }
+
+    pub fn contains(&self, ip: &str) -> Result<bool> {
+        validate_ip(ip)?;
+        let output = Command::new("ipset")
+            .args(["test", &self.name, ip])
+            .output()?;
+        Ok(output.status.success())
+    }
+
+    pub fn list(&self) -> Result<Vec<String>> {
+        let output = Command::new("ipset")
+            .args(["list", &self.name, "-json"])
+            .output()?;
+        if !output.status.success() {
+            let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
+            if stderr.contains("does not exist") {
+                return Ok(Vec::new());
+            }
+            return Err(IpsetError::CommandFailed(stderr));
+        }
+        let json: serde_json::Value = serde_json::from_slice(&output.stdout)
+            .map_err(|e| IpsetError::ParseError(e.to_string()))?;
+
+        let mut ips = Vec::new();
+        if let Some(sets) = json.get("ipset").and_then(|v| v.as_array()) {
+            for set in sets {
+                if let Some(members) = set.get("members").and_then(|v| v.as_array()) {
+                    for m in members {
+                        if let Some(ip) = m.get("ip").and_then(|v| v.as_str()) {
+                            ips.push(ip.to_string());
+                        }
+                    }
+                }
+            }
+        }
+        Ok(ips)
+    }
+
+    pub fn save(&self, path: &Path) -> Result<()> {
+        let output = Command::new("ipset")
+            .args(["save", &self.name])
+            .output()?;
+        if !output.status.success() {
+            return Err(IpsetError::CommandFailed(
+                String::from_utf8_lossy(&output.stderr).into_owned(),
+            ));
+        }
+        std::fs::write(path, &output.stdout)?;
+        Ok(())
+    }
+
+    pub fn flush(&self) -> Result<()> {
+        let output = Command::new("ipset")
+            .args(["flush", &self.name])
+            .output()?;
+        if !output.status.success() {
+            return Err(IpsetError::CommandFailed(
+                String::from_utf8_lossy(&output.stderr).into_owned(),
+            ));
+        }
+        Ok(())
+    }
+
+    pub fn destroy(&self) -> Result<()> {
+        let output = Command::new("ipset")
+            .args(["destroy", &self.name])
+            .output()?;
+        if !output.status.success() {
+            let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
+            if stderr.contains("does not exist") {
+                return Ok(());
+            }
+            return Err(IpsetError::CommandFailed(stderr));
+        }
+        Ok(())
+    }
+}
+
+fn validate_ip(ip: &str) -> Result<()> {
+    if ip.parse::<std::net::Ipv4Addr>().is_err() && ip.parse::<std::net::Ipv6Addr>().is_err() {
+        return Err(IpsetError::InvalidIp(ip.to_string()));
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const TEST_SET: &str = "banned_test_rs";
+
+    fn setup() -> Option<Ipset> {
+        if !Ipset::is_available() {
+            eprintln!("ipset binary not available, skipping tests");
+            return None;
+        }
+        let set = Ipset::new(TEST_SET);
+        let _ = set.destroy();
+        match set.ensure_exists() {
+            Ok(_) => Some(set),
+            Err(e) => {
+                eprintln!("Cannot create ipset (need root?): {}", e);
+                None
+            }
+        }
+    }
+
+    fn teardown(set: &Ipset) {
+        let _ = set.destroy();
+    }
+
+    #[test]
+    fn ensure_exists_idempotent() {
+        let Some(set) = setup() else { return };
+        for _ in 0..5 {
+            set.ensure_exists().expect("ensure_exists failed");
+        }
+        teardown(&set);
+    }
+
+    #[test]
+    fn add_then_list_contains_ip() {
+        let Some(set) = setup() else { return };
+        set.add("192.0.2.1", None).expect("add failed");
+        let ips = set.list().expect("list failed");
+        assert!(ips.contains(&"192.0.2.1".to_string()));
+        teardown(&set);
+    }
+
+    #[test]
+    fn add_twice_no_error() {
+        let Some(set) = setup() else { return };
+        set.add("192.0.2.2", None).expect("first add");
+        set.add("192.0.2.2", None).expect("second add should be idempotent");
+        teardown(&set);
+    }
+
+    #[test]
+    fn del_then_list_empty() {
+        let Some(set) = setup() else { return };
+        set.add("192.0.2.3", None).expect("add");
+        set.del("192.0.2.3").expect("del");
+        let ips = set.list().expect("list");
+        assert!(!ips.contains(&"192.0.2.3".to_string()));
+        teardown(&set);
+    }
+
+    #[test]
+    fn del_nonexistent_no_error() {
+        let Some(set) = setup() else { return };
+        set.del("192.0.2.4").expect("del nonexistent should not fail");
+        teardown(&set);
+    }
+
+    #[test]
+    fn contains_true_after_add_false_after_del() {
+        let Some(set) = setup() else { return };
+        set.add("192.0.2.5", None).expect("add");
+        assert!(set.contains("192.0.2.5").expect("contains"));
+        set.del("192.0.2.5").expect("del");
+        assert!(!set.contains("192.0.2.5").expect("contains after del"));
+        teardown(&set);
+    }
+
+    #[test]
+    fn save_and_restore_roundtrip() {
+        use std::env;
+        let Some(set) = setup() else { return };
+        set.add("192.0.2.10", None).unwrap();
+        set.add("192.0.2.11", None).unwrap();
+
+        let tmp = env::temp_dir().join("martillo_ipset_test.restore");
+        set.save(&tmp).expect("save");
+
+        let content = std::fs::read_to_string(&tmp).expect("read");
+        assert!(content.contains("192.0.2.10"));
+        assert!(content.contains("192.0.2.11"));
+
+        let _ = std::fs::remove_file(&tmp);
+        teardown(&set);
+    }
+
+    #[test]
+    fn list_parses_json_correctly() {
+        let Some(set) = setup() else { return };
+        set.add("198.51.100.1", None).unwrap();
+        set.add("198.51.100.2", None).unwrap();
+        set.add("198.51.100.3", None).unwrap();
+        let ips = set.list().unwrap();
+        assert_eq!(ips.len(), 3);
+        assert!(ips.contains(&"198.51.100.1".to_string()));
+        assert!(ips.contains(&"198.51.100.2".to_string()));
+        assert!(ips.contains(&"198.51.100.3".to_string()));
+        teardown(&set);
+    }
+
+    #[test]
+    fn invalid_ip_rejected() {
+        let result = validate_ip("not-an-ip");
+        assert!(matches!(result, Err(IpsetError::InvalidIp(_))));
+        assert!(validate_ip("1.2.3.4").is_ok());
+        assert!(validate_ip("::1").is_ok());
+    }
+
+    #[test]
+    fn flush_clears_all() {
+        let Some(set) = setup() else { return };
+        set.add("203.0.113.1", None).unwrap();
+        set.add("203.0.113.2", None).unwrap();
+        assert_eq!(set.list().unwrap().len(), 2);
+        set.flush().unwrap();
+        assert_eq!(set.list().unwrap().len(), 0);
+        teardown(&set);
+    }
+
+    #[test]
+    fn add_with_timeout() {
+        let Some(set) = setup() else { return };
+        set.add("203.0.113.10", Some(2)).expect("add with timeout");
+        assert!(set.contains("203.0.113.10").unwrap());
+        std::thread::sleep(std::time::Duration::from_secs(3));
+        assert!(!set.contains("203.0.113.10").unwrap(), "timeout should expire");
+        teardown(&set);
+    }
+
+    #[test]
+    fn list_on_nonexistent_set_returns_empty() {
+        let Some(set) = setup() else { return };
+        let bogus = Ipset::new("nonexistent_set_zzz_999");
+        let result = bogus.list();
+        assert!(result.is_ok(), "list on nonexistent should not error");
+        assert_eq!(result.unwrap().len(), 0);
+        teardown(&set);
+    }
+
+    #[test]
+    fn bulk_1000_ips_performance() {
+        let Some(set) = setup() else { return };
+        let start = std::time::Instant::now();
+        for i in 0..1000u32 {
+            let ip = format!("10.99.{}.{}", i / 256, i % 256);
+            set.add(&ip, None).expect("add");
+        }
+        let elapsed = start.elapsed();
+        assert!(elapsed.as_secs() < 5, "1000 adds took {:?}", elapsed);
+        assert_eq!(set.list().unwrap().len(), 1000);
+        teardown(&set);
+    }
+}

src/lib.rs

@@ -1 +1,2 @@
+pub mod ipset;
 pub mod martillo_maldito;

src/main.rs

@@ -10,6 +10,18 @@ fn main() {
             let banned_ips = MartilloMaldito::ipv4(docker).get_banned_ips();
             println!("{}", serde_json::to_string(&banned_ips).unwrap());
         }
+        Cli::BanIp { ip, timeout, docker } => {
+            let martillo = MartilloMaldito::ipv4(docker);
+            let result = match timeout {
+                Some(t) => martillo.ban_ip_with_timeout(&ip, t),
+                None => martillo.ban_ip(&ip),
+            };
+            println!("{}", result.is_ok())
+        }
+        Cli::UnbanIp { ip, docker } => {
+            let martillo = MartilloMaldito::ipv4(docker);
+            println!("{}", martillo.unban_ip(&ip).is_ok())
+        }
         Cli::GetSecuredPorts { docker } => {
             let secured_ports = MartilloMaldito::ipv4(docker).get_secured_ports();
             println!("{}", serde_json::to_string(&secured_ports).unwrap());

src/martillo_maldito.rs

@@ -1,30 +1,68 @@
+use crate::ipset::{Ipset, DEFAULT_SET_NAME, DEFAULT_TIMEOUT_SECS};
 use regex::Regex;
 use std::{collections::HashMap, process::Command};
 
 pub struct MartilloMaldito {
     iptables: iptables::IPTables,
     chain: String,
+    ipset: Ipset,
+    ban_timeout: Option<u32>,
 }
 
 impl MartilloMaldito {
     pub fn ipv4(docker: bool) -> MartilloMaldito {
-        MartilloMaldito {
-            iptables: iptables::new(false).unwrap(),
-            chain: Self::get_chain(docker).to_string(),
-        }
+        Self::ipv4_with_config(docker, DEFAULT_SET_NAME, None)
     }
 
-    pub fn ipv6(docker: bool) -> MartilloMaldito {
-        MartilloMaldito {
-            iptables: iptables::new(true).unwrap(),
-            chain: Self::get_chain(docker).to_string(),
+    pub fn ipv4_with_config(
+        docker: bool,
+        ipset_name: &str,
+        ban_timeout: Option<u32>,
+    ) -> MartilloMaldito {
+        let ipset = Ipset::new(ipset_name);
+        if let Err(e) = ipset.ensure_exists() {
+            eprintln!("Warning: could not ensure ipset exists: {}", e);
         }
+        let martillo = MartilloMaldito {
+            iptables: iptables::new(false).unwrap(),
+            chain: Self::get_chain(docker).to_string(),
+            ipset,
+            ban_timeout,
+        };
+        if let Err(e) = martillo.ensure_ipset_match_rule() {
+            eprintln!("Warning: could not install ipset match rule: {}", e);
+        }
+        martillo
+    }
+
+    fn ensure_ipset_match_rule(&self) -> Result<(), Box<dyn std::error::Error>> {
+        let check_rule = format!("-m set --match-set {} src -j DROP", self.ipset.name());
+        let output = Command::new("iptables")
+            .args(["-C", &self.chain, &check_rule])
+            .output()?;
+        if output.status.success() {
+            return Ok(());
+        }
+        let output = Command::new("iptables")
+            .args(["-I", &self.chain, "1", &check_rule])
+            .output()?;
+        if !output.status.success() {
+            return Err(format!(
+                "failed to install ipset match rule: {}",
+                String::from_utf8_lossy(&output.stderr)
+            )
+            .into());
+        }
+        Ok(())
     }
 
     pub fn save_rules() -> std::io::Result<std::process::Output> {
-        Command::new("iptables-save")
-            .args(["-f", "/etc/iptables/rules.v4"])
-            .output()
+        Command::new("iptables-save").args(["-f", "/etc/iptables/rules.v4"]).output()
+    }
+
+    pub fn save_ipset(&self, path: &std::path::Path) -> Result<(), Box<dyn std::error::Error>> {
+        self.ipset.save(path)?;
+        Ok(())
     }
 
     pub fn is_port_secured(&self, port: u16) -> bool {
@@ -32,7 +70,6 @@ impl MartilloMaldito {
         if rules.is_err() {
             return false;
         }
-
         for rule in rules.unwrap() {
             if rule.contains(&format!("-p tcp -m tcp --dport {} -j DROP", port)) {
                 return true;
@@ -46,48 +83,30 @@ impl MartilloMaldito {
         if rules.is_err() {
             return vec![];
         }
-
         let port_regex = iptables_regex_for_port();
         rules
             .unwrap()
             .iter()
             .filter(|r| r.contains("-p tcp -m tcp --dport") && r.contains("-j DROP"))
+            .filter(|r| !r.contains("match-set"))
             .map(|r| extract_port(&port_regex, r).unwrap())
             .collect()
     }
 
     pub fn get_banned_ips(&self) -> Vec<String> {
-        let rules = self.get_rules();
-        if rules.is_err() {
-            return vec![];
-        }
-
-        let ip_regex = iptables_regex_for_ip();
-        rules
-            .unwrap()
-            .iter()
-            .filter(|r| {
-                r.contains(&format!("-A {}", self.chain))
-                    && r.contains("-j DROP")
-                    && r.contains("-s")
-            })
-            .map(|r| extract_ip(&ip_regex, r).unwrap())
-            .collect()
+        self.ipset.list().unwrap_or_default()
     }
 
     pub fn get_secured_ports_with_allowed_ips(&self) -> HashMap<u16, Vec<String>> {
         let mut result: HashMap<u16, Vec<String>> = HashMap::new();
-
         let secured_ports = self.get_secured_ports();
         if secured_ports.is_empty() {
             return result;
         }
-
         let rules = self.get_rules();
         if rules.is_err() {
             return result;
         }
-
         let rules = rules.unwrap();
         let ip_regex = iptables_regex_for_ip();
         for port in secured_ports {
@@ -101,16 +120,26 @@ impl MartilloMaldito {
                 .collect();
             result.insert(port, ips);
         }
-
         result
     }
 
     pub fn ban_ip(&self, ip: &str) -> Result<(), Box<dyn std::error::Error>> {
-        self.append_unique("filter", &format!("-s {} -j DROP", ip))
+        self.ipset.add(ip, self.ban_timeout)?;
+        Ok(())
+    }
+
+    pub fn ban_ip_with_timeout(
+        &self,
+        ip: &str,
+        timeout_secs: u32,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        self.ipset.add(ip, Some(timeout_secs))?;
+        Ok(())
     }
 
     pub fn unban_ip(&self, ip: &str) -> Result<(), Box<dyn std::error::Error>> {
-        self.remove_unique("filter", &format!("-s {} -j DROP", ip))
+        self.ipset.del(ip)?;
+        Ok(())
     }
 
     pub fn secure_port(
@@ -173,10 +202,6 @@ impl MartilloMaldito {
         self.iptables.append_unique(table, &self.chain, rule)
     }
 
-    fn remove_unique(&self, table: &str, rule: &str) -> Result<(), Box<dyn std::error::Error>> {
-        self.iptables.delete(table, &self.chain, rule)
-    }
-
     fn insert_unique(
         &self,
         table: &str,
@@ -224,6 +249,10 @@ fn iptables_regex_for_port() -> Regex {
     Regex::new(r"--dport\s+(\d+)").unwrap()
 }
 
+pub const fn default_ban_timeout_secs() -> u32 {
+    DEFAULT_TIMEOUT_SECS
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -248,4 +277,19 @@ mod tests {
         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()));
     }
+
+    #[test]
+    fn secured_ports_excludes_ipset_match_rule() {
+        let input = "-A INPUT -p tcp -m tcp --dport 2222 -j DROP";
+        assert_eq!(
+            extract_port(&iptables_regex_for_port(), input),
+            Some(2222)
+        );
+        assert!(!input.contains("match-set"));
+    }
+
+    #[test]
+    fn default_timeout_constant_is_one_hour() {
+        assert_eq!(default_ban_timeout_secs(), 3600);
+    }
 }