forked from midefos/martillo-maldito
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:
1 parent
9af186eea3
commit
03a52925d4
9 files changed
+835
-44
No files matched your search
+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);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user