From 841ad289736813aaf0ad54b7f9ff9f91d4ab2219 Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Sun, 2 Dec 2018 16:12:10 +0200
Subject: [PATCH] Add netstat

---
 .../infection_monkey/system_info/__init__.py  | 11 +++--
 .../system_info/netstat_collector.py          | 44 +++++++++++++++++++
 .../system_info/windows_info_collector.py     |  2 +-
 3 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 monkey/infection_monkey/system_info/netstat_collector.py

diff --git a/monkey/infection_monkey/system_info/__init__.py b/monkey/infection_monkey/system_info/__init__.py
index e3892abac..56d7fca8b 100644
--- a/monkey/infection_monkey/system_info/__init__.py
+++ b/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,
-        containing host ip and the subnet range.
+        Currently updates with netstat and a list of networks accessible from host
+        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):
         """
diff --git a/monkey/infection_monkey/system_info/netstat_collector.py b/monkey/infection_monkey/system_info/netstat_collector.py
new file mode 100644
index 000000000..361bf0d81
--- /dev/null
+++ b/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
+            }
diff --git a/monkey/infection_monkey/system_info/windows_info_collector.py b/monkey/infection_monkey/system_info/windows_info_collector.py
index fb2261572..ced13de4e 100644
--- a/monkey/infection_monkey/system_info/windows_info_collector.py
+++ b/monkey/infection_monkey/system_info/windows_info_collector.py
@@ -2,7 +2,7 @@ import os
 import logging
 import sys
 
-sys.coinit_flags = 0 # needed for proper destruction of the wmi python module
+sys.coinit_flags = 0  # needed for proper destruction of the wmi python module
 
 import infection_monkey.config
 from infection_monkey.system_info.mimikatz_collector import MimikatzCollector