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 = std::result::Result; impl Ipset { pub fn new(name: impl Into) -> 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", "timeout", "0", "-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) -> 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 { validate_ip(ip)?; let output = Command::new("ipset") .args(["test", &self.name, ip]) .output()?; Ok(output.status.success()) } pub fn list(&self) -> Result> { let output = Command::new("ipset") .args(["list", &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(Vec::new()); } return Err(IpsetError::CommandFailed(stderr)); } let stdout = String::from_utf8_lossy(&output.stdout); let mut in_members = false; let mut ips = Vec::new(); for line in stdout.lines() { if line.starts_with("Members:") { in_members = true; continue; } if in_members && !line.is_empty() { let entry = line.split_whitespace().next().unwrap_or(""); if entry.is_empty() { continue; } if entry.parse::().is_ok() || entry.parse::().is_ok() || entry.contains('/') { ips.push(entry.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::().is_err() && ip.parse::().is_err() { return Err(IpsetError::InvalidIp(ip.to_string())); } Ok(()) } #[cfg(test)] mod tests { use super::*; use serial_test::serial; const TEST_SET: &str = "banned_test_rs"; fn setup() -> Option { 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] #[serial] fn ensure_exists_idempotent() { let Some(set) = setup() else { return }; for _ in 0..5 { set.ensure_exists().expect("ensure_exists failed"); } teardown(&set); } #[test] #[serial] 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] #[serial] 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] #[serial] 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] #[serial] 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] #[serial] 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] #[serial] 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] #[serial] fn list_parses_plain_format_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] #[serial] fn list_empty_set_returns_empty_vec() { let Some(set) = setup() else { return }; let ips = set.list().unwrap(); assert_eq!(ips.len(), 0); teardown(&set); } #[test] #[serial] fn list_filters_non_ip_lines() { let Some(set) = setup() else { return }; set.add("203.0.113.1", None).unwrap(); set.add("203.0.113.2", None).unwrap(); let ips = set.list().unwrap(); for ip in &ips { assert!( ip.parse::().is_ok() || ip.parse::().is_ok() || ip.contains('/'), "got non-IP entry: {ip}" ); } assert_eq!(ips.len(), 2); 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] #[serial] 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] #[serial] 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 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); } #[test] #[serial] 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); } }