Merge pull request #1 from kbknapp/peer-alias

Adds peer aliases

Cargo.toml

@@ -31,6 +31,7 @@ prometheus-hyper = "0.1.3"
 tokio = { version = "1.5.0", features = ["full"] }
 tracing = "0.1.25"
 tracing-subscriber = "0.3.1"
+base64 = "0.13.0"
 
 [build-dependencies]
 clap = "3.0.0-beta.5"

src/cli.rs

@@ -1,4 +1,4 @@
-use std::{env, net::IpAddr};
+use std::{collections::HashMap, env, net::IpAddr, str::FromStr};
 
 use clap::{crate_authors, Parser};
 
@@ -8,20 +8,81 @@ static AUTHORS: &str = crate_authors!();
 /// A Prometheus exporter for WireGuard
 #[derive(Parser)]
 #[clap(author = AUTHORS, version = VERSION)]
-pub(crate) struct Args {
+pub struct Args {
     /// How often metrics are gathered
     #[clap(long, default_value = "5", value_name = "SECS")]
-    pub(crate) collect_interval: u64,
+    pub collect_interval: u64,
     /// The listen port for scraping metrics
     #[clap(short = 'p', long, default_value = "9586", value_name = "PORT")]
-    pub(crate) listen_port: u16,
+    pub listen_port: u16,
     /// The listen address scraping metrics
     #[clap(short, long, default_value = "0.0.0.0", value_name = "ADDR")]
-    pub(crate) listen_address: IpAddr,
+    pub listen_address: IpAddr,
     /// Show verbose output at a level or higher. -v:  DEBUG, -vv: TRACE
     #[clap(long, short, parse(from_occurrences))]
-    pub(crate) verbose: u8,
+    pub verbose: u8,
     /// Supress output at a level or lower. -q: INFO, -qq: WARN, -qqq: ERROR (i.e. everything)
     #[clap(long, short, overrides_with = "verbose", parse(from_occurrences))]
-    pub(crate) quiet: u8,
+    pub quiet: u8,
+    /// Add an alias for a given public key in the form of 'alias:pubkey' (separate multiple with commas)
+    #[clap(long, short, value_delimiter = ',', multiple_occurrences = true)]
+    pub alias: Vec<Alias>,
+}
+
+impl Args {
+    pub fn aliases(&self) -> HashMap<&str, &str> {
+        let mut map = HashMap::new();
+        for alias in &self.alias {
+            let Alias {
+                inner: (pubkey, alias),
+            } = alias;
+            map.insert(pubkey.as_ref(), alias.as_ref());
+        }
+
+        map
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Alias {
+    // A base64 encoded public key and a human readable alias
+    // (pubkey, alias)
+    pub inner: (String, String),
+}
+
+impl FromStr for Alias {
+    type Err = String;
+    fn from_str(s: &str) -> Result<Alias, Self::Err> {
+        let mut parts = s.split(':');
+        let pubkey = parts.next();
+        let alias = parts.next();
+
+        match (pubkey, alias) {
+            (Some(pubkey), None) => Err(format!(
+                "must be in the format 'PUBKEY:ALIAS' but found '{}'",
+                pubkey
+            )),
+            (None, _) => unreachable!(),
+            (Some(pubkey), Some(alias)) => {
+                if pubkey.is_empty() || alias.is_empty() {
+                    return Err(format!(
+                        "\t\nMust be in the format 'PUBKEY:ALIAS' but found '{}:{}'",
+                        pubkey, alias
+                    ));
+                }
+
+                if pubkey.len() != 44 {
+                    return Err(format!("\t\nPUBKEY '{}' has an invalid length", pubkey,));
+                }
+
+                if base64::decode(pubkey).is_err() {
+                    return Err(format!("\n\t'{}' is not a valid public key", pubkey,));
+                }
+
+                Ok(Alias {
+                    inner: (pubkey.into(), alias.into()),
+                })
+            }
+        }
+    }
 }

src/main.rs

@@ -55,6 +55,8 @@ async fn main() -> Result<()> {
 async fn try_main(args: Args) -> Result<()> {
     tracing_subscriber::fmt::init();
 
+    let aliases = args.aliases();
+
     let running = Arc::new(AtomicBool::new(true));
 
     info!("Registering metrics...");
@@ -86,7 +88,9 @@ async fn try_main(args: Args) -> Result<()> {
         let before = Instant::now();
 
         debug!("Updating metrics...");
-        metrics.update(&WireguardState::scrape().await?).await;
+        metrics
+            .update(&WireguardState::scrape(&aliases).await?)
+            .await;
         let after = Instant::now();
 
         let elapsed = after.duration_since(before);

src/metrics.rs

@@ -43,15 +43,15 @@ impl Metrics {
                 "wireguard_peer_bytes_total",
                 "Total number of bytes per direction for a peer",
             ),
-            &["interface", "peer", "direction"],
+            &["interface", "peer", "direction", "alias"],
         )?;
 
         let duration_since_latest_handshake = IntGaugeVec::new(
             Opts::new(
                 "wireguard_duration_since_latest_handshake",
                 "During since latest handshake for a peer",
             ),
-            &["interface", "peer"],
+            &["interface", "peer", "alias"],
         )?;
 
         debug!("Registering wireguard metrics");
@@ -112,6 +112,10 @@ impl Metrics {
                 &state.interfaces[p.interface],
                 &p.pubkey,
                 "tx",
+                &p.alias
+                    .as_ref()
+                    .map(ToOwned::to_owned)
+                    .unwrap_or_else(String::new),
             ]);
             let diff = p.tx_bytes - pbtt.get();
             pbtt.inc_by(diff);
@@ -120,6 +124,10 @@ impl Metrics {
                 &state.interfaces[p.interface],
                 &p.pubkey,
                 "rx",
+                &p.alias
+                    .as_ref()
+                    .map(ToOwned::to_owned)
+                    .unwrap_or_else(String::new),
             ]);
             let diff = p.rx_bytes - pbtr.get();
             pbtr.inc_by(diff);
@@ -131,7 +139,14 @@ impl Metrics {
                 let elapsed = now.signed_duration_since(hs_ts);
 
                 self.duration_since_latest_handshake
-                    .with_label_values(&[&state.interfaces[p.interface], &p.pubkey])
+                    .with_label_values(&[
+                        &state.interfaces[p.interface],
+                        &p.pubkey,
+                        &p.alias
+                            .as_ref()
+                            .map(ToOwned::to_owned)
+                            .unwrap_or_else(String::new),
+                    ])
                     .set(elapsed.num_milliseconds());
             }
         }

src/wireguard.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use color_eyre::eyre::{Result, WrapErr};
 use tokio::process::Command;
 
@@ -20,7 +22,7 @@ pub struct WireguardState {
 }
 
 impl WireguardState {
-    pub async fn scrape() -> Result<Self> {
+    pub async fn scrape(aliases: &HashMap<&str, &str>) -> Result<Self> {
         let mut peers = Vec::with_capacity(32); // Picked by fair dice roll
         let mut interfaces = Vec::new();
 
@@ -56,6 +58,7 @@ impl WireguardState {
                     let ts = handshake_ts.parse()?;
                     peers.push(Peer {
                         interface: interfaces.len() - 1,
+                        alias: aliases.get(pubkey).map(|&s| s.into()),
                         pubkey: pubkey.into(),
                         handshake_timestamp: if ts == 0 { None } else { Some(ts) },
                         tx_bytes: tx_bytes.parse()?,
@@ -84,6 +87,7 @@ impl WireguardState {
 #[derive(Clone, PartialEq, Debug, Default)]
 pub struct Peer {
     pub pubkey: String,
+    pub alias: Option<String>,
     pub interface: usize,
     pub tx_bytes: u64,
     pub rx_bytes: u64,