1 files changed, 34 insertions, 6 deletions

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>,