Refactoring & compatibility with aggregate6 options

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",

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

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<Ipv4Net>,
+    v6: IpRange<Ipv6Net>,
+}
+
+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<Self::Item> {
+        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<u8, Box<dyn Error>> {
+        let mask = p.parse::<Ipv4Addr>();
+        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<Self, Box<dyn Error>> {
+        let ip = ip.parse::<IpAddr>()?;
+        let pfxlenp = pfxlen.parse::<u8>();
+
+        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<dyn Error>;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let parts = s.split_once('/');
+        match parts {
+            Some((ip, pfxlen)) => IpOrNet::from_parts(ip, pfxlen),
+            None => Ok(s.parse::<IpAddr>()?.into()),
+        }
+    }
+}
+
+impl From<IpNet> for IpOrNet {
+    fn from(net: IpNet) -> Self {
+        IpOrNet::IpNet(net)
+    }
+}
+
+impl From<IpAddr> for IpOrNet {
+    fn from(addr: IpAddr) -> Self {
+        IpOrNet::IpAddr(addr)
+    }
+}

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<Ipv4Net>,
-    v6: IpRange<Ipv6Net>,
-}
-
 type Errors = Vec<IpParseError>;
 
-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,
 }
-                    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(),
-                            });
-                        }
-                    }
-                }
-            }
-        }
-    }
-    res.v4.simplify();
-    res.v6.simplify();
 
-    (res, 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;
+                    }
+                }
+                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::<IpOrNet>();
+                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();
+    }
+
+    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();
 }