diff --git a/src/ipset.rs b/src/ipset.rs index 9d72c7d..56bffc3 100644 --- a/src/ipset.rs +++ b/src/ipset.rs @@ -1,40 +1,74 @@ use rust_util::{opt_result, simple_error, XResult}; +use std::error::Error; use std::ffi::OsStr; use std::process::Command; const CMD_IPSET: &str = "ipset"; -// TODO -// list names -// list ipsetname -// add -// del - -struct IpSet { +pub struct IpSet { name: String, } impl IpSet { - pub fn list_names() -> Vec { - vec![] + pub fn new(name: &str) -> XResult { + if name.is_empty() { + return simple_error!("ipset name cannot be empty"); + } + Ok(IpSet { + name: String::from(name), + }) } - pub fn list(&self) -> Vec { - vec![] + pub fn list_names() -> XResult> { + list_names() } - - pub fn add(&self, ip: &str) -> () {} - - pub fn del(&self, ip: &str) -> () {} } -fn list_names() {} +impl IpSet { + pub fn list(&self) -> XResult> { + let (stdout, _stderr) = execute_ipset(&["list", &self.name])?; + let ipset_list_ips = opt_result!( + String::from_utf8(stdout.clone()), + "Parse output: {} failed: {}", + String::from_utf8_lossy(&stdout) + ); + Ok(parse_ipset_list_ips(&ipset_list_ips)) + } + + pub fn add(&self, ip: &str) -> XResult<()> { + add_ipset(&self.name, ip) + } + + pub fn del(&self, ip: &str) -> XResult<()> { + del_ipset(&self.name, ip) + } +} + +fn list_names() -> XResult> { + let (stdout, _stderr) = execute_ipset(&["list", "-n"])?; + let ipset_list_names = opt_result!( + String::from_utf8(stdout.clone()), + "Parse output: {} failed: {}", + String::from_utf8_lossy(&stdout) + ); + Ok(parse_ipset_list(&ipset_list_names)) +} fn list_ipset(ipset_name: &str) {} -fn add_ipset(ipset_name: &str, ip: &str) {} +fn add_ipset(ipset_name: &str, ip: &str) -> XResult<()> { + if let Err(err) = execute_ipset(&["add", ipset_name, ip]) { + return check_result(err, &["cannot be added to the set: it's already added"]); + } + Ok(()) +} -fn del_ipset(ipset_name: &str, ip: &str) {} +fn del_ipset(ipset_name: &str, ip: &str) -> XResult<()> { + if let Err(err) = execute_ipset(&["del", ipset_name, ip]) { + return check_result(err, &["cannot be deleted from the set: it's not added"]); + } + Ok(()) +} fn execute_ipset(args: I) -> XResult<(Vec, Vec)> where @@ -46,9 +80,100 @@ where let cmd_output = opt_result!(cmd_ipset.output(), "Execute ipset with failed: {}"); if !cmd_output.status.success() { return simple_error!( - "Execute ipset not failed, exit code: {:?}", - cmd_output.status.code() + "Execute ipset not failed, exit code: {:?}, stdout: {}, stderr: {}", + cmd_output.status.code(), + String::from_utf8_lossy(&cmd_output.stdout), + String::from_utf8_lossy(&cmd_output.stderr) ); } Ok((cmd_output.stdout, cmd_output.stderr)) } + +fn check_result(e: Box, allowed_messages: &[&str]) -> XResult<()> { + let e_msg = e.to_string(); + for allowed_message in allowed_messages { + if e_msg.contains(allowed_message) { + return Ok(()); + } + } + Err(e) +} + +fn parse_ipset_list(ipset_list: &str) -> Vec { + ipset_list + .split('\n') + .map(|name| name.trim()) + .filter(|name| !name.is_empty()) + .map(ToString::to_string) + .collect::>() +} + +fn parse_ipset_list_ips(ipset_list: &str) -> Vec { + let mut result = vec![]; + let mut is_after_members = false; + ipset_list.split('\n').for_each(|ln| { + let ln = ln.trim(); + if !ln.is_empty() { + if is_after_members { + result.push(ln.to_string()); + } else { + if ln.to_lowercase() == "members:" { + is_after_members = true; + } + } + } + }); + return result; +} + +#[test] +fn test_parse_ipset_list() { + let empty_vec: Vec = vec![]; + assert_eq!(empty_vec, parse_ipset_list("")); + assert_eq!(vec!["test".to_string()], parse_ipset_list("test\n")); + assert_eq!( + vec!["allowipset".to_string()], + parse_ipset_list("allowipset") + ); + assert_eq!( + vec!["allowipset".to_string()], + parse_ipset_list("allowipset\n") + ); +} + +#[test] +fn test_parse_ipset_list_ips_1() { + let ipset_list_ips = r"Name: allowipset +Type: hash:ip +Revision: 1 +Header: family inet hashsize 1024 maxelem 65536 +Size in memory: 21264 +References: 6 +Members: +115.195.128.60 +36.18.233.45 +36.27.0.166 +60.186.198.141 +223.104.246.130"; + let ips = parse_ipset_list_ips(ipset_list_ips); + assert_eq!(5, ips.len()); + assert_eq!("115.195.128.60", &ips[0]); + assert_eq!("36.18.233.45", &ips[1]); + assert_eq!("36.27.0.166", &ips[2]); + assert_eq!("60.186.198.141", &ips[3]); + assert_eq!("223.104.246.130", &ips[4]); +} + +#[test] +fn test_parse_ipset_list_ips_2() { + let ipset_list_ips = r"Size in memory: 30688 +References: 6 +Members: +36.28.155.64 +36.20.56.144 +"; + let ips = parse_ipset_list_ips(ipset_list_ips); + assert_eq!(2, ips.len()); + assert_eq!("36.28.155.64", &ips[0]); + assert_eq!("36.20.56.144", &ips[1]); +}