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
This commit is contained in:
@@ -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
|
||||
Generated
+243
-3
@@ -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",
|
||||
]
|
||||
|
||||
+6
-1
@@ -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"
|
||||
|
||||
+17
-1
@@ -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")]
|
||||
|
||||
+333
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -1 +1,2 @@
|
||||
pub mod ipset;
|
||||
pub mod martillo_maldito;
|
||||
|
||||
+12
@@ -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());
|
||||
|
||||
+83
-39
@@ -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(®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);
|
||||
}
|
||||
}
|
||||
|
||||
Executable
+129
@@ -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."
|
||||
Reference in New Issue
Block a user