netdig 0.0.5

#[cfg(feature = "json")]
mod json;
mod tabular;
mod warn;
mod whois;

use anyhow::{Context, Result};
use ipnet::Ipv4Net;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use std::{
    collections::BTreeSet, fmt::Debug, net::Ipv4Addr, ops::Deref,
    thread::sleep, time::Duration,
};
#[cfg(feature = "tracing")]
use tracing::{debug, instrument, warn};
use warn::{Warn, Warning};
use whois::Whois;

#[cfg(feature = "json")]
pub use json::Json;
pub use tabular::Tabular;

#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
pub struct Output<T> {
    pub input: Option<T>,
    pub valid: bool,
    pub canonical: Ipv4Net,
    pub contained_by: Option<Ipv4Net>,
    pub network: Ipv4Addr,
    pub address: Ipv4Addr,
    pub netmask: Ipv4Addr,
    pub hosts: usize,
    pub whois: Vec<String>,
    pub warnings: Vec<Warning>,
}

#[derive(Debug)]
pub struct IpSet<T> {
    ips: BTreeSet<(Option<T>, Ipv4Net)>,
    warnings: Vec<Box<dyn Warn>>,
}

impl<T> Default for IpSet<T> {
    fn default() -> Self {
        Self {
            ips: BTreeSet::default(),
            warnings: warn::all(),
        }
    }
}

impl<T: AsRef<str> + Ord> IpSet<T> {
    /// Attempt to add a CIDR block from a string. This will first try to parse
    /// it as an `Ipv4Net`, and then fall back to trying it as an `Ipv4Addr`,
    /// which can be converted to a /32.
    ///
    /// # Errors
    ///
    /// Returns an error if the IP fails to parse.
    pub fn insert(&mut self, s: T) -> Result<&mut Self> {
        let ip = if let Ok(ip) = s.as_ref().parse() {
            ip
        } else {
            #[cfg(feature = "tracing")]
            debug!("parsing as a CIDR block failed; trying as an IP");
            let ip: Ipv4Addr =
                s.as_ref().parse().context("parsing as Ipv4Addr")?;

            ip.into()
        };

        self.ips.insert((Some(s), ip));

        Ok(self)
    }
}

impl From<Vec<Ipv4Net>> for IpSet<String> {
    fn from(value: Vec<Ipv4Net>) -> Self {
        Self {
            ips: value
                .into_iter()
                .map(|i| (Some(i.to_string()), i))
                .collect(),
            ..Self::default()
        }
    }
}

impl TryFrom<Vec<String>> for IpSet<String> {
    type Error = anyhow::Error;

    fn try_from(value: Vec<String>) -> std::result::Result<Self, Self::Error> {
        let mut ips = Self::default();
        for ip in value {
            ips.insert(ip)?;
        }

        Ok(ips)
    }
}

impl<T> Deref for IpSet<T> {
    type Target = BTreeSet<(Option<T>, Ipv4Net)>;

    fn deref(&self) -> &Self::Target {
        &self.ips
    }
}

impl<T: std::fmt::Debug + PartialEq<String> + Clone + ToString> IpSet<T> {
    fn ips(&self) -> Vec<Ipv4Net> {
        self.iter().map(|ip| ip.1).collect()
    }

    /// # Errors
    ///
    /// Returns an error if anything goes wrong talking to whois.
    #[cfg_attr(feature = "tracing", instrument(ret, err))]
    pub fn check_ips(&self, whois: bool) -> Result<Vec<Output<T>>> {
        // Aggregate the IPs so we can look for supersets of ranges.
        let aggregated_ips = Ipv4Net::aggregate(&self.ips());

        // If the lengths between passed IPs and aggregated IPs don't match, it
        // means we've got redundant subnets. Warn the user.
        #[cfg(feature = "tracing")]
        if aggregated_ips.len() != self.len() {
            warn!(
                aggregated = aggregated_ips.len(),
                ips = self.len(),
                "one or more subnets are not aggregated"
            );
        }

        let mut outputs = vec![];
        for (raw, ip) in self.iter() {
            let canonical = ip.trunc();

            // Figure out whether this is a subnet that's part of a greater range in
            // the set.
            let contained_by =
                aggregated_ips.iter().find(|i| *i != ip && i.contains(ip));

            // Check to see whether we have any warnings.
            let warnings = self
                .warnings
                .iter()
                .filter_map(|w| w.check(canonical))
                .collect();

            let output = Output {
                input: raw.clone(),
                valid: raw
                    .as_ref()
                    .map_or(true, |s| *s == canonical.to_string()),
                canonical,
                contained_by: contained_by.copied(),
                network: ip.network(),
                address: ip.addr(),
                netmask: ip.netmask(),
                hosts: ip.hosts().count(),
                whois: if whois {
                    whois::Arin::lookup(*ip)?
                } else {
                    vec![]
                },
                warnings,
            };
            outputs.push(output);

            sleep(Duration::from_millis(100));
        }

        Ok(outputs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_whois() {
        assert_eq!(
            whois::Mock::lookup(Ipv4Addr::new(127, 0, 0, 1).into()).unwrap(),
            vec!["WHOIS RESULT"]
        );
    }
}