subnet inversion calculator

src/main.rs

@@ -0,0 +1,168 @@
+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);
+    }
+}