diff --git a/src/display/components/total_bandwidth.rs b/src/display/components/total_bandwidth.rs index cd6a6f285..ba78d2dd9 100644 --- a/src/display/components/total_bandwidth.rs +++ b/src/display/components/total_bandwidth.rs @@ -8,18 +8,29 @@ use crate::display::{DisplayBandwidth, UIState}; pub struct TotalBandwidth<'a> { pub state: &'a UIState, + pub paused: bool, } impl<'a> TotalBandwidth<'a> { pub fn render(&self, frame: &mut Frame, rect: Rect) { - let title_text = [Text::styled( - format!( - " Total Rate Up / Down: {} / {}", - DisplayBandwidth(self.state.total_bytes_uploaded as f64), - DisplayBandwidth(self.state.total_bytes_downloaded as f64) - ), - Style::default().fg(Color::Green).modifier(Modifier::BOLD), - )]; + let title_text = { + let paused_str = if self.paused { "[PAUSED]" } else { "" }; + let color = if self.paused { + Color::Yellow + } else { + Color::Green + }; + + [Text::styled( + format!( + " Total Rate Up / Down: {} / {} {}", + DisplayBandwidth(self.state.total_bytes_uploaded as f64), + DisplayBandwidth(self.state.total_bytes_downloaded as f64), + paused_str + ), + Style::default().fg(color).modifier(Modifier::BOLD), + )] + }; Paragraph::new(title_text.iter()) .block(Block::default().borders(Borders::NONE)) .alignment(Alignment::Left) diff --git a/src/display/ui.rs b/src/display/ui.rs index e61f5e7af..e3464bb5c 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -74,7 +74,7 @@ where )); } } - pub fn draw(&mut self) { + pub fn draw(&mut self, paused: bool) { let state = &self.state; let ip_to_host = &self.ip_to_host; self.terminal @@ -83,7 +83,10 @@ where let connections = Table::create_connections_table(&state, &ip_to_host); let processes = Table::create_processes_table(&state); let remote_addresses = Table::create_remote_addresses_table(&state, &ip_to_host); - let total_bandwidth = TotalBandwidth { state: &state }; + let total_bandwidth = TotalBandwidth { + state: &state, + paused, + }; let layout = Layout { header: total_bandwidth, children: vec![processes, connections, remote_addresses], diff --git a/src/main.rs b/src/main.rs index ea40de00a..30a9b25f1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -100,6 +100,7 @@ where B: Backend + Send + 'static, { let running = Arc::new(AtomicBool::new(true)); + let paused = Arc::new(AtomicBool::new(false)); let mut active_threads = vec![]; @@ -121,11 +122,12 @@ where .name("resize_handler".to_string()) .spawn({ let ui = ui.clone(); + let paused = paused.clone(); move || { on_winch({ Box::new(move || { let mut ui = ui.lock().unwrap(); - ui.draw(); + ui.draw(paused.load(Ordering::SeqCst)); }) }); } @@ -138,6 +140,7 @@ where .name("display_handler".to_string()) .spawn({ let running = running.clone(); + let paused = paused.clone(); let network_utilization = network_utilization.clone(); move || { while running.load(Ordering::Acquire) { @@ -160,11 +163,14 @@ where } { let mut ui = ui.lock().unwrap(); - ui.update_state(sockets_to_procs, utilization, ip_to_host); + let paused = paused.load(Ordering::SeqCst); + if !paused { + ui.update_state(sockets_to_procs, utilization, ip_to_host); + } if raw_mode { ui.output_text(&mut write_to_stdout); } else { - ui.draw(); + ui.draw(paused); } } let render_duration = render_start_time.elapsed(); @@ -195,6 +201,10 @@ where display_handler.unpark(); break; } + Event::Key(Key::Char(' ')) => { + paused.fetch_xor(true, Ordering::SeqCst); + display_handler.unpark(); + } _ => (), }; } diff --git a/src/tests/cases/snapshots/ui__pause_by_space-2.snap b/src/tests/cases/snapshots/ui__pause_by_space-2.snap new file mode 100644 index 000000000..204efd1c2 --- /dev/null +++ b/src/tests/cases/snapshots/ui__pause_by_space-2.snap @@ -0,0 +1,55 @@ +--- +source: src/tests/cases/ui.rs +expression: "&terminal_draw_events_mirror[1]" +--- + Total Rate Up / Down: 0Bps / 0Bps [PAUSED] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/tests/cases/snapshots/ui__pause_by_space-3.snap b/src/tests/cases/snapshots/ui__pause_by_space-3.snap new file mode 100644 index 000000000..d710afbc2 --- /dev/null +++ b/src/tests/cases/snapshots/ui__pause_by_space-3.snap @@ -0,0 +1,55 @@ +--- +source: src/tests/cases/ui.rs +expression: "&terminal_draw_events_mirror[2]" +--- + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/tests/cases/snapshots/ui__pause_by_space-4.snap b/src/tests/cases/snapshots/ui__pause_by_space-4.snap new file mode 100644 index 000000000..5488196df --- /dev/null +++ b/src/tests/cases/snapshots/ui__pause_by_space-4.snap @@ -0,0 +1,55 @@ +--- +source: src/tests/cases/ui.rs +expression: "&terminal_draw_events_mirror[3]" +--- + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/tests/cases/snapshots/ui__pause_by_space-5.snap b/src/tests/cases/snapshots/ui__pause_by_space-5.snap new file mode 100644 index 000000000..1524b669b --- /dev/null +++ b/src/tests/cases/snapshots/ui__pause_by_space-5.snap @@ -0,0 +1,55 @@ +--- +source: src/tests/cases/ui.rs +expression: "&terminal_draw_events_mirror[4]" +--- + Total Rate Up / Down: 0Bps / 0Bps + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/tests/cases/snapshots/ui__pause_by_space.snap b/src/tests/cases/snapshots/ui__pause_by_space.snap new file mode 100644 index 000000000..ec8d97c42 --- /dev/null +++ b/src/tests/cases/snapshots/ui__pause_by_space.snap @@ -0,0 +1,55 @@ +--- +source: src/tests/cases/ui.rs +expression: "&terminal_draw_events_mirror[0]" +--- + Total Rate Up / Down: 0Bps / 0Bps +┌Utilization by process name──────────────────────────────────────────────────────────────────┐┌Utilization by connection────────────────────────────────────────────────────────────────────┐ +│Process Connection count Rate Up / Down ││Connection Process Rate Up / Down │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ │└─────────────────────────────────────────────────────────────────────────────────────────────┘ +│ │┌Utilization by remote address────────────────────────────────────────────────────────────────┐ +│ ││Remote Address Connection Count Rate Up / Down │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +│ ││ │ +└─────────────────────────────────────────────────────────────────────────────────────────────┘└─────────────────────────────────────────────────────────────────────────────────────────────┘ + diff --git a/src/tests/cases/test_utils.rs b/src/tests/cases/test_utils.rs index 22e6c6341..18440cc1c 100644 --- a/src/tests/cases/test_utils.rs +++ b/src/tests/cases/test_utils.rs @@ -24,9 +24,9 @@ pub fn os_input_output( ) -> OsInputOutput { os_input_output_factory( network_frames, - sleep_num, None, create_fake_dns_client(HashMap::new()), + sleep_and_quit_events(sleep_num), ) } pub fn os_input_output_stdout( @@ -36,9 +36,9 @@ pub fn os_input_output_stdout( ) -> OsInputOutput { os_input_output_factory( network_frames, - sleep_num, stdout, create_fake_dns_client(HashMap::new()), + sleep_and_quit_events(sleep_num), ) } @@ -48,14 +48,19 @@ pub fn os_input_output_dns( stdout: Option>>>, dns_client: Option, ) -> OsInputOutput { - os_input_output_factory(network_frames, sleep_num, stdout, dns_client) + os_input_output_factory( + network_frames, + stdout, + dns_client, + sleep_and_quit_events(sleep_num), + ) } -fn os_input_output_factory( +pub fn os_input_output_factory( network_frames: Vec>, - sleep_num: usize, stdout: Option>>>, dns_client: Option, + keyboard_events: Box + Send>, ) -> OsInputOutput { let on_winch = create_fake_on_winch(false); let cleanup = Box::new(|| {}); @@ -74,7 +79,7 @@ fn os_input_output_factory( network_interfaces: get_interfaces(), network_frames, get_open_sockets, - keyboard_events: sleep_and_quit_events(sleep_num), + keyboard_events, dns_client, on_winch, cleanup, diff --git a/src/tests/cases/ui.rs b/src/tests/cases/ui.rs index cd5aad90b..c5addcf8b 100644 --- a/src/tests/cases/ui.rs +++ b/src/tests/cases/ui.rs @@ -9,13 +9,17 @@ use ::std::collections::HashMap; use ::std::net::IpAddr; use crate::tests::cases::test_utils::{ - opts_ui, os_input_output, sleep_and_quit_events, test_backend_factory, + opts_ui, os_input_output, os_input_output_factory, sleep_and_quit_events, test_backend_factory, }; +use ::termion::event::{Event, Key}; use packet_builder::payload::PayloadData; use packet_builder::*; use pnet_bandwhich_fork::datalink::DataLinkReceiver; use pnet_bandwhich_fork::packet::Packet; use pnet_base::MacAddr; +use std::iter; + +use crate::tests::fakes::KeyboardEvents; use crate::{start, Opt, OsInputOutput}; @@ -59,6 +63,55 @@ fn basic_startup() { assert_snapshot!(&terminal_draw_events_mirror[0]); } +#[test] +fn pause_by_space() { + let network_frames = vec![NetworkFrames::new(vec![ + Some(build_tcp_packet( + "1.1.1.1", + "10.0.0.2", + 12345, + 443, + b"I have come from 1.1.1.1", + )), + None, // sleep + None, // sleep + None, // sleep + Some(build_tcp_packet( + "1.1.1.1", + "10.0.0.2", + 12345, + 443, + b"Same here, but one second later", + )), + ]) as Box]; + + // sleep for 1s, then press space, sleep for 2s, then quit + let mut events: Vec> = iter::repeat(None).take(1).collect(); + events.push(Some(Event::Key(Key::Char(' ')))); + events.push(None); + events.push(None); + events.push(Some(Event::Key(Key::Char(' ')))); + events.push(Some(Event::Key(Key::Ctrl('c')))); + + let events = Box::new(KeyboardEvents::new(events)); + let os_input = os_input_output_factory(network_frames, None, None, events); + let (terminal_events, terminal_draw_events, backend) = test_backend_factory(190, 50); + let opts = opts_ui(); + start(backend, os_input, opts); + let terminal_draw_events_mirror = terminal_draw_events.lock().unwrap(); + let expected_terminal_events = vec![ + Clear, HideCursor, Draw, Flush, Draw, Flush, Draw, Flush, Clear, ShowCursor, + ]; + assert_eq!( + &terminal_events.lock().unwrap()[..], + &expected_terminal_events[..] + ); + assert_eq!(terminal_draw_events_mirror.len(), 3); + assert_snapshot!(&terminal_draw_events_mirror[0]); + assert_snapshot!(&terminal_draw_events_mirror[1]); + assert_snapshot!(&terminal_draw_events_mirror[2]); +} + #[test] fn one_packet_of_traffic() { let network_frames = vec![NetworkFrames::new(vec![Some(build_tcp_packet(