diff --git a/src/cli.rs b/src/cli.rs index d65a985..1c86130 100644 --- a/src/cli.rs +++ b/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, } diff --git a/src/ipset.rs b/src/ipset.rs new file mode 100644 index 0000000..4c16edb --- /dev/null +++ b/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 { + 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, + num_ips: u32, +} + +#[derive(Debug)] +pub(crate) struct IpsetState { + lists: Vec, + ignore_list_regex: Regex, +} + +enum ParserState { + OutsideList, + InsideList, +} + +fn get_prefix_length(cidr: &str) -> Option { + 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>(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>(&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(()) + } +} diff --git a/src/ipset/metrics.rs b/src/ipset/metrics.rs new file mode 100644 index 0000000..8a9adb7 --- /dev/null +++ b/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, +} + +impl MetricsIpset { + pub(crate) fn new(targets: &[ScrapeTarget], r: &Registry) -> Result { + 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); + } + } +} diff --git a/src/iptables.rs b/src/iptables.rs index 6ab78d5..e55ca3d 100644 --- a/src/iptables.rs +++ b/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; diff --git a/src/iptables/metrics.rs b/src/iptables/metrics.rs index c3e15a5..ecefc7b 100644 --- a/src/iptables/metrics.rs +++ b/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, +pub(crate) struct MetricsIptables { + map: HashMap, } -impl Metrics { +impl MetricsIptables { pub(crate) fn new(targets: &[ScrapeTarget], r: &Registry) -> Result { - 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, diff --git a/src/main.rs b/src/main.rs index fe4fc0f..5fb88b6 100644 --- a/src/main.rs +++ b/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, + ®istry + )))); + + let metrics_ipset = Arc::new(Mutex::new(unwrap_or_exit!(MetricsIpset::new( &args.scrape_targets, ®istry )))); @@ -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);