add support for ipset

src/cli.rs

@@ -29,6 +29,9 @@ pub(crate) struct Args {
     /// (i.e. everything)
     #[arg(long, short, overrides_with = "verbose", action = ArgAction::Count)]
     pub(crate) quiet: u8,
+    /// Ignore ipset matching this regex
+    #[arg(long, value_name = "REGEX")]
+    pub(crate) ipset_ignore_regex: String,
 }
 
 #[derive(clap::ValueEnum, PartialEq, Eq, Copy, Clone, Debug, strum::AsRefStr, strum::Display)]
@@ -45,4 +48,7 @@ pub(crate) enum ScrapeTarget {
     /// enable 'ip6tables-legacy-save' for metrics
     #[strum(serialize = "ip6tables-legacy")]
     Ip6tablesLegacy,
+    /// enable 'ipset' for metrics
+    #[strum(serialize = "ipset")]
+    Ipset,
 }

src/ipset.rs

@@ -0,0 +1,136 @@
+mod metrics;
+
+pub(crate) use metrics::MetricsIpset;
+
+use std::str::FromStr;
+
+use anyhow::{Context, Result};
+use regex::Regex;
+use tokio::process::Command;
+
+use std::net::Ipv4Addr;
+
+pub(crate) async fn ipset() -> Result<String> {
+    let cmd = format!("ipset");
+
+    String::from_utf8(
+        Command::new(&cmd)
+            .arg("list")
+            .output()
+            .await
+            .with_context(|| format!("Failed to run {cmd}"))?
+            .stdout,
+    )
+    .with_context(|| format!("Failed {cmd} output to valid UTF-8"))
+}
+
+#[derive(Debug)]
+struct IpsetData {
+    name: String,
+    entries: Vec<String>,
+    num_ips: u32,
+}
+
+#[derive(Debug)]
+pub(crate) struct IpsetState {
+    lists: Vec<IpsetData>,
+    ignore_list_regex: Regex,
+}
+
+enum ParserState {
+    OutsideList,
+    InsideList,
+}
+
+fn get_prefix_length(cidr: &str) -> Option<u8> {
+    let parts: Vec<&str> = cidr.split('/').collect();
+    if parts.len() != 2 {
+        return None; // Invalid CIDR format
+    }
+
+    let _ip = Ipv4Addr::from_str(parts[0]).ok()?; // Validate IP part
+    let prefix_length: u8 = parts[1].parse().ok()?; // Parse prefix length
+
+    if prefix_length <= 32 {
+        Some(prefix_length)
+    } else {
+        None // Invalid prefix length
+    }
+}
+
+fn calculate_usable_ip_count(prefix_length: u8) -> u32 {
+    let total_ips = 2u32.pow(32 - prefix_length as u32);
+
+    // Subtract 2 if the prefix length is less than 31 (for network and broadcast)
+    if prefix_length < 31 {
+        total_ips - 2
+    } else {
+        total_ips // /31 and /32 don't have broadcast or network address to subtract
+    }
+}
+
+impl IpsetState {
+    pub(crate) fn new<S: AsRef<str>>(ignore_regex: S) -> Self {
+        let re = Regex::new(ignore_regex.as_ref()).unwrap();
+
+        Self {
+            lists: Vec::new(),
+            ignore_list_regex: re,
+        }
+    }
+
+    pub(crate) fn filter_by_regex(&mut self) {
+        self.lists
+            .retain(|x| !self.ignore_list_regex.is_match(&x.name));
+    }
+
+    pub(crate) async fn parse<S: AsRef<str>>(&mut self, out: S) -> Result<()> {
+        let out = out.as_ref();
+
+        let mut state = ParserState::OutsideList;
+
+        for line in out.lines() {
+            match line {
+                s if s.starts_with("Name:") => {
+                    if let Some((_, right)) = line.split_once(": ") {
+                        self.lists.push(IpsetData {
+                            name: right.to_string(),
+                            entries: Vec::new(),
+                            num_ips: 0,
+                        });
+                    } else {
+                        ()
+                    }
+                }
+                "Members:" => {
+                    state = ParserState::InsideList;
+                }
+                "" => {
+                    state = ParserState::OutsideList;
+                }
+                _ if matches!(state, ParserState::InsideList) => {
+                    if let Some(cur) = self.lists.last_mut() {
+                        cur.entries.push(line.to_string());
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        for list in self.lists.iter_mut() {
+            for entry in list.entries.iter() {
+                let mut num_ips = 1;
+                let mut parsed_entry = entry.as_str();
+                if let Some((front, _)) = entry.split_once(' ') {
+                    parsed_entry = front;
+                }
+                if let Some(pref_len) = get_prefix_length(parsed_entry) {
+                    num_ips = calculate_usable_ip_count(pref_len);
+                }
+
+                list.num_ips = num_ips;
+            }
+        }
+        Ok(())
+    }
+}

src/ipset/metrics.rs

@@ -0,0 +1,76 @@
+use std::collections::HashMap;
+
+use anyhow::Result;
+use prometheus::{IntGaugeVec, Opts, Registry};
+use tracing::trace;
+
+use crate::{cli::ScrapeTarget, ipset::IpsetState};
+
+pub(crate) struct TargetMetricsIpset {
+    entries_total: IntGaugeVec,
+    ips_total: IntGaugeVec,
+}
+
+impl TargetMetricsIpset {
+    fn update(&mut self, state: &IpsetState) {
+        for l in &state.lists {
+            let entries_total_gauge = self.entries_total.with_label_values(&[&l.name]);
+            entries_total_gauge.set(l.entries.len() as i64);
+
+            let ips_total_gauge = self.ips_total.with_label_values(&[&l.name]);
+            ips_total_gauge.set(l.entries.len() as i64);
+        }
+    }
+}
+
+pub(crate) struct MetricsIpset {
+    map: HashMap<String, TargetMetricsIpset>,
+}
+
+impl MetricsIpset {
+    pub(crate) fn new(targets: &[ScrapeTarget], r: &Registry) -> Result<Self> {
+        trace!("MetricsIpset::new");
+
+        let mut map = HashMap::new();
+        for tgt in targets {
+            if !matches!(tgt, ScrapeTarget::Ipset) {
+                continue;
+            }
+            let prefix = String::from("ipset");
+            let entries_total = IntGaugeVec::new(
+                Opts::new(
+                    &format!("{prefix}_entries_total"),
+                    "Total number of entries in the ipset",
+                ),
+                &["list"],
+            )?;
+
+            let ips_total = IntGaugeVec::new(
+                Opts::new(
+                    &format!("{prefix}_ips_total"),
+                    "Total number individual IPs",
+                ),
+                &["list"],
+            )?;
+
+            r.register(Box::new(entries_total.clone()))?;
+            r.register(Box::new(ips_total.clone()))?;
+
+            map.insert(
+                prefix,
+                TargetMetricsIpset {
+                    entries_total,
+                    ips_total,
+                },
+            );
+        }
+
+        Ok(Self { map })
+    }
+
+    pub(crate) fn update(&mut self, tgt: ScrapeTarget, state: &IpsetState) {
+        if let Some(tgt_metrics) = self.map.get_mut(tgt.as_ref()) {
+            tgt_metrics.update(state);
+        }
+    }
+}

src/iptables.rs

@@ -6,7 +6,7 @@ mod table;
 
 pub(crate) use chain::Chain;
 pub(crate) use counter::Counter;
-pub(crate) use metrics::Metrics;
+pub(crate) use metrics::MetricsIptables;
 pub(crate) use rule::Rule;
 pub(crate) use table::Table;
 

src/iptables/metrics.rs

@@ -6,7 +6,7 @@ use tracing::{debug, trace};
 
 use crate::{cli::ScrapeTarget, iptables::IptablesState};
 
-pub(crate) struct TargetMetrics {
+pub(crate) struct TargetMetricsIptables {
     chains_total: IntGaugeVec,
     rules_total: IntGaugeVec,
     chain_bytes_total: IntCounterVec,
@@ -17,7 +17,7 @@ pub(crate) struct TargetMetrics {
     rule_packets_total: IntCounterVec,
 }
 
-impl TargetMetrics {
+impl TargetMetricsIptables {
     fn update(&mut self, state: &IptablesState) {
         for t in &state.tables {
             let ct = self.chains_total.with_label_values(&[&t.name]);
@@ -75,13 +75,13 @@ impl TargetMetrics {
     }
 }
 
-pub(crate) struct Metrics {
-    map: HashMap<String, TargetMetrics>,
+pub(crate) struct MetricsIptables {
+    map: HashMap<String, TargetMetricsIptables>,
 }
 
-impl Metrics {
+impl MetricsIptables {
     pub(crate) fn new(targets: &[ScrapeTarget], r: &Registry) -> Result<Self> {
-        trace!("Metrics::new");
+        trace!("MetricsIptables::new");
 
         let mut map = HashMap::new();
         for tgt in targets {
@@ -162,7 +162,7 @@ impl Metrics {
             r.register(Box::new(chains_total.clone()))?;
             map.insert(
                 tgt.to_string(),
-                TargetMetrics {
+                TargetMetricsIptables {
                     chains_total,
                     rules_total,
                     chain_bytes_total,

src/main.rs

@@ -2,6 +2,7 @@
 mod macros;
 mod cli;
 mod error;
+mod ipset;
 mod iptables;
 mod parse;
 
@@ -21,11 +22,13 @@ use prometheus_hyper::Server;
 use tokio::time::{Duration, Instant};
 use tracing::{debug, info};
 
-use crate::iptables::{iptables_save, IptablesState, Metrics};
+use crate::cli::ScrapeTarget;
+use crate::ipset::{ipset, IpsetState, MetricsIpset};
+use crate::iptables::{iptables_save, IptablesState, MetricsIptables};
 
 #[tokio::main(flavor = "current_thread")]
 async fn main() {
-    let args = cli::Args::parse();
+    let args: cli::Args = cli::Args::parse();
 
     match args.verbose {
         0 => match args.quiet {
@@ -44,7 +47,12 @@ async fn main() {
 
     info!("Registering metrics...");
     let registry = Arc::new(Registry::new());
-    let metrics = Arc::new(Mutex::new(unwrap_or_exit!(Metrics::new(
+    let metrics_iptables = Arc::new(Mutex::new(unwrap_or_exit!(MetricsIptables::new(
+        &args.scrape_targets,
+        &registry
+    ))));
+
+    let metrics_ipset = Arc::new(Mutex::new(unwrap_or_exit!(MetricsIpset::new(
         &args.scrape_targets,
         &registry
     ))));
@@ -86,19 +94,39 @@ async fn main() {
         for tgt in args.scrape_targets.iter().cloned() {
             let scrape_durations = scrape_durations.clone();
             let scrape_successes = scrape_successes.clone();
-            let metrics = metrics.clone();
+            let metrics_iptables = metrics_iptables.clone();
+            let metrics_ipset = metrics_ipset.clone();
+            let ipset_regex = args.ipset_ignore_regex.clone();
             tokio::task::spawn(async move {
                 info!("Collecting {tgt} metrics...");
                 let before = Instant::now();
-                let out = unwrap_or_exit!(iptables_save(tgt).await);
 
-                let mut state = IptablesState::new();
-                unwrap_or_exit!(state.parse(&*out).await);
+                let out: String;
+                match tgt {
+                    ScrapeTarget::Ipset => {
+                        out = unwrap_or_exit!(ipset().await);
+
+                        let mut state = IpsetState::new(ipset_regex);
+                        unwrap_or_exit!(state.parse(&*out).await);
+                        state.filter_by_regex();
+
+                        debug!("Updating {tgt} metrics...");
+                        if let Ok(mut guard) = metrics_ipset.lock() {
+                            guard.update(tgt, &state);
+                        };
+                    }
+                    _ => {
+                        out = unwrap_or_exit!(iptables_save(tgt).await);
+                        let mut state = IptablesState::new();
+                        unwrap_or_exit!(state.parse(&*out).await);
+
+                        debug!("Updating {tgt} metrics...");
+                        if let Ok(mut guard) = metrics_iptables.lock() {
+                            guard.update(tgt, &state);
+                        };
+                    }
+                }
 
-                debug!("Updating {tgt} metrics...");
-                if let Ok(mut guard) = metrics.lock() {
-                    guard.update(tgt, &state);
-                };
                 let after = Instant::now();
 
                 let elapsed = after.duration_since(before);