diff --git a/src/display/components/table.rs b/src/display/components/table.rs index 74e207a..aef914e 100644 --- a/src/display/components/table.rs +++ b/src/display/components/table.rs @@ -64,7 +64,7 @@ impl<'a> Table<'a> { .iter() .map(|(connection, connection_data)| { vec![ - display_connection_string(&connection, &ip_to_host), + display_connection_string(&connection, &ip_to_host, &connection_data.interface), connection_data.process_name.to_string(), display_upload_and_download(*connection_data), ] diff --git a/src/display/ui.rs b/src/display/ui.rs index 3fcb092..fd137bc 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -53,7 +53,11 @@ where write_to_stdout(format!( "connection: <{}> {} up/down Bps: {}/{} process: \"{}\"", timestamp, - display_connection_string(connection, ip_to_host), + display_connection_string( + connection, + ip_to_host, + &connection_network_data.interface + ), connection_network_data.total_bytes_uploaded, connection_network_data.total_bytes_downloaded, connection_network_data.process_name diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index a027ff4..d374c01 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -20,6 +20,7 @@ pub struct ConnectionData { pub total_bytes_downloaded: u128, pub total_bytes_uploaded: u128, pub process_name: String, + pub interface: String, } impl Bandwidth for ConnectionData { @@ -52,7 +53,7 @@ pub struct UIState { impl UIState { pub fn new( connections_to_procs: HashMap, - network_utilization: Utilization, + mut network_utilization: Utilization, ) -> Self { let mut processes: BTreeMap = BTreeMap::new(); let mut remote_addresses: BTreeMap = BTreeMap::new(); @@ -60,32 +61,27 @@ impl UIState { let mut total_bytes_downloaded: u128 = 0; let mut total_bytes_uploaded: u128 = 0; for (connection, process_name) in connections_to_procs { - if let Some(connection_bandwidth_utilization) = - network_utilization.connections.get(&connection) - { + if let Some(connection_info) = network_utilization.connections.remove(&connection) { let data_for_remote_address = remote_addresses .entry(connection.remote_socket.ip) .or_default(); let connection_data = connections.entry(connection).or_default(); let data_for_process = processes.entry(process_name.clone()).or_default(); - data_for_process.total_bytes_downloaded += - &connection_bandwidth_utilization.total_bytes_downloaded; - data_for_process.total_bytes_uploaded += - &connection_bandwidth_utilization.total_bytes_uploaded; + data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded; + data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded; data_for_process.connection_count += 1; - connection_data.total_bytes_downloaded += - &connection_bandwidth_utilization.total_bytes_downloaded; - connection_data.total_bytes_uploaded += - &connection_bandwidth_utilization.total_bytes_uploaded; + connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded; + connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded; connection_data.process_name = process_name; + connection_data.interface = connection_info.interface; data_for_remote_address.total_bytes_downloaded += - connection_bandwidth_utilization.total_bytes_downloaded; + connection_info.total_bytes_downloaded; data_for_remote_address.total_bytes_uploaded += - connection_bandwidth_utilization.total_bytes_uploaded; + connection_info.total_bytes_uploaded; data_for_remote_address.connection_count += 1; - total_bytes_downloaded += connection_bandwidth_utilization.total_bytes_downloaded; - total_bytes_uploaded += connection_bandwidth_utilization.total_bytes_uploaded; + total_bytes_downloaded += connection_info.total_bytes_downloaded; + total_bytes_uploaded += connection_info.total_bytes_uploaded; } } UIState { diff --git a/src/main.rs b/src/main.rs index 9022cfa..34cc09c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,7 +36,7 @@ use structopt::StructOpt; pub struct Opt { #[structopt(short, long)] /// The network interface to listen on, eg. eth0 - interface: String, + interface: Option, #[structopt(short, long)] /// Machine friendlier output raw: bool, @@ -78,8 +78,8 @@ fn try_main() -> Result<(), failure::Error> { } pub struct OsInputOutput { - pub network_interface: NetworkInterface, - pub network_frames: Box, + pub network_interfaces: Vec, + pub network_frames: Vec>, pub get_open_sockets: fn() -> HashMap, pub keyboard_events: Box + Send>, pub dns_client: Option, @@ -105,7 +105,13 @@ where let raw_mode = opts.raw; - let mut sniffer = Sniffer::new(os_input.network_interface, os_input.network_frames); + let mut sniffers = os_input + .network_interfaces + .into_iter() + .zip(os_input.network_frames.into_iter()) + .map(|(iface, frames)| Sniffer::new(iface, frames)) + .collect::>(); + let network_utilization = Arc::new(Mutex::new(Utilization::new())); let ui = Arc::new(Mutex::new(Ui::new(terminal_backend))); @@ -196,10 +202,13 @@ where active_threads.push( thread::Builder::new() .name("sniffing_handler".to_string()) - .spawn(move || { - while running.load(Ordering::Acquire) { + .spawn(move || 'sniffing: loop { + for sniffer in sniffers.iter_mut() { if let Some(segment) = sniffer.next() { - network_utilization.lock().unwrap().update(&segment) + network_utilization.lock().unwrap().update(segment); + } + if !running.load(Ordering::Acquire) { + break 'sniffing; } } }) diff --git a/src/network/connection.rs b/src/network/connection.rs index f931d54..bd42525 100644 --- a/src/network/connection.rs +++ b/src/network/connection.rs @@ -52,9 +52,11 @@ pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap) pub fn display_connection_string( connection: &Connection, ip_to_host: &HashMap, + interface: &str, ) -> String { format!( - ":{} => {}:{} ({})", + "<{}>:{} => {}:{} ({})", + interface, connection.local_port, display_ip_or_host(connection.remote_socket.ip, ip_to_host), connection.remote_socket.port, diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs index 3ec6886..7b2d0b0 100644 --- a/src/network/sniffer.rs +++ b/src/network/sniffer.rs @@ -14,6 +14,7 @@ use ::std::net::{IpAddr, SocketAddr}; use crate::network::{Connection, Protocol}; pub struct Segment { + pub interface: String, pub connection: Connection, pub direction: Direction, pub data_length: u128, @@ -81,6 +82,7 @@ impl Sniffer { } _ => return None, }; + let interface = self.network_interface.name.clone(); let direction = Direction::new(&self.network_interface.ips, &ip_packet); let from = SocketAddr::new(IpAddr::V4(ip_packet.get_source()), source_port); let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port); @@ -90,6 +92,7 @@ impl Sniffer { Direction::Upload => Connection::new(to, source_port, protocol)?, }; Some(Segment { + interface, connection, data_length, direction, diff --git a/src/network/utilization.rs b/src/network/utilization.rs index cb0f12b..6cd732f 100644 --- a/src/network/utilization.rs +++ b/src/network/utilization.rs @@ -3,14 +3,15 @@ use crate::network::{Connection, Direction, Segment}; use ::std::collections::HashMap; #[derive(Clone)] -pub struct TotalBandwidth { +pub struct ConnectionInfo { + pub interface: String, pub total_bytes_downloaded: u128, pub total_bytes_uploaded: u128, } #[derive(Clone)] pub struct Utilization { - pub connections: HashMap, + pub connections: HashMap, } impl Utilization { @@ -23,14 +24,15 @@ impl Utilization { self.connections.clear(); clone } - pub fn update(&mut self, seg: &Segment) { - let total_bandwidth = - self.connections - .entry(seg.connection.clone()) - .or_insert(TotalBandwidth { - total_bytes_downloaded: 0, - total_bytes_uploaded: 0, - }); + pub fn update(&mut self, seg: Segment) { + let total_bandwidth = self + .connections + .entry(seg.connection) + .or_insert(ConnectionInfo { + interface: seg.interface, + total_bytes_downloaded: 0, + total_bytes_uploaded: 0, + }); match seg.direction { Direction::Download => { total_bandwidth.total_bytes_downloaded += seg.data_length; diff --git a/src/os/shared.rs b/src/os/shared.rs index 8947ecb..3a9aeee 100644 --- a/src/os/shared.rs +++ b/src/os/shared.rs @@ -34,11 +34,11 @@ fn get_datalink_channel( interface: &NetworkInterface, ) -> Result, failure::Error> { let mut config = Config::default(); - config.read_timeout = Some(time::Duration::new(0, 1)); + config.read_timeout = Some(time::Duration::new(0, 2_000_000)); match datalink::channel(interface, config) { Ok(Ethernet(_tx, rx)) => Ok(rx), Ok(_) => failure::bail!("Unknown interface type"), - Err(e) => failure::bail!("Failed to listen to network interface: {}", e), + Err(e) => failure::bail!("Failed to listen on network interface: {}", e), } } @@ -76,15 +76,27 @@ fn create_write_to_stdout() -> Box { }) } -pub fn get_input(interface_name: &str, resolve: bool) -> Result { - let keyboard_events = Box::new(KeyboardEvents); - let network_interface = match get_interface(interface_name) { - Some(interface) => interface, - None => { - failure::bail!("Cannot find interface {}", interface_name); +pub fn get_input( + interface_name: &Option, + resolve: bool, +) -> Result { + let network_interfaces = if let Some(name) = interface_name { + match get_interface(&name) { + Some(interface) => vec![interface], + None => { + failure::bail!("Cannot find interface {}", name); + } } + } else { + datalink::interfaces() }; - let network_frames = get_datalink_channel(&network_interface)?; + + let network_frames = network_interfaces + .iter() + .map(|iface| get_datalink_channel(iface)) + .collect::, _>>()?; + + let keyboard_events = Box::new(KeyboardEvents); let write_to_stdout = create_write_to_stdout(); let (on_winch, cleanup) = sigwinch(); let dns_client = if resolve { @@ -96,7 +108,7 @@ pub fn get_input(interface_name: &str, resolve: bool) -> Result>>) -> Box { + pub fn new(packets: Vec>>) -> Box { Box::new(NetworkFrames { packets, current_index: 0, @@ -135,14 +135,14 @@ pub fn get_open_sockets() -> HashMap { open_sockets } -pub fn get_interface() -> NetworkInterface { - NetworkInterface { +pub fn get_interfaces() -> Vec { + vec![NetworkInterface { name: String::from("interface_name"), index: 42, mac: None, ips: vec![IpNetwork::V4("10.0.0.2".parse().unwrap())], flags: 42, - } + }] } pub fn create_fake_on_winch(should_send_winch_event: bool) -> Box {