Merge pull request #220 from guardicore/feature/netstat

Add netstat

monkey/infection_monkey/system_info/__init__.py

@@ -8,6 +8,7 @@ from enum import IntEnum
 from infection_monkey.network.info import get_host_subnets
 from infection_monkey.system_info.aws_collector import AwsCollector
 from infection_monkey.system_info.azure_cred_collector import AzureCollector
+from infection_monkey.system_info.netstat_collector import NetstatCollector
 
 LOG = logging.getLogger(__name__)
 
@@ -107,12 +108,16 @@ class InfoCollector(object):
     def get_network_info(self):
         """
         Adds network information from the host to the system information.
-        Currently updates with a list of networks accessible from host,
+        Currently updates with netstat and a list of networks accessible from host
-        containing host ip and the subnet range.
+        containing host ip and the subnet range
         :return: None. Updates class information
         """
         LOG.debug("Reading subnets")
-        self.info['network_info'] = {'networks': get_host_subnets()}
+        self.info['network_info'] =\
+            {
+                'networks': get_host_subnets(),
+                'netstat': NetstatCollector.get_netstat_info()
+            }
 
     def get_azure_info(self):
         """

monkey/infection_monkey/system_info/netstat_collector.py

@@ -0,0 +1,44 @@
+# Inspired by Giampaolo Rodola's psutil example from https://github.com/giampaolo/psutil/blob/master/scripts/netstat.py
+
+import logging
+import psutil
+import socket
+
+from socket import AF_INET, SOCK_STREAM, SOCK_DGRAM
+
+__author__ = 'itay.mizeretz'
+
+LOG = logging.getLogger(__name__)
+
+
+class NetstatCollector(object):
+    """
+    Extract netstat info
+    """
+
+    AF_INET6 = getattr(socket, 'AF_INET6', object())
+
+    proto_map = {
+        (AF_INET, SOCK_STREAM): 'tcp',
+        (AF_INET6, SOCK_STREAM): 'tcp6',
+        (AF_INET, SOCK_DGRAM): 'udp',
+        (AF_INET6, SOCK_DGRAM): 'udp6',
+    }
+
+    @staticmethod
+    def get_netstat_info():
+        LOG.info("Collecting netstat info")
+        return [NetstatCollector._parse_connection(c) for c in psutil.net_connections(kind='inet')]
+
+    @staticmethod
+    def _parse_connection(c):
+        return \
+            {
+                'proto': NetstatCollector.proto_map[(c.family, c.type)],
+                'local_address': c.laddr[0],
+                'local_port': c.laddr[1],
+                'remote_address': c.raddr[0] if c.raddr else None,
+                'remote_port': c.raddr[1] if c.raddr else None,
+                'status': c.status,
+                'pid': c.pid
+            }