diff --git a/CHANGELOG.md b/CHANGELOG.md index 92be8d57f..4110fb34c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) * Display bandwidth in different unit families #328 - @cyqsimon * CI: ensure a changelog entry exists for each PR #331 - @cyqsimon * Show interface names #340 - @ilyes-ced +* Add option to exclude IPs #341 - @ilyes-ced ### Changed diff --git a/src/cli.rs b/src/cli.rs index 644a500ad..a0dff269d 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,4 +1,36 @@ -use std::{net::Ipv4Addr, path::PathBuf}; +use std::{ + net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6}, + path::PathBuf, + str::FromStr, +}; + +#[derive(Clone, Debug)] +pub enum HostFilter { + Ipv4Addr(Ipv4Addr), + Ipv6Addr(Ipv6Addr), + SocketAddrV4(SocketAddrV4), + SocketAddrV6(SocketAddrV6), + Hostname(String), +} + +impl FromStr for HostFilter { + type Err = String; + + fn from_str(s: &str) -> Result { + if let Ok(ipv4) = s.parse() { + Ok(HostFilter::Ipv4Addr(ipv4)) + } else if let Ok(ipv6) = s.parse() { + Ok(HostFilter::Ipv6Addr(ipv6)) + } else if let Ok(socketv4) = s.parse() { + Ok(HostFilter::SocketAddrV4(socketv4)) + } else if let Ok(socketv6) = s.parse() { + Ok(HostFilter::SocketAddrV6(socketv6)) + } else { + // might need validation + Ok(HostFilter::Hostname(s.to_string())) + } + } +} use clap::{Args, Parser}; use clap_verbosity_flag::{InfoLevel, Verbosity}; @@ -38,6 +70,17 @@ pub struct Opt { #[derivative(Default(value = "Verbosity::new(0, 0)"))] pub verbosity: Verbosity, + #[arg(short, long)] + /// exclude ip addres with <-e x.x.x.x> + /// exclude multiple ip addresses with <-e x.x.x.x -e y.y.y.y> + /// examples: + /// IpV4: 127.0.0.1 + /// IpV6: 2001:db8::1 OR 2001:0db8:85a3:0000:0000:8a2e:0370:7334 + /// SocketAddrV4: 127.0.0.1:8080 + /// SocketAddrV6: "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:8080" + /// hostname: String + pub excluded: Option>, + #[command(flatten)] pub render_opts: RenderOpts, } diff --git a/src/display/ui.rs b/src/display/ui.rs index 4b6d71698..56e9d97de 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -4,7 +4,7 @@ use chrono::prelude::*; use ratatui::{backend::Backend, Terminal}; use crate::{ - cli::{Opt, RenderOpts}, + cli::{HostFilter, Opt, RenderOpts}, display::{ components::{HeaderDetails, HelpText, Layout, Table}, UIState, @@ -177,9 +177,39 @@ where utilization: Utilization, ip_to_host: HashMap, ) { + let hostnames: Vec = self + .state + .excluded_ips + .clone() + .unwrap_or_default() + .iter() + .filter_map(|hf| match hf { + HostFilter::Hostname(s) => Some(s.clone()), + _ => None, + }) + .collect(); + for (k, v) in &ip_to_host { + if hostnames.contains(v) { + match &self.state.excluded_ips { + None => {} + Some(_) => match k { + IpAddr::V4(ip) => self.push_to_excluded_ips(HostFilter::Ipv4Addr(*ip)), + IpAddr::V6(ip) => self.push_to_excluded_ips(HostFilter::Ipv6Addr(*ip)), + }, + } + } + } self.state.update(connections_to_procs, utilization); self.ip_to_host.extend(ip_to_host); } + fn push_to_excluded_ips(&mut self, ip: HostFilter) { + let mut vec = self.state.excluded_ips.take().unwrap_or_default(); + vec.push(ip); + self.state.excluded_ips = Some(vec); + } + pub fn set_excluded(&mut self, ex: Option>) { + self.state.excluded_ips = ex; + } pub fn end(&mut self) { self.terminal.show_cursor().unwrap(); } diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index 37ff63977..77e5c4f98 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -3,10 +3,11 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, hash::Hash, iter::FromIterator, - net::{IpAddr, Ipv4Addr, Ipv6Addr}, + net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, }; use crate::{ + cli::HostFilter, display::BandwidthUnitFamily, mt_log, network::{Connection, LocalSocket, Utilization}, @@ -94,14 +95,60 @@ pub struct UIState { pub connections_map: HashMap, /// Used for reducing logging noise. known_orphan_sockets: VecDeque, + pub excluded_ips: Option>, } impl UIState { pub fn update( &mut self, connections_to_procs: HashMap, - network_utilization: Utilization, + mut network_utilization: Utilization, ) { + if let Some(excluded_addresses) = &self.excluded_ips { + for ex in excluded_addresses { + network_utilization.connections.retain(|k, _| { + let ip_address = k.remote_socket.ip; + let port = k.remote_socket.port; + let socket = SocketAddr::new(ip_address, port); + + match ex { + HostFilter::Ipv4Addr(ipaddr) => { + if let IpAddr::V4(ipv4) = ip_address { + &ipv4 != ipaddr + } else { + true + } + } + HostFilter::Ipv6Addr(ipaddr) => { + if let IpAddr::V6(ipv6) = ip_address { + &ipv6 != ipaddr + } else { + true + } + } + HostFilter::SocketAddrV4(socketaddr) => { + if let SocketAddr::V4(socketv4) = socket { + &socketv4 != socketaddr + } else { + true + } + } + HostFilter::SocketAddrV6(socketaddr) => { + if let SocketAddr::V6(socketv6) = socket { + &socketv6 != socketaddr + } else { + true + } + } + HostFilter::Hostname(_name) => { + // not implemented yet + true + } + } + }); + } + } + self.utilization_data.push_back(UtilizationData { connections_to_procs, network_utilization, diff --git a/src/main.rs b/src/main.rs index a6a00d973..2584b267b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -155,6 +155,7 @@ where let mut ui = ui.lock().unwrap(); let paused = paused.load(Ordering::SeqCst); let ui_offset = ui_offset.load(Ordering::SeqCst); + ui.set_excluded(opts.excluded.clone()); if !paused { ui.update_state(sockets_to_procs, utilization, ip_to_host); }