Harden Connections plugin

1 files changed, 59 insertions, 43 deletions

diff --git a/glances/plugins/glances_connections.py b/glances/plugins/glances_connections.py
index cda7a1ba..5cb721cd 100644
--- a/glances/plugins/glances_connections.py
+++ b/glances/plugins/glances_connections.py
@@ -65,7 +65,10 @@ class Plugin(GlancesPlugin):
         # We want to display the stat in the curse interface
         self.display_curse = True
 
-        # @TODO the plugin should be enable only for Linux OS
+        # This plugin is composed if net_connections and nf_conntrack
+        # Enabled by default
+        self.net_connections_enabled = True
+        self.nf_conntrack_enabled = True
 
     @GlancesPlugin._check_decorator
     @GlancesPlugin._log_result_decorator
@@ -81,30 +84,40 @@ class Plugin(GlancesPlugin):
             # Update stats using the PSUtils lib
 
             # Grab network interface stat using the psutil net_connections method
-            try:
-                net_connections = psutil.net_connections(kind="tcp")
-            except Exception as e:
-                logger.debug('Can not get network connections stats ({})'.format(e))
-                return self.stats
-
-            for s in self.status_list:
-                stats[s] = len([c for c in net_connections if c.status == s])
-            initiated = 0
-            for s in self.initiated_states:
-                stats[s] = len([c for c in net_connections if c.status == s])
-                initiated += stats[s]
-            stats['initiated'] = initiated
-            terminated = 0
-            for s in self.initiated_states:
-                stats[s] = len([c for c in net_connections if c.status == s])
-                terminated += stats[s]
-            stats['terminated'] = terminated
-
-            # Grab connections track directly from the /proc file
-            for i in self.conntrack:
-                with open(self.conntrack[i], 'r') as f:
-                    stats[i] = float(f.readline().rstrip("\n"))
-            stats['nf_conntrack_percent'] = stats['nf_conntrack_count'] * 100 / stats['nf_conntrack_max']
+            if self.net_connections_enabled:
+                try:
+                    net_connections = psutil.net_connections(kind="tcp")
+                except Exception as e:
+                    logger.debug('Can not get network connections stats ({})'.format(e))
+                    self.net_connections_enabled = False
+                    self.stats = stats
+                    return self.stats
+
+                for s in self.status_list:
+                    stats[s] = len([c for c in net_connections if c.status == s])
+                initiated = 0
+                for s in self.initiated_states:
+                    stats[s] = len([c for c in net_connections if c.status == s])
+                    initiated += stats[s]
+                stats['initiated'] = initiated
+                terminated = 0
+                for s in self.initiated_states:
+                    stats[s] = len([c for c in net_connections if c.status == s])
+                    terminated += stats[s]
+                stats['terminated'] = terminated
+
+            if self.nf_conntrack_enabled:
+                # Grab connections track directly from the /proc file
+                for i in self.conntrack:
+                    try:
+                        with open(self.conntrack[i], 'r') as f:
+                            stats[i] = float(f.readline().rstrip("\n"))
+                    except IOError as e:
+                        logger.debug('Can not get network connections track ({})'.format(e))
+                        self.nf_conntrack_enabled = False
+                        self.stats = stats
+                        return self.stats
+                stats['nf_conntrack_percent'] = stats['nf_conntrack_count'] * 100 / stats['nf_conntrack_max']
 
         elif self.input_method == 'snmp':
             # Update stats using SNMP
@@ -112,7 +125,6 @@ class Plugin(GlancesPlugin):
 
         # Update the stats
         self.stats = stats
-
         return self.stats
 
     def update_views(self):
@@ -123,7 +135,8 @@ class Plugin(GlancesPlugin):
         # Add specifics informations
         try:
             # Alert and log
-            self.views['nf_conntrack_percent']['decoration'] = self.get_alert(header='nf_conntrack_percent')
+            if self.nf_conntrack_enabled:
+                self.views['nf_conntrack_percent']['decoration'] = self.get_alert(header='nf_conntrack_percent')
         except KeyError:
             # try/except mandatory for Windows compatibility (no conntrack stats)
             pass
@@ -140,25 +153,28 @@ class Plugin(GlancesPlugin):
             return ret
 
         # Header
-        msg = '{}'.format('TCP CONNECTIONS')
-        ret.append(self.curse_add_line(msg, "TITLE"))
+        if self.net_connections_enabled or self.nf_conntrack_enabled:
+            msg = '{}'.format('TCP CONNECTIONS')
+            ret.append(self.curse_add_line(msg, "TITLE"))
         # Connections status
-        for s in [psutil.CONN_LISTEN, 'initiated', psutil.CONN_ESTABLISHED, 'terminated']:
+        if self.net_connections_enabled:
+            for s in [psutil.CONN_LISTEN, 'initiated', psutil.CONN_ESTABLISHED, 'terminated']:
+                ret.append(self.curse_new_line())
+                msg = '{:{width}}'.format(nativestr(s).capitalize(), width=len(s))
+                ret.append(self.curse_add_line(msg))
+                msg = '{:>{width}}'.format(self.stats[s], width=max_width - len(s) + 2)
+                ret.append(self.curse_add_line(msg))
+        # Connections track
+        if self.nf_conntrack_enabled:
+            s = 'Tracked'
             ret.append(self.curse_new_line())
             msg = '{:{width}}'.format(nativestr(s).capitalize(), width=len(s))
             ret.append(self.curse_add_line(msg))
-            msg = '{:>{width}}'.format(self.stats[s], width=max_width - len(s) + 2)
-            ret.append(self.curse_add_line(msg))
-        # Connections track
-        s = 'Tracked'
-        ret.append(self.curse_new_line())
-        msg = '{:{width}}'.format(nativestr(s).capitalize(), width=len(s))
-        ret.append(self.curse_add_line(msg))
-        msg = '{:>{width}}'.format('{:0.0f}/{:0.0f}'.format(self.stats['nf_conntrack_count'],
-                                                            self.stats['nf_conntrack_max']),
-                                   width=max_width - len(s) + 2)
-        ret.append(self.curse_add_line(msg,
-                                       self.get_views(key='nf_conntrack_percent',
-                                                      option='decoration')))
+            msg = '{:>{width}}'.format('{:0.0f}/{:0.0f}'.format(self.stats['nf_conntrack_count'],
+                                                                self.stats['nf_conntrack_max']),
+                                       width=max_width - len(s) + 2)
+            ret.append(self.curse_add_line(msg,
+                                           self.get_views(key='nf_conntrack_percent',
+                                                          option='decoration')))
 
         return ret