diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1f7e6fe..4ff90e3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,7 +11,7 @@ jobs: rust: - stable - nightly - - 1.56.1 # MSRV + - 1.58.1 # MSRV steps: - uses: actions/checkout@v2 diff --git a/Cargo.toml b/Cargo.toml index 1702ea8..2bfff46 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ tracing = "0.1.25" tracing-subscriber = "0.3.1" base64 = "0.13.0" time = { version = "0.3.5", features = ["local-offset"] } +maxminddb = "0.24.0" [build-dependencies] clap = "3.0.0-beta.5" diff --git a/src/cli.rs b/src/cli.rs index 17a09e9..c7ba4c6 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, env, net::IpAddr, str::FromStr}; +use std::{collections::HashMap, env, net::IpAddr, path::PathBuf, str::FromStr}; use clap::{crate_authors, Parser}; @@ -27,6 +27,9 @@ pub struct Args { /// Add an alias for a given public key in the form of 'pubkey:alias' (separate multiple with commas) #[clap(long, short, value_delimiter = ',', multiple_occurrences = true)] pub alias: Vec, + /// Do geoip lookup using Country MMDB from the PATH for 'endpoint_ip' attribute in the 'wireguard_peer_endpoint' metric and add attribute 'endpoint_country' + #[clap(short, long, value_name = "PATH")] + pub geoip_path: Option, } impl Args { diff --git a/src/main.rs b/src/main.rs index 4be1ae6..0b989f5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,7 @@ use color_eyre::{ config::{HookBuilder, Theme}, eyre::Result, }; +use maxminddb; use prometheus::{IntGauge, Registry}; use prometheus_hyper::Server; use tokio::time::{Duration, Instant}; @@ -57,11 +58,15 @@ async fn try_main(args: Args) -> Result<()> { let aliases = args.aliases(); + let maxminddb_reader = args.geoip_path.as_ref().map_or(None, |path| { + Some(unwrap_or_exit!(maxminddb::Reader::open_readfile(path))) + }); + let running = Arc::new(AtomicBool::new(true)); info!("Registering metrics..."); let registry = Arc::new(Registry::new()); - let mut metrics = unwrap_or_exit!(Metrics::new(®istry)); + let mut metrics = unwrap_or_exit!(Metrics::new(®istry, maxminddb_reader)); let scrape_duration = unwrap_or_exit!(IntGauge::new( "wireguard_scrape_duration_milliseconds", "Duration in milliseconds of the scrape", diff --git a/src/metrics.rs b/src/metrics.rs index 24207c6..ff8185f 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -1,4 +1,5 @@ use color_eyre::eyre::Result; +use maxminddb::geoip2; use prometheus::{IntCounterVec, IntGaugeVec, Opts, Registry}; use time::OffsetDateTime; use tracing::{debug, trace}; @@ -12,10 +13,11 @@ pub struct Metrics { bytes_total: IntCounterVec, peer_bytes_total: IntCounterVec, duration_since_latest_handshake: IntGaugeVec, + maxminddb_reader: Option>>, } impl Metrics { - pub fn new(r: &Registry) -> Result { + pub fn new(r: &Registry, maxminddb_reader: Option>>) -> Result { trace!("Metrics::new"); let interfaces_total = IntGaugeVec::new( @@ -33,7 +35,16 @@ impl Metrics { let peer_endpoint = IntGaugeVec::new( Opts::new("wireguard_peer_endpoint", "Peers info. static value"), - &["interface", "endpoint_ip", "peer", "alias"], + match maxminddb_reader { + Some(_) => &[ + "interface", + "endpoint_ip", + "endpoint_country", + "peer", + "alias", + ], + None => &["interface", "endpoint_ip", "peer", "alias"], + }, )?; let bytes_total = IntCounterVec::new( @@ -75,6 +86,7 @@ impl Metrics { peer_endpoint, peer_bytes_total, duration_since_latest_handshake, + maxminddb_reader, }) } @@ -117,8 +129,32 @@ impl Metrics { assert!(p.interface < state.interfaces.len()); if let Some(endpoint) = p.endpoint { - self.peer_endpoint - .with_label_values(&[ + match &self.maxminddb_reader { + Some(reader) => { + let endpoint_country = match reader.lookup::(endpoint.ip()) + { + Ok(reader_response) => { + reader_response.country.map_or_else(String::new, |c| { + c.iso_code + .map(|code| code.to_string()) + .unwrap_or_else(String::new) + }) + } + _ => String::new(), + }; + + self.peer_endpoint.with_label_values(&[ + &state.interfaces[p.interface], + &endpoint.ip().to_string(), + &endpoint_country.to_string(), + &p.pubkey, + &p.alias + .as_ref() + .map(ToOwned::to_owned) + .unwrap_or_else(String::new), + ]) + } + None => self.peer_endpoint.with_label_values(&[ &state.interfaces[p.interface], &endpoint.ip().to_string(), &p.pubkey, @@ -126,8 +162,9 @@ impl Metrics { .as_ref() .map(ToOwned::to_owned) .unwrap_or_else(String::new), - ]) - .set(1); + ]), + } + .set(1); }; let pbtt = self.peer_bytes_total.with_label_values(&[