From 81e95f74f5ac10c60f12b33880d3fe65a1eac549 Mon Sep 17 00:00:00 2001 From: etoledom Date: Mon, 18 Nov 2019 17:37:56 +0100 Subject: refactor(macos): handle lsof in macos in a different module * Refactor lsof connections as iterator * Fix build issues and warnings * Fix iterator to not finish on first wrong line * Remove println * Fix overflow. Creating the regex instance on every iteration is too bad performant * Adding connection regex as static * Small syntax improvement --- Cargo.lock | 1 + Cargo.toml | 1 + src/os/lsof_utils.rs | 167 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/os/macos.rs | 46 ++++---------- src/os/mod.rs | 4 +- 5 files changed, 183 insertions(+), 36 deletions(-) create mode 100644 src/os/lsof_utils.rs diff --git a/Cargo.lock b/Cargo.lock index 55262ce..3bb5901 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1253,6 +1253,7 @@ dependencies = [ "failure 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", "insta 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)", "ipnetwork 0.15.0 (registry+https://github.com/rust-lang/crates.io-index)", + "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)", "packet-builder 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)", "pnet 0.23.0 (registry+https://github.com/rust-lang/crates.io-index)", "pnet_base 0.23.0 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/Cargo.toml b/Cargo.toml index 94ebb0d..f58f3eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ signal-hook = "0.1.10" failure = "0.1.6" chrono = "0.4" regex = "1.3.1" +lazy_static = "1.4.0" [target.'cfg(target_os="linux")'.dependencies] procfs = "0.5.3" diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs new file mode 100644 index 0000000..caabd26 --- /dev/null +++ b/src/os/lsof_utils.rs @@ -0,0 +1,167 @@ +use std::process::{Command}; +use std::ffi::OsStr; +use regex::{Regex}; +use crate::network::Protocol; +use std::net::IpAddr; +use lazy_static::lazy_static; + +#[derive(Debug, Clone)] +pub struct RawConnection { + ip: String, + local_port: String, + remote_port: String, + protocol: String, + pub process_name: String, +} + +lazy_static! { + static ref CONNECTION_REGEX: Regex = Regex::new(r"([^\s]+).*(TCP|UDP).*:(.*)->(.*):(\d*)(\s|$)").unwrap(); +} + +impl RawConnection { + pub fn new(raw_line: &str) -> Option { + let raw_connection_iter = CONNECTION_REGEX.captures_iter(raw_line).filter_map(|cap| { + let process_name = String::from(cap.get(1).unwrap().as_str()); + let protocol = String::from(cap.get(2).unwrap().as_str()); + let local_port = String::from(cap.get(3).unwrap().as_str()); + let ip = String::from(cap.get(4).unwrap().as_str()); + let remote_port = String::from(cap.get(5).unwrap().as_str()); + let connection = RawConnection { process_name, ip, local_port, remote_port, protocol }; + Some(connection) + }); + let raw_connection_vec = raw_connection_iter.map(|m| m).collect::>(); + if raw_connection_vec.is_empty() { + None + } else { + Some(raw_connection_vec[0].clone()) + } + } + + pub fn get_protocol(&self) -> Protocol { + return Protocol::from_string(&self.protocol).unwrap(); + } + + pub fn get_ip_address(&self) -> IpAddr { + return IpAddr::V4(self.ip.parse().unwrap()); + } + + pub fn get_remote_port(&self) -> u16 { + return self.remote_port.parse::().unwrap(); + } + + pub fn get_local_port(&self) -> u16 { + return self.local_port.parse::().unwrap(); + } +} + +pub fn get_connections<'a>() -> RawConnections { + let content = run(&["-n","-P", "-i4"]); + RawConnections::new(content) +} + +fn run<'a, I, S>(args: I) -> String + where I: IntoIterator, S: AsRef +{ + let output = Command::new("lsof") + .args(args) + .output() + .expect("failed to execute process"); + + String::from_utf8(output.stdout).unwrap() +} + +pub struct RawConnections { + content: Vec, +} + +impl RawConnections { + pub fn new(content: String) -> RawConnections { + let lines: Vec = content.lines() + .flat_map(|string| RawConnection::new(string)) + .collect(); + + RawConnections{ content: lines } + } +} + +impl Iterator for RawConnections { + type Item = RawConnection; + + fn next(&mut self) -> Option { + self.content.pop() + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + const LINE_RAW_OUTPUT: &str = "ProcessName 29266 user 39u IPv4 0x28ffb9c0021196bf 0t0 UDP 192.168.0.1:1111->198.252.206.25:2222"; + const FULL_RAW_OUTPUT: &str = r#" +com.apple 590 etoledom 193u IPv4 0x28ffb9c041115627 0t0 TCP 192.168.1.37:60298->31.13.83.36:443 (ESTABLISHED) +com.apple 590 etoledom 198u IPv4 0x28ffb9c04110ea8f 0t0 TCP 192.168.1.37:60299->31.13.83.8:443 (ESTABLISHED) +com.apple 590 etoledom 203u IPv4 0x28ffb9c04110ea8f 0t0 TCP 192.168.1.37:60299->31.13.83.8:443 (ESTABLISHED) +com.apple 590 etoledom 204u IPv4 0x28ffb9c04111253f 0t0 TCP 192.168.1.37:60374->140.82.114.26:443 +"#; + + #[test] + fn test_iterator_multiline() { + let iterator = RawConnections::new(String::from(FULL_RAW_OUTPUT)); + let connections: Vec = iterator.collect(); + assert_eq!(connections.len(), 4); + } + + #[test] + fn test_iterator() { + let iterator = RawConnections::new(String::from(LINE_RAW_OUTPUT)); + let mut worked = false; + for raw_connection in iterator { + worked = true; + assert_eq!(raw_connection.process_name, String::from("ProcessName")); + } + assert!(worked); + } + + #[test] + fn test_raw_connection_is_created_from_raw_output() { + let connection = RawConnection::new(LINE_RAW_OUTPUT); + assert!(connection.is_some()); + } + + #[test] + fn test_raw_connection_is_not_created_from_wrong_raw_output() { + let connection = RawConnection::new("not a process"); + assert!(connection.is_none()); + } + + #[test] + fn test_raw_connection_parse_remote_port() { + let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); + assert_eq!(connection.get_remote_port(), 2222); + } + + #[test] + fn test_raw_connection_parse_local_port() { + let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); + assert_eq!(connection.get_local_port(), 1111); + } + + #[test] + fn test_raw_connection_parse_ip_address() { + let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); + assert_eq!(connection.get_ip_address().to_string(), String::from("198.252.206.25")); + } + + #[test] + fn test_raw_connection_parse_protocol() { + let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); + assert_eq!(connection.get_protocol(), Protocol::Udp); + } + + #[test] + fn test_raw_connection_parse_process_name() { + let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); + assert_eq!(connection.process_name, String::from("ProcessName")); + } +} diff --git a/src/os/macos.rs b/src/os/macos.rs index 0f05e37..2bba346 100644 --- a/src/os/macos.rs +++ b/src/os/macos.rs @@ -8,15 +8,13 @@ use ::termion::input::TermRead; use ::std::collections::HashMap; use ::std::net::IpAddr; -use std::process::Command; -use regex::{Regex}; - use signal_hook::iterator::Signals; -use crate::network::{Connection, Protocol}; +use crate::network::{Connection}; use crate::OsInputOutput; use std::net::{SocketAddr}; +use super::lsof_utils; struct KeyboardEvents; @@ -58,40 +56,18 @@ struct RawConnection { fn get_open_sockets() -> HashMap { let mut open_sockets = HashMap::new(); - let output = Command::new("lsof") - .args(&["-n","-P", "-i4"])//"4tcp" - .output() - .expect("failed to execute process"); - - let regex = Regex::new(r"([^\s]+).*(TCP|UDP).*:(.*)->(.*):(\d*)(\s|$)").unwrap(); - - let output_string = String::from_utf8(output.stdout).unwrap(); - let lines = output_string.lines(); + let connections = lsof_utils::get_connections(); - for line in lines { - let raw_connection_iter = regex.captures_iter(line).filter_map(|cap| { - let process_name = String::from(cap.get(1).unwrap().as_str()); - let protocol = String::from(cap.get(2).unwrap().as_str()); - let local_port = String::from(cap.get(3).unwrap().as_str()); - let ip = String::from(cap.get(4).unwrap().as_str()); - let remote_port = String::from(cap.get(5).unwrap().as_str()); - let connection = RawConnection{process_name, ip,local_port, remote_port, protocol}; - Some(connection) - }); + for raw_connection in connections { + let protocol = raw_connection.get_protocol(); + let ip_address = raw_connection.get_ip_address(); + let remote_port = raw_connection.get_remote_port(); + let local_port = raw_connection.get_local_port(); - let raw_connection_vec = raw_connection_iter.map(|m| m).collect::>(); + let socket_addr = SocketAddr::new(ip_address, remote_port); + let connection = Connection::new(socket_addr, local_port, protocol).unwrap(); - if let Some(raw_connection) = raw_connection_vec.first() { - let protocol = Protocol::from_string(&raw_connection.protocol).unwrap(); - let ip_address = IpAddr::V4(raw_connection.ip.parse().unwrap()); - let remote_port = raw_connection.remote_port.parse::().unwrap(); - let local_port = raw_connection.local_port.parse::().unwrap(); - - let socket_addr = SocketAddr::new(ip_address, remote_port); - let connection = Connection::new(socket_addr, local_port, protocol).unwrap(); - - open_sockets.insert(connection, raw_connection.process_name.clone()); - } + open_sockets.insert(connection, raw_connection.process_name.clone()); } return open_sockets; diff --git a/src/os/mod.rs b/src/os/mod.rs index 0a391f5..f490ae6 100644 --- a/src/os/mod.rs +++ b/src/os/mod.rs @@ -8,4 +8,6 @@ pub use linux::*; mod macos; #[cfg(target_os = "macos")] -pub use macos::*; \ No newline at end of file +pub use macos::*; + +mod lsof_utils; -- cgit v1.2.3