evenmoreutils/src/bin/domnet.rs

169 lines
4.9 KiB
Rust

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);
}
}