diff options
-rw-r--r-- | src/os/lsof.rs | 25 | ||||
-rw-r--r-- | src/os/lsof_utils.rs | 38 |
2 files changed, 34 insertions, 29 deletions
diff --git a/src/os/lsof.rs b/src/os/lsof.rs index f6d1e3f..088d9cf 100644 --- a/src/os/lsof.rs +++ b/src/os/lsof.rs @@ -1,24 +1,9 @@ -use std::collections::HashMap; - -use crate::{network::LocalSocket, os::lsof_utils::get_connections, OpenSockets}; +use crate::{os::lsof_utils::get_connections, OpenSockets}; pub(crate) fn get_open_sockets() -> OpenSockets { - let mut open_sockets = HashMap::new(); - - let connections = get_connections(); - - for raw_connection in connections { - open_sockets.insert( - LocalSocket { - ip: raw_connection.get_local_ip(), - port: raw_connection.get_local_port(), - protocol: raw_connection.get_protocol(), - }, - raw_connection.process_name.clone(), - ); - } + let sockets_to_procs = get_connections() + .filter_map(|raw| raw.as_local_socket().map(|s| (s, raw.process_name))) + .collect(); - OpenSockets { - sockets_to_procs: open_sockets, - } + OpenSockets { sockets_to_procs } } diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs index 21735f1..5e55dad 100644 --- a/src/os/lsof_utils.rs +++ b/src/os/lsof_utils.rs @@ -1,8 +1,9 @@ use std::{ffi::OsStr, net::IpAddr, process::Command, sync::OnceLock}; +use log::warn; use regex::Regex; -use crate::network::Protocol; +use crate::network::{LocalSocket, Protocol}; #[allow(dead_code)] #[derive(Debug, Clone)] @@ -103,16 +104,35 @@ impl RawConnection { } } - pub fn get_protocol(&self) -> Protocol { - Protocol::from_str(&self.protocol).unwrap() + pub fn get_protocol(&self) -> Option<Protocol> { + Protocol::from_str(&self.protocol) } - pub fn get_local_ip(&self) -> IpAddr { - self.local_ip.parse().unwrap() + pub fn get_local_ip(&self) -> Option<IpAddr> { + self.local_ip.parse().ok() } - pub fn get_local_port(&self) -> u16 { - self.local_port.parse::<u16>().unwrap() + pub fn get_local_port(&self) -> Option<u16> { + self.local_port.parse::<u16>().ok() + } + + pub fn as_local_socket(&self) -> Option<LocalSocket> { + let process = &self.process_name; + + let Some(ip) = self.get_local_ip() else { + warn!(r#"Failed to get the local IP of a connection belonging to "{process}"."#); + return None; + }; + let Some(port) = self.get_local_port() else { + warn!(r#"Failed to get the local port of a connection belonging to "{process}"."#); + return None; + }; + let Some(protocol) = self.get_protocol() else { + warn!(r#"Failed to get the protocol of a connection belonging to "{process}"."#); + return None; + }; + + Some(LocalSocket { ip, port, protocol }) } } @@ -203,7 +223,7 @@ com.apple 590 etoledom 204u IPv4 0x28ffb9c04111253f 0t0 TCP 192.168.1. } fn test_raw_connection_parse_local_port(raw_output: &str) { let connection = RawConnection::new(raw_output).unwrap(); - assert_eq!(connection.get_local_port(), 1111); + assert_eq!(connection.get_local_port(), Some(1111)); } #[test] @@ -216,7 +236,7 @@ com.apple 590 etoledom 204u IPv4 0x28ffb9c04111253f 0t0 TCP 192.168.1. } fn test_raw_connection_parse_protocol(raw_line: &str) { let connection = RawConnection::new(raw_line).unwrap(); - assert_eq!(connection.get_protocol(), Protocol::Udp); + assert_eq!(connection.get_protocol(), Some(Protocol::Udp)); } #[test] |