ipfilter/cli/src/main.rs

use core::net::{Ipv4Addr, Ipv6Addr};
use std::collections::HashSet;
use std::fs;

use ipfilter::{v4, v6};

use anyhow::Result;
use clap::Parser;
use ipnet::{Ipv4Subnets, Ipv6Subnets};

mod cli;
use cli::{Cli, Commands};
use itertools::Itertools;

fn main() -> Result<()> {
	let cli = Cli::parse();

	let input = fs::File::open(cli.input)?;

	match cli.command {
		Commands::List => {
			let mut reader = csv::ReaderBuilder::new()
				.has_headers(false)
				.from_reader(input);

			let mut countries = HashSet::<(String, String)>::new();

			for result in reader.records() {
				let record = result?;
				countries.insert((record[2].to_owned(), record[3].to_owned()));
			}

			let mut countries: Vec<_> = countries.drain().collect();
			countries.sort_by(|lhs, rhs| lhs.0.cmp(&rhs.0));

			for (code, country) in countries {
				println!("{code} - {country}");
			}
		}
		Commands::Merge {
			countries,
			output,
			_6,
		} => {
			let mut reader = csv::ReaderBuilder::new()
				.has_headers(false)
				.from_reader(input);

			let countries: HashSet<String> = countries.split(',').map(ToOwned::to_owned).collect();

			let records = reader
				.records()
				.filter_ok(|r| countries.contains(&r[2]))
				.map(|r| r.map_err(anyhow::Error::from));

			macro_rules! merge_ip {
				($v:ident, $net:ty, $addr:ty) => {
					let r = records
						.map(|r| {
							r.and_then(|record| {
								Ok(<$net>::new(
									<$addr>::from_bits(record[0].parse()?),
									<$addr>::from_bits(record[1].parse()?),
									0,
								))
							})
						})
						.flatten_ok();

					let range = $v::from_fallible_iter(r)?;

					if let Some(o) = output {
						let mut writer = fs::File::create(o)?;
						$v::write_to(range, &mut writer)?;
					} else {
						for subnet in &range {
							println!("{subnet}");
						}
					};
				};
			}

			if _6 {
				merge_ip!(v6, Ipv6Subnets, Ipv6Addr);
			} else {
				merge_ip!(v4, Ipv4Subnets, Ipv4Addr);
			}
		}
		Commands::Load { _6 } => {
			macro_rules! load_ip {
				($v:ident, $net:ty, $addr:ty) => {
					let range = $v::read_from(&mut &input)?;

					for subnet in &range {
						println!("{subnet}");
					}
				};
			}

			if _6 {
				load_ip!(v6, Ipv6Subnets, Ipv6Addr);
			} else {
				load_ip!(v4, Ipv4Subnets, Ipv4Addr);
			}
		}
	}

	Ok(())
}