feat(ui): select windows to display from CLI (#107)

* added first version (no tests)

* fixed layout

* added support for 2 windows

* comments

* breathing windows

* added tests

* format

* fixed rebase

* review

* added more tests

* simplified table creation

* fix(table): do not overflow on max size

Co-authored-by: Aram Drevekenin <aram@poor.dev>

54 files changed, 1046 insertions, 187 deletions

diff --git a/src/display/components/layout.rs b/src/display/components/layout.rs
index 9854e3c..bd1fbfd 100644
--- a/src/display/components/layout.rs
+++ b/src/display/components/layout.rs
@@ -47,19 +47,52 @@ impl<'a> Layout<'a> {
                 layout
             })
     }
-    fn build_layout(&self, rect: Rect) -> Vec<Rect> {
+
+    fn build_two_children_layout(&self, rect: Rect) -> Vec<Rect> {
+        // if there are two elements
+        if rect.height < FIRST_HEIGHT_BREAKPOINT && rect.width < FIRST_WIDTH_BREAKPOINT {
+            //if the space is not enough, we drop one element
+            self.progressive_split(rect, vec![])
+        } else if rect.width < FIRST_WIDTH_BREAKPOINT {
+            // if the horizontal space is not enough, we drop one element and we split horizontally
+            self.progressive_split(rect, vec![Direction::Vertical])
+        } else {
+            // by default we display two elements splitting vertically
+            self.progressive_split(rect, vec![Direction::Horizontal])
+        }
+    }
+
+    fn build_three_children_layout(&self, rect: Rect) -> Vec<Rect> {
+        //if there are three elements
         if rect.height < FIRST_HEIGHT_BREAKPOINT && rect.width < FIRST_WIDTH_BREAKPOINT {
+            //if the space is not enough, we drop two elements
             self.progressive_split(rect, vec![])
         } else if rect.height < FIRST_HEIGHT_BREAKPOINT {
+            // if the vertical space is not enough, we drop one element and we split vertically
             self.progressive_split(rect, vec![Direction::Horizontal])
         } else if rect.width < FIRST_WIDTH_BREAKPOINT {
+            // if the horizontal space is not enough, we drop one element and we split horizontally
             self.progressive_split(rect, vec![Direction::Vertical])
         } else if rect.width < SECOND_WIDTH_BREAKPOINT {
+            // if the horizontal space is not enough for the default layout, we display one wide element
+            // on top and we split horizontally the bottom
             self.progressive_split(rect, vec![Direction::Vertical, Direction::Horizontal])
         } else {
+            // default layout
             self.progressive_split(rect, vec![Direction::Horizontal, Direction::Vertical])
         }
     }
+
+    fn build_layout(&self, rect: Rect) -> Vec<Rect> {
+        if self.children.len() == 1 {
+            // if there's only one element to render, it can take the whole frame
+            vec![rect]
+        } else if self.children.len() == 2 {
+            self.build_two_children_layout(rect)
+        } else {
+            self.build_three_children_layout(rect)
+        }
+    }
     pub fn render(&self, frame: &mut Frame<impl Backend>, rect: Rect) {
         let (top, app, bottom) = top_app_and_bottom_split(rect);
         let layout_slots = self.build_layout(app);
diff --git a/src/display/components/table.rs b/src/display/components/table.rs
index ff5a148..1ca56d7 100644
--- a/src/display/components/table.rs
+++ b/src/display/components/table.rs
@@ -16,8 +16,11 @@ const FIRST_WIDTH_BREAKPOINT: u16 = 50;
 const SECOND_WIDTH_BREAKPOINT: u16 = 71;
 const THIRD_WIDTH_BREAKPOINT: u16 = 95;
 
+const MAX_FIRST_COLUMN_WIDTH_PERCENTAGE: u16 = 53;
+const MAX_SECOND_COLUMN_WIDTH_PERCENTAGE: u16 = 21;
+const MAX_THIRD_COLUMN_WIDTH_PERCENTAGE: u16 = 22;
+
 const FIRST_COLUMN_WIDTHS: [u16; 4] = [10, 30, 40, 50];
-const SECOND_COLUMN_WIDTHS: [u16; 1] = [20];
 const THIRD_COLUMN_WIDTHS: [u16; 4] = [20, 20, 20, 20];
 
 fn display_upload_and_download(bandwidth: &impl Bandwidth) -> String {
@@ -137,9 +140,9 @@ impl<'a> Table<'a> {
             vec![FIRST_COLUMN_WIDTHS[2], THIRD_COLUMN_WIDTHS[2]]
         } else {
             vec![
-                FIRST_COLUMN_WIDTHS[3],
-                SECOND_COLUMN_WIDTHS[0],
-                THIRD_COLUMN_WIDTHS[2],
+                rect.width * MAX_FIRST_COLUMN_WIDTH_PERCENTAGE / 100,
+                rect.width * MAX_SECOND_COLUMN_WIDTH_PERCENTAGE / 100,
+                rect.width * MAX_THIRD_COLUMN_WIDTH_PERCENTAGE / 100 - 1,
             ]
         };
 
diff --git a/src/display/ui.rs b/src/display/ui.rs
index 2c44d23..4b193d5 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -9,6 +9,7 @@ use crate::network::{display_connection_string, display_ip_or_host, LocalSocket,
 
 use ::std::net::IpAddr;
 
+use crate::RenderOpts;
 use chrono::prelude::*;
 
 pub struct Ui<B>
@@ -18,13 +19,14 @@ where
     terminal: Terminal<B>,
     state: UIState,
     ip_to_host: HashMap<IpAddr, String>,
+    opts: RenderOpts,
 }
 
 impl<B> Ui<B>
 where
     B: Backend,
 {
-    pub fn new(terminal_backend: B) -> Self {
+    pub fn new(terminal_backend: B, opts: RenderOpts) -> Self {
         let mut terminal = Terminal::new(terminal_backend).unwrap();
         terminal.clear().unwrap();
         terminal.hide_cursor().unwrap();
@@ -32,6 +34,7 @@ where
             terminal,
             state: Default::default(),
             ip_to_host: Default::default(),
+            opts,
         }
     }
     pub fn output_text(&mut self, write_to_stdout: &mut (dyn FnMut(String) + Send)) {
@@ -76,13 +79,10 @@ where
     }
     pub fn draw(&mut self, paused: bool) {
         let state = &self.state;
-        let ip_to_host = &self.ip_to_host;
+        let children = self.get_tables_to_display();
         self.terminal
             .draw(|mut frame| {
                 let size = frame.size();
-                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,
                     paused,
@@ -90,13 +90,41 @@ where
                 let help_text = HelpText { paused };
                 let layout = Layout {
                     header: total_bandwidth,
-                    children: vec![processes, connections, remote_addresses],
+                    children,
                     footer: help_text,
                 };
                 layout.render(&mut frame, size);
             })
             .unwrap();
     }
+
+    fn get_tables_to_display(&self) -> Vec<Table<'static>> {
+        let opts = &self.opts;
+        let mut children: Vec<Table> = Vec::new();
+        if opts.processes {
+            children.push(Table::create_processes_table(&self.state));
+        }
+        if opts.addresses {
+            children.push(Table::create_remote_addresses_table(
+                &self.state,
+                &self.ip_to_host,
+            ));
+        }
+        if opts.connections {
+            children.push(Table::create_connections_table(
+                &self.state,
+                &self.ip_to_host,
+            ));
+        }
+        if !(opts.processes || opts.addresses || opts.connections) {
+            children = vec![
+                Table::create_processes_table(&self.state),
+                Table::create_connections_table(&self.state, &self.ip_to_host),
+                Table::create_remote_addresses_table(&self.state, &self.ip_to_host),
+            ];
+        }
+        children
+    }
     pub fn update_state(
         &mut self,
         connections_to_procs: HashMap<LocalSocket, String>,
diff --git a/src/main.rs b/src/main.rs
index 30a9b25..1e35289 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,7 +28,6 @@ use ::std::io;
 use ::std::time::Instant;
 use ::termion::raw::IntoRawMode;
 use ::tui::backend::TermionBackend;
-
 use structopt::StructOpt;
 
 const DISPLAY_DELTA: time::Duration = time::Duration::from_millis(1000);
@@ -45,6 +44,15 @@ pub struct Opt {
     #[structopt(short, long)]
     /// Do not attempt to resolve IPs to their hostnames
     no_resolve: bool,
+    #[structopt(flatten)]
+    render_opts: RenderOpts,
+}
+
+#[derive(StructOpt, Debug)]
+pub struct RenderOpts {
+    processes: bool,
+    connections: bool,
+    addresses: bool,
 }
 
 fn main() {
@@ -114,7 +122,7 @@ where
     let raw_mode = opts.raw;
 
     let network_utilization = Arc::new(Mutex::new(Utilization::new()));
-    let ui = Arc::new(Mutex::new(Ui::new(terminal_backend)));
+    let ui = Arc::new(Mutex::new(Ui::new(terminal_backend, opts.render_opts)));
 
     if !raw_mode {
         active_threads.push(
diff --git a/src/tests/cases/raw_mode.rs b/src/tests/cases/raw_mode.rs
index 03c3684..14fae65 100644
--- a/src/tests/cases/raw_mode.rs
+++ b/src/tests/cases/raw_mode.rs
@@ -10,31 +10,12 @@ 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 crate::tests::cases::test_utils::{
-    opts_raw, os_input_output_dns, os_input_output_stdout, test_backend_factory,
+    build_tcp_packet, opts_raw, os_input_output_dns, os_input_output_stdout, test_backend_factory,
 };
 
-use crate::{start, Opt};
-
-fn build_tcp_packet(
-    source_ip: &str,
-    destination_ip: &str,
-    source_port: u16,
-    destination_port: u16,
-    payload: &'static [u8],
-) -> Vec<u8> {
-    let mut pkt_buf = [0u8; 1500];
-    let pkt = packet_builder!(
-         pkt_buf,
-         ether({set_destination => MacAddr(0,0,0,0,0,0), set_source => MacAddr(0,0,0,0,0,0)}) /
-         ipv4({set_source => ipv4addr!(source_ip), set_destination => ipv4addr!(destination_ip) }) /
-         tcp({set_source => source_port, set_destination => destination_port }) /
-         payload(payload)
-    );
-    pkt.packet().to_vec()
-}
+use crate::{start, Opt, RenderOpts};
 
 fn build_ip_tcp_packet(
     source_ip: &str,
@@ -595,6 +576,11 @@ fn no_resolve_mode() {
         interface: Some(String::from("interface_name")),
         raw: true,
         no_resolve: true,
+        render_opts: RenderOpts {
+            addresses: false,
+            connections: false,
+            processes: false,
+        },
     };
     start(backend, os_input, opts);
     let stdout = Arc::try_unwrap(stdout).unwrap().into_inner().unwrap();
diff --git a/src/tests/cases/snapshots/ui__basic_only_addresses.snap b/src/tests/cases/snapshots/ui__basic_only_addresses.snap
new file mode 100644
index 0000000..c3c5e53
--- /dev/null
+++ b/src/tests/cases/snapshots/ui__basic_only_addresses.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 remote address───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
+│Remote Address                                                                                        Connection Count                         Rate Up / Down                               │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+└────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
+ Press <SPACE> to pause.                                                                                                                                                                      
+
diff --git a/src/tests/cases/snapshots/ui__basic_only_connections.snap b/src/tests/cases/snapshots/ui__basic_only_connections.snap
new file mode 100644
index 0000000..27f8036
--- /dev/null
+++ b/src/tests/cases/snapshots/ui__basic_only_connections.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 connection───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
+│Connection                                                                                            Process                                  Rate Up / Down                               │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│                                                                                                                                                                                            │
+│