Add option to exclude IPs

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
 

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<Self, Self::Err> {
+        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<InfoLevel>,
 
+    #[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<Vec<HostFilter>>,
+
     #[command(flatten)]
     pub render_opts: RenderOpts,
 }

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<IpAddr, String>,
     ) {
+        let hostnames: Vec<String> = 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<Vec<HostFilter>>) {
+        self.state.excluded_ips = ex;
+    }
     pub fn end(&mut self) {
         self.terminal.show_cursor().unwrap();
     }

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<Connection, ConnectionData>,
     /// Used for reducing logging noise.
     known_orphan_sockets: VecDeque<LocalSocket>,
+    pub excluded_ips: Option<Vec<HostFilter>>,
 }
 
 impl UIState {
     pub fn update(
         &mut self,
         connections_to_procs: HashMap<LocalSocket, String>,
-        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,

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