arpbox 0.2.0

Enrich ARP data with device information from NetBox
use serde::Deserialize;
use std::collections::HashMap;
use std::io::{self, BufRead};

#[derive(Deserialize)]
struct Config {
    netbox_url: String,
    netbox_token: String,
    netbox_key: Option<String>,
}

#[derive(Deserialize)]
struct NetBoxResponse {
    results: Vec<NetBoxInterface>,
}

#[derive(Deserialize)]
struct NetBoxInterface {
    name: String,
    device: Option<NetBoxDevice>,
}

#[derive(Deserialize)]
struct NetBoxDevice {
    url: String,
    #[serde(rename = "name")]
    _name: String,
}

#[derive(Deserialize)]
struct NetBoxDeviceFull {
    name: String,
    asset_tag: Option<String>,
}

fn load_config() -> Result<Config, Box<dyn std::error::Error>> {
    let config_dir = dirs::config_dir().ok_or("could not determine config directory")?;
    let config_path = config_dir.join("arpbox").join("config.toml");
    let contents = std::fs::read_to_string(&config_path)
        .map_err(|e| format!("failed to read config at {}: {}", config_path.display(), e))?;
    let config: Config = toml::from_str(&contents)?;
    Ok(config)
}

fn lookup_mac(
    client: &reqwest::blocking::Client,
    config: &Config,
    mac: &str,
    cache: &mut HashMap<String, Option<String>>,
) -> Option<String> {
    if let Some(cached) = cache.get(mac) {
        return cached.clone();
    }

    let url = format!(
        "{}/api/dcim/interfaces/?mac_address={}",
        config.netbox_url.trim_end_matches('/'),
        mac
    );

    let auth = match &config.netbox_key {
        Some(key) => format!("Bearer nbt_{}.{}", key, config.netbox_token),
        None => format!("Token {}", config.netbox_token),
    };

    let result = client
        .get(&url)
        .header("Authorization", &auth)
        .header("Accept", "application/json")
        .send()
        .ok()
        .and_then(|resp| resp.json::<NetBoxResponse>().ok())
        .and_then(|body| {
            body.results.into_iter().next().and_then(|iface| {
                iface.device.and_then(|d| {
                    let device: NetBoxDeviceFull = client
                        .get(&d.url)
                        .header("Authorization", &auth)
                        .header("Accept", "application/json")
                        .send()
                        .ok()?
                        .json()
                        .ok()?;
                    let name = match device.asset_tag {
                        Some(tag) => format!("{}:{}:{}", device.name, tag, iface.name),
                        None => format!("{}:{}", device.name, iface.name),
                    };
                    Some(name)
                })
            })
        });

    cache.insert(mac.to_string(), result.clone());
    result
}

fn lookup_mac_verbose(
    client: &reqwest::blocking::Client,
    config: &Config,
    mac: &str,
) -> Result<Option<String>, String> {
    let url = format!(
        "{}/api/dcim/interfaces/?mac_address={}",
        config.netbox_url.trim_end_matches('/'),
        mac
    );

    let auth = match &config.netbox_key {
        Some(key) => format!("Bearer nbt_{}.{}", key, config.netbox_token),
        None => format!("Token {}", config.netbox_token),
    };

    let resp = client
        .get(&url)
        .header("Authorization", &auth)
        .header("Accept", "application/json")
        .send()
        .map_err(|e| format!("request failed: {}", e))?;

    if !resp.status().is_success() {
        return Err(format!("HTTP {}: {}", resp.status(), resp.status().canonical_reason().unwrap_or("unknown")));
    }

    let body: NetBoxResponse = resp
        .json()
        .map_err(|e| format!("failed to parse response: {}", e))?;

    let iface = match body.results.into_iter().next() {
        Some(i) => i,
        None => return Ok(None),
    };

    let device_url = match &iface.device {
        Some(d) => &d.url,
        None => return Ok(None),
    };

    let resp = client
        .get(device_url)
        .header("Authorization", &auth)
        .header("Accept", "application/json")
        .send()
        .map_err(|e| format!("device request failed: {}", e))?;

    if !resp.status().is_success() {
        return Err(format!("device HTTP {}: {}", resp.status(), resp.status().canonical_reason().unwrap_or("unknown")));
    }

    let device: NetBoxDeviceFull = resp
        .json()
        .map_err(|e| format!("failed to parse device response: {}", e))?;

    let name = match device.asset_tag {
        Some(tag) => format!("{}:{}:{}", device.name, tag, iface.name),
        None => format!("{}:{}", device.name, iface.name),
    };

    Ok(Some(name))
}

fn main() {
    let config = match load_config() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Error: {}", e);
            std::process::exit(1);
        }
    };

    let client = reqwest::blocking::Client::new();

    let args: Vec<String> = std::env::args().collect();
    if args.len() == 2 {
        let mac = &args[1];
        match lookup_mac_verbose(&client, &config, mac) {
            Ok(Some(device)) => println!("(netbox://{})", device),
            Ok(None) => println!("No device found for {}", mac),
            Err(e) => {
                eprintln!("Error: {}", e);
                std::process::exit(1);
            }
        }
        return;
    }

    let mut cache: HashMap<String, Option<String>> = HashMap::new();
    let stdin = io::stdin();

    for line in stdin.lock().lines() {
        let line = match line {
            Ok(l) => l,
            Err(_) => continue,
        };

        let parts: Vec<&str> = line.split('\t').collect();
        if parts.len() < 2 {
            println!("{}", line);
            continue;
        }

        // Validate MAC-like format in second field
        let mac = parts[1];
        if mac.len() < 11 || !mac.contains(':') {
            println!("{}", line);
            continue;
        }

        if let Some(device_name) = lookup_mac(&client, &config, mac, &mut cache) {
            let vendor = if parts.len() >= 3 {
                format!("\t{}", parts[2..].join("\t"))
            } else {
                String::new()
            };
            println!(
                "{}\t{}\t(netbox://{}){}",
                parts[0], mac, device_name, vendor
            );
        } else {
            println!("{}", line);
        }
    }
}