evenmoreutils/src/bin/domnet.rs

use std::fmt;
use std::net::{AddrParseError, Ipv4Addr, Ipv6Addr};
use std::num::ParseIntError;
use std::str::FromStr;

use thiserror::Error;

struct V4Subnet(Vec<bool>);
struct V6Subnet(Vec<bool>);

#[derive(Error, Debug)]
enum SubnetParseError {
    #[error("failed to parse CIDR block")]
    Cidr,
    #[error("failed to parse IP address")]
    Addr(#[from] AddrParseError),
    #[error("failed to parse prefix length")]
    PrefixLength(#[from] ParseIntError),
}

fn bitstring_to_int(b: &[bool], size: usize) -> u128 {
    let mut result = 0;
    for idx in 0..size {
        result *= 2;
        if let Some(true) = b.get(idx) {
            result += 1;
        }
    }
    result
}

fn int_to_bitstring(val: u128, size: usize, prefix_len: usize) -> Vec<bool> {
    let mut result = Vec::new();
    for idx in 0..prefix_len {
        result.push((val >> (size - 1 - idx)) % 2 == 1);
    }
    result
}

impl fmt::Display for V4Subnet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let addr = Ipv4Addr::from(bitstring_to_int(&self.0, 32) as u32);
        write!(f, "{}/{}", addr, self.0.len())
    }
}

impl FromStr for V4Subnet {
    type Err = SubnetParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (addr, len) = s.split_once('/').ok_or(SubnetParseError::Cidr)?;
        let val = u32::from(Ipv4Addr::from_str(addr)?);
        let prefix_len = usize::from_str(len)?;
        Ok(V4Subnet(int_to_bitstring(val as u128, 32, prefix_len)))
    }
}

impl fmt::Display for V6Subnet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let addr = Ipv6Addr::from(bitstring_to_int(&self.0, 128));
        write!(f, "{}/{}", addr, self.0.len())
    }
}

impl FromStr for V6Subnet {
    type Err = SubnetParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (addr, len) = s.split_once('/').ok_or(SubnetParseError::Cidr)?;
        let val = u128::from(Ipv6Addr::from_str(addr)?);
        let prefix_len = usize::from_str(len)?;
        Ok(V6Subnet(int_to_bitstring(val, 128, prefix_len)))
    }
}

/// Finds the smallest set of subnets that complement the set of subnets provided
fn invert_subnets(subnets: &Vec<Vec<bool>>, size: usize) -> Vec<Vec<bool>> {
    invert_subnets_helper(&mut Vec::new(), subnets, size)
}

// Finds the smallest set of subnets that complement, within `root`, the set of subnets provided
fn invert_subnets_helper(
    root: &mut Vec<bool>,
    subnets: &Vec<Vec<bool>>,
    size: usize,
) -> Vec<Vec<bool>> {
    if subnets.iter().any(|subnet| root.starts_with(subnet)) {
        // We're already within one of the provided subnets, bail out
        return Vec::new();
    }

    if subnets.iter().any(|subnet| subnet.starts_with(root)) {
        // We need to narrow down further

        if root.len() == size {
            // There's nowhere further to go
            return Vec::new();
        }

        // left branch
        root.push(false);
        let mut results = invert_subnets_helper(root, subnets, size);
        root.pop();

        // right branch
        root.push(true);
        results.append(&mut invert_subnets_helper(root, subnets, size));
        root.pop();

        results
    } else {
        vec![root.to_owned()]
    }
}

fn usage(name: Option<String>) {
    eprintln!("usage: {} SUBNET ...", name.as_deref().unwrap_or("domnet"));
    eprintln!("Calculate the inverse of a set of IP subnets.");
    eprintln!("Arguments must be all IPv4 subnets or all IPv6 subnets, in CIDR block format.");
    std::process::exit(1);
}

fn main() {
    let mut args = std::env::args();
    let name = args.next();
    let first = args.next();

    let subnet = if let Some(s) = first {
        s
    } else {
        return usage(name);
    };
    if let Ok(v4subnet) = V4Subnet::from_str(&subnet) {
        let rest = args
            .map(|s| V4Subnet::from_str(&s))
            .collect::<Result<Vec<V4Subnet>, SubnetParseError>>();
        let subnets = if let Ok(r) = rest {
            std::iter::once(v4subnet)
                .chain(r.into_iter())
                .map(|s| s.0)
                .collect::<Vec<Vec<bool>>>()
        } else {
            return usage(name);
        };
        let inverse = invert_subnets(&subnets, 32);
        for subnet in inverse.into_iter() {
            println!("{}", V4Subnet(subnet));
        }
    } else if let Ok(v6subnet) = V6Subnet::from_str(&subnet) {
        let rest = args
            .map(|s| V6Subnet::from_str(&s))
            .collect::<Result<Vec<V6Subnet>, SubnetParseError>>();
        let subnets = if let Ok(r) = rest {
            std::iter::once(v6subnet)
                .chain(r.into_iter())
                .map(|s| s.0)
                .collect::<Vec<Vec<bool>>>()
        } else {
            return usage(name);
        };
        let inverse = invert_subnets(&subnets, 128);
        for subnet in inverse.into_iter() {
            println!("{}", V6Subnet(subnet));
        }
    } else {
        usage(name);
    }
}