diff --git a/.woodpecker.yml b/.woodpecker.yml new file mode 100644 index 0000000..2b88f26 --- /dev/null +++ b/.woodpecker.yml @@ -0,0 +1,11 @@ +when: + - event: push + branch: main + - event: pull_request + +steps: + - name: test + image: rustlang/rust:nightly + commands: + - apt-get update && apt-get install -y ipset iptables + - cargo test --lib \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index a32183a..759c7fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -70,6 +70,68 @@ dependencies = [ "vec_map", ] +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + [[package]] name = "heck" version = "0.3.3" @@ -113,19 +175,43 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.169" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "martillo_maldito" -version = "0.1.2" +version = "0.2.0" dependencies = [ "iptables", "regex", "serde", "serde_json", + "serial_test", "structopt", + "tempfile", + "thiserror", ] [[package]] @@ -146,6 +232,41 @@ dependencies = [ "libc", ] +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "proc-macro-error" version = "1.0.4" @@ -188,6 +309,21 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags 2.6.0", +] + [[package]] name = "regex" version = "1.11.1" @@ -217,12 +353,31 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags 2.6.0", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + [[package]] name = "ryu" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "serde" version = "1.0.217" @@ -255,6 +410,43 @@ dependencies = [ "serde", ] +[[package]] +name = "serial_test" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "699f4197115b8a7e7ff19c9a315a4bd6fffec26cc4626ef45ecaea389e081c6d" +dependencies = [ + "futures-executor", + "futures-util", + "log", + "once_cell", + "parking_lot", + "serial_test_derive", +] + +[[package]] +name = "serial_test_derive" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e153fc76e1c6a068703d6d29c508a0b15c061c4b7e43da59cc097bc342673c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.93", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + [[package]] name = "strsim" version = "0.8.0" @@ -307,6 +499,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + [[package]] name = "textwrap" version = "0.11.0" @@ -316,6 +521,26 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.93", +] + [[package]] name = "unicode-ident" version = "1.0.14" @@ -367,3 +592,18 @@ name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] diff --git a/Cargo.toml b/Cargo.toml index 5212634..2f8c8c2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "martillo_maldito" -version = "0.1.2" +version = "0.2.0" edition = "2021" [dependencies] @@ -11,6 +11,11 @@ serde = {version = "1.0.217", features = ["derive"]} serde_json = "1.0.134" regex = "1.11.1" +thiserror = "1.0" + +[dev-dependencies] +serial_test = "3.2.0" +tempfile = "3.13" [lib] name = "martillo_maldito" diff --git a/src/cli.rs b/src/cli.rs index 52776f5..ddf2c53 100644 --- a/src/cli.rs +++ b/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, + #[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")] diff --git a/src/ipset.rs b/src/ipset.rs new file mode 100644 index 0000000..36d9701 --- /dev/null +++ b/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 = 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", "-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, "-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::().is_err() && ip.parse::().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 { + 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); + } +} diff --git a/src/lib.rs b/src/lib.rs index 2316816..34b953f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1 +1,2 @@ +pub mod ipset; pub mod martillo_maldito; diff --git a/src/main.rs b/src/main.rs index 8a2595d..9e58ee6 100644 --- a/src/main.rs +++ b/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()); diff --git a/src/martillo_maldito.rs b/src/martillo_maldito.rs index a98a4c3..8fede8f 100644 --- a/src/martillo_maldito.rs +++ b/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, } 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, + ) -> 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> { + 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 { - 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> { + 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 { - 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> { let mut result: HashMap> = 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> { - 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> { + self.ipset.add(ip, Some(timeout_secs))?; + Ok(()) } pub fn unban_ip(&self, ip: &str) -> Result<(), Box> { - 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> { - 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(®ex, 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); + } } diff --git a/tests/integration_netns.sh b/tests/integration_netns.sh new file mode 100755 index 0000000..84fdaab --- /dev/null +++ b/tests/integration_netns.sh @@ -0,0 +1,129 @@ +#!/bin/bash +# Integration tests for martillo-maldito CLI in isolated network namespaces. +# These exercise the full iptables+ipset stack end-to-end. +# +# Requires: ipset, iptables, sudo, cargo (already built binary at $BIN) +# Usage: BIN=./target/release/martillo_maldito ./tests/integration_netns.sh + +set -euo pipefail + +BIN="${BIN:-./target/release/martillo_maldito}" +SET="${IPSET_NAME:-banned_it_test}" +NS="martillo-it-$RANDOM" + +if [[ ! -x "$BIN" ]]; then + echo "ERROR: binary not found at $BIN" + echo "Build it first: cargo build --release" + exit 1 +fi + +if ! command -v ipset >/dev/null; then + echo "ERROR: ipset not installed" + exit 1 +fi + +if ! command -v sudo >/dev/null; then + echo "ERROR: sudo not installed (required for netns + iptables)" + exit 1 +fi + +cleanup() { + sudo ip netns del "$NS" 2>/dev/null || true + sudo ipset destroy "$SET" 2>/dev/null || true +} +trap cleanup EXIT + +echo "==> Creating namespace $NS" +sudo ip netns add "$NS" + +echo "==> Cleaning any pre-existing set" +sudo ipset destroy "$SET" 2>/dev/null || true + +PASS=0 +FAIL=0 + +assert_eq() { + local desc="$1" expected="$2" actual="$3" + if [[ "$expected" == "$actual" ]]; then + echo " PASS: $desc" + PASS=$((PASS+1)) + else + echo " FAIL: $desc (expected: '$expected', got: '$actual')" + FAIL=$((FAIL+1)) + fi +} + +run_cli() { + sudo ip netns exec "$NS" "$BIN" "$@" +} + +echo "" +echo "==> Test 1: ban then unban cycle" +run_cli ban-ip -i 192.0.2.10 +assert_eq "ipset contains banned IP" "0" "$(sudo ip netns exec "$NS" ipset test "$SET" 192.0.2.10 >/dev/null 2>&1; echo $?)" +run_cli unban-ip -i 192.0.2.10 +assert_eq "ipset no longer contains IP" "1" "$(sudo ip netns exec "$NS" ipset test "$SET" 192.0.2.10 >/dev/null 2>&1; echo $?)" + +echo "" +echo "==> Test 2: get-banned-ips returns JSON array" +sudo ip netns exec "$NS" "$BIN" ban-ip -i 192.0.2.20 +sudo ip netns exec "$NS" "$BIN" ban-ip -i 192.0.2.21 +output=$(sudo ip netns exec "$NS" "$BIN" get-banned-ips) +assert_eq "output is valid JSON" "true" "$(echo "$output" | jq -e '. | type == "array"' >/dev/null 2>&1 && echo true || echo false)" +count=$(echo "$output" | jq 'length') +assert_eq "contains 2 IPs" "2" "$count" + +echo "" +echo "==> Test 3: unban is idempotent" +run_cli unban-ip -i 192.0.2.99 +output=$(run_cli unban-ip -i 192.0.2.99) +assert_eq "second unban returns ok" "true" "$output" + +echo "" +echo "==> Test 4: iptables rule installed for ipset match" +rule_check=$(sudo ip netns exec "$NS" iptables -C INPUT -m set --match-set "$SET" src -j DROP 2>&1; echo $?) +assert_eq "iptables rule exists" "0" "$rule_check" + +echo "" +echo "==> Test 5: secured ports with allowed IPs" +run_cli secure-port -p 9999 +assert_eq "port 9999 is secured" "true" "$(run_cli is-port-secured -p 9999)" +assert_eq "port 8888 is NOT secured" "false" "$(run_cli is-port-secured -p 8888)" +run_cli allow-ip-for-port -i 10.0.0.5 -p 9999 +allowed=$(run_cli get-secured-ports-with-allowed-ips) +assert_eq "10.0.0.5 allowed for 9999" "10.0.0.5" "$(echo "$allowed" | jq -r '."9999"[]')" +run_cli unsecure-port -p 9999 +run_cli remove-allow-ip-port -i 10.0.0.5 -p 9999 + +echo "" +echo "==> Test 6: bulk ban performance (1000 IPs)" +start=$(date +%s%N) +for i in $(seq 1 1000); do + a=$((i / 256)) + b=$((i % 256)) + sudo ip netns exec "$NS" ipset add "$SET" "10.50.$a.$b" -exist >/dev/null +done +end=$(date +%s%N) +elapsed_ms=$(( (end - start) / 1000000 )) +assert_eq "1000 ipset adds under 5s" "true" "$([[ $elapsed_ms -lt 5000 ]] && echo true || echo false)" +echo " (1000 adds took ${elapsed_ms}ms)" + +echo "" +echo "==> Test 7: cli list matches ipset list directly" +cli_list=$(sudo ip netns exec "$NS" "$BIN" get-banned-ips | jq -r '.[]' | sort) +ipset_list=$(sudo ip netns exec "$NS" ipset list "$SET" -json | jq -r '.ipset[0].members[].ip' | sort) +# ipset -json uses different format, compare counts instead +cli_count=$(echo "$cli_list" | wc -l) +ipset_count=$(echo "$ipset_list" | wc -l) +assert_eq "cli and ipset report same count" "$ipset_count" "$cli_count" + +echo "" +echo "================================================" +echo "Results: $PASS passed, $FAIL failed" +echo "================================================" + +if [[ $FAIL -gt 0 ]]; then + exit 1 +fi + +echo "All integration tests passed." \ No newline at end of file