From d8b48aba9a19e984fc55b9401d7a27b517b373e8 Mon Sep 17 00:00:00 2001 From: Keenan Tims Date: Sat, 18 Mar 2023 11:44:41 -0700 Subject: [PATCH] Refactoring & compatibility with aggregate6 options --- Cargo.lock | 23 ++++--- Cargo.toml | 2 +- src/iputils.rs | 179 +++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 139 +++++++++++++++++++++----------------- 4 files changed, 271 insertions(+), 72 deletions(-) create mode 100644 src/iputils.rs diff --git a/Cargo.lock b/Cargo.lock index 9e156eb..183a109 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -16,9 +16,9 @@ checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f" [[package]] name = "clap" -version = "4.1.8" +version = "4.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3d7ae14b20b94cb02149ed21a86c423859cbe18dc7ed69845cace50e52b40a5" +checksum = "ce38afc168d8665cfc75c7b1dd9672e50716a137f433f070991619744a67342a" dependencies = [ "bitflags", "clap_derive", @@ -31,9 +31,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.1.8" +version = "4.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44bec8e5c9d09e439c4335b1af0abaab56dcf3b94999a936e1bb47b9134288f0" +checksum = "fddf67631444a3a3e3e5ac51c36a5e01335302de677bd78759eaa90ab1f46644" dependencies = [ "heck", "proc-macro-error", @@ -44,9 +44,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "350b9cf31731f9957399229e9b2adc51eeabdfbe9d71d9a0552275fd12710d09" +checksum = "033f6b7a4acb1f358c742aaca805c939ee73b4c6209ae4318ec7aca81c42e646" dependencies = [ "os_str_bytes", ] @@ -97,10 +97,11 @@ checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286" [[package]] name = "io-lifetimes" -version = "1.0.6" +version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfa919a82ea574332e2de6e74b4c36e74d41982b335080fa59d4ef31be20fdf3" +checksum = "0dd6da19f25979c7270e70fa95ab371ec3b701cd0eefc47667a09785b3c59155" dependencies = [ + "hermit-abi", "libc", "windows-sys 0.45.0", ] @@ -122,9 +123,9 @@ dependencies = [ [[package]] name = "is-terminal" -version = "0.4.4" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21b6b32576413a8e69b90e952e4a026476040d81017b80445deda5f2d3921857" +checksum = "8687c819457e979cc940d09cb16e42a1bf70aa6b60a549de6d3a62a0ee90c69e" dependencies = [ "hermit-abi", "io-lifetimes", @@ -200,7 +201,7 @@ dependencies = [ [[package]] name = "rs-aggregate" -version = "0.1.0" +version = "0.2.0" dependencies = [ "clap", "clio", diff --git a/Cargo.toml b/Cargo.toml index 8ab911e..3da2d65 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rs-aggregate" -version = "0.1.0" +version = "0.2.0" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/src/iputils.rs b/src/iputils.rs new file mode 100644 index 0000000..62bd455 --- /dev/null +++ b/src/iputils.rs @@ -0,0 +1,179 @@ +use std::{ + error::Error, + fmt::Display, + net::{IpAddr, Ipv4Addr}, + str::FromStr, +}; + +use ipnet::{IpNet, Ipv4Net, Ipv6Net}; +use iprange::{IpRange, IpRangeIter}; + +#[derive(Default)] +pub struct IpBothRange { + v4: IpRange, + v6: IpRange, +} + +impl IpBothRange { + pub fn new() -> IpBothRange { + IpBothRange::default() + } + pub fn add(&mut self, net: IpOrNet) { + match net { + IpOrNet::IpNet(net) => match net { + IpNet::V4(v4_net) => drop(self.v4.add(v4_net)), + IpNet::V6(v6_net) => drop(self.v6.add(v6_net)), + }, + IpOrNet::IpAddr(addr) => match addr { + IpAddr::V4(v4_addr) => drop(self.v4.add(v4_addr.into())), + IpAddr::V6(v6_addr) => drop(self.v6.add(v6_addr.into())), + }, + } + } + pub fn simplify(&mut self) { + self.v4.simplify(); + self.v6.simplify(); + } +} + +pub struct IpBothRangeIter<'a> { + v4_iter: IpRangeIter<'a, Ipv4Net>, + v6_iter: IpRangeIter<'a, Ipv6Net>, + _v4_done: bool, +} + +impl<'a> Iterator for IpBothRangeIter<'a> { + type Item = IpNet; + fn next(&mut self) -> Option { + if self._v4_done { + match self.v6_iter.next() { + Some(net) => return Some(net.into()), + None => return None, + } + } + match self.v4_iter.next() { + Some(net) => Some(net.into()), + None => { + self._v4_done = true; + match self.v6_iter.next() { + Some(net) => Some(net.into()), + None => None, + } + } + } + } +} + +impl<'a> IntoIterator for &'a IpBothRange { + type Item = IpNet; + type IntoIter = IpBothRangeIter<'a>; + fn into_iter(self) -> Self::IntoIter { + IpBothRangeIter { + v4_iter: self.v4.iter(), + v6_iter: self.v6.iter(), + _v4_done: false, + } + } +} + +pub enum IpOrNet { + IpNet(IpNet), + IpAddr(IpAddr), +} + +#[derive(Debug, Clone)] +pub struct NetParseError { + #[allow(dead_code)] + msg: String, +} + +impl Display for NetParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("Unable to parse address") + } +} + +impl Error for NetParseError {} + +impl IpOrNet { + // Accepted formats: + // netmask - 1.1.1.0/255.255.255.0 + // wildcard mask - 1.1.1.0/0.0.0.255 + fn parse_mask(p: &str) -> Result> { + let mask = p.parse::(); + match mask { + Ok(mask) => { + let intrep: u32 = mask.into(); + let lead_ones = intrep.leading_ones(); + if lead_ones > 0 { + if lead_ones + intrep.trailing_zeros() == 32 { + Ok(lead_ones.try_into()?) + } else { + Err(Box::new(NetParseError { + msg: "Invalid subnet mask".to_owned(), + })) + } + } else { + let lead_zeros = intrep.leading_zeros(); + if lead_zeros + intrep.trailing_ones() == 32 { + Ok(lead_zeros.try_into()?) + } else { + Err(Box::new(NetParseError { + msg: "Invalid wildcard mask".to_owned(), + })) + } + } + } + Err(e) => Err(Box::new(e)), + } + } + fn from_parts(ip: &str, pfxlen: &str) -> Result> { + let ip = ip.parse::()?; + let pfxlenp = pfxlen.parse::(); + + match pfxlenp { + Ok(pfxlen) => Ok(IpNet::new(ip, pfxlen)?.into()), + Err(_) => { + if ip.is_ipv4() { + Ok(IpNet::new(ip, IpOrNet::parse_mask(pfxlen)?)?.into()) + } else { + Err(Box::new(NetParseError { + msg: "Mask form is not valid for IPv6 address".to_owned(), + })) + } + } + } + } + pub fn prefix_len(&self) -> u8 { + match self { + Self::IpNet(net) => net.prefix_len(), + Self::IpAddr(addr) => match addr { + IpAddr::V4(_) => 32, + IpAddr::V6(_) => 128, + }, + } + } +} + +impl FromStr for IpOrNet { + type Err = Box; + fn from_str(s: &str) -> Result { + let parts = s.split_once('/'); + match parts { + Some((ip, pfxlen)) => IpOrNet::from_parts(ip, pfxlen), + None => Ok(s.parse::()?.into()), + } + } +} + +impl From for IpOrNet { + fn from(net: IpNet) -> Self { + IpOrNet::IpNet(net) + } +} + +impl From for IpOrNet { + fn from(addr: IpAddr) -> Self { + IpOrNet::IpAddr(addr) + } +} diff --git a/src/main.rs b/src/main.rs index 2ed9ea7..f916c96 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,11 @@ extern crate ipnet; extern crate iprange; +mod iputils; +use iputils::{IpBothRange, IpOrNet}; + use clio::*; -use ipnet::{IpNet, Ipv4Net, Ipv6Net}; -use iprange::IpRange; -use std::{io::BufRead, net::IpAddr}; +use std::io::BufRead; use clap::Parser; @@ -13,80 +14,98 @@ use clap::Parser; struct Args { #[clap(value_parser, default_value = "-")] input: Input, + #[arg( + short, + long, + default_value = "128", + help = "Sets the maximum prefix length for entries read. Longer prefixes will be discarded prior to processing." + )] + max_prefixlen: u8, + #[arg(short, long, help = "truncate IP/mask to network/mask (else ignore)")] + truncate: bool, + #[arg(id="4", short, help = "Only output IPv4 prefixes", conflicts_with("6"))] + only_v4: bool, + #[arg(id="6", short, help = "Only output IPv6 prefixes", conflicts_with("4"))] + only_v6: bool, } +impl Default for Args { + fn default() -> Self { + Args { + input: clio::Input::default(), + max_prefixlen: 128, + truncate: false, + only_v4: false, + only_v6: false, + } + } +} + +#[derive(Parser)] + struct IpParseError { ip: String, problem: String, } -struct IpBothRange { - v4: IpRange, - v6: IpRange, -} - type Errors = Vec; -fn simplify_input(mut input: Input) -> (IpBothRange, Errors) { - let mut res = IpBothRange { - v4: IpRange::new(), - v6: IpRange::new(), - }; - let mut errors = Errors::new(); - for line in input.lock().lines() { - for net in line.unwrap().split_whitespace().to_owned() { - match net.parse() { - Ok(ipnet) => match ipnet { - IpNet::V4(v4_net) => { - res.v4.add(v4_net.trunc()); - () +#[derive(Default)] +struct App { + args: Args, + prefixes: IpBothRange, + errors: Errors, +} + +impl App { + fn add_prefix(&mut self, pfx: IpOrNet) { + // Parser accepts host bits set, so detect that case and error if not truncate mode + if !self.args.truncate { + match pfx { + IpOrNet::IpNet(net) => { + if net.addr() != net.network() { + eprintln!("ERROR: '{}' is not a valid IP network, ignoring.", net); + return; } - IpNet::V6(v6_net) => { - res.v6.add(v6_net.trunc()); - () - } - }, - Err(_) => { - // First try to add it as a bare IP - match net.parse() { - Ok(ip) => match ip { - IpAddr::V4(v4_ip) => { - res.v4.add(Ipv4Net::new(v4_ip, 32).unwrap()); - () - } - IpAddr::V6(v6_ip) => { - res.v6.add(Ipv6Net::new(v6_ip, 128).unwrap()); - () - } - }, - Err(error) => { - eprintln!("ERROR: {} - {}, ignoring.", net, error.to_string()); - errors.push(IpParseError { - ip: net.to_string(), - problem: error.to_string(), - }); - } + } + IpOrNet::IpAddr(_) => (), + } + } + if pfx.prefix_len() <= self.args.max_prefixlen { + self.prefixes.add(pfx); + } + } + fn simplify_input(&mut self) { + for line in self.args.input.to_owned().lock().lines() { + for net in line.unwrap().split_whitespace().to_owned() { + let pnet = net.parse::(); + match pnet { + Ok(pnet) => self.add_prefix(pnet), + Err(e) => { + self.errors.push(IpParseError { + ip: net.to_string(), + problem: e.to_string(), + }); + eprintln!("ERROR: '{}' is not a valid IP network, ignoring.", net); } } } } + self.prefixes.simplify(); } - res.v4.simplify(); - res.v6.simplify(); - (res, errors) + fn main(&mut self) { + self.args = Args::parse(); + + self.simplify_input(); + + for net in &self.prefixes { + println!("{}", net); + } + } } fn main() { - let args = Args::parse(); - let input = args.input; - - let (res, _) = simplify_input(input); - - for net in &res.v4 { - println!("{}", net); - } - for net in &res.v6 { - println!("{}", net); - } + let mut app = App::default(); + app.main(); }