netdig 0.0.5

Utilities for analyzing and aggregating CIDR blocks
Documentation
use anyhow::Result;
use ipnet::Ipv4Net;
use std::{
    io::{Read, Write},
    net::TcpStream,
    time::Duration,
};
#[cfg(feature = "tracing")]
use tracing::trace;

pub(crate) trait Whois {
    fn lookup(ip: Ipv4Net) -> Result<Vec<String>>;
}

pub(crate) struct Arin {}

impl Whois for Arin {
    fn lookup(ip: Ipv4Net) -> Result<Vec<String>> {
        // Look up the whois info.
        let mut stream = TcpStream::connect("whois.arin.net:43")?;
        stream.set_read_timeout(Some(Duration::from_secs(5)))?;
        stream.set_write_timeout(Some(Duration::from_secs(5)))?;
        let ip_addr = ip.addr().to_string();
        stream.write_all(&[b"n + ", ip_addr.as_bytes(), b"\n"].concat())?;
        #[cfg(feature = "tracing")]
        trace!(?ip_addr, "wrote request");
        let mut buf = String::new();
        stream.read_to_string(&mut buf)?;
        #[cfg(feature = "tracing")]
        trace!(?buf, "read response");

        // Parse the response into key/val pairs.
        Ok(buf
            .lines()
            .filter_map(|l| {
                if l.starts_with('#') {
                    return None;
                }

                let (k, v) = l.split_once(':')?;

                match k.trim() {
                    "NetRange" | "CIDR" | "Organization" | "City"
                    | "StateProv" => Some(v.trim().to_string()),
                    _ => None,
                }
            })
            .collect())
    }
}

#[cfg(test)]
pub(crate) struct Mock;

#[cfg(test)]
impl Whois for Mock {
    fn lookup(_ip: Ipv4Net) -> Result<Vec<String>> {
        Ok(vec!["WHOIS RESULT".into()])
    }
}