Merge pull request #478 from guardicore/feature/refactor_fingerprinting

Feature/refactor fingerprinting

monkey/infection_monkey/config.py

@@ -33,9 +33,6 @@ class Configuration(object):
             if self._depth_from_commandline and key == "depth":
                 continue
             # handle in cases
-            if key == 'finger_classes':
-                class_objects = [getattr(network_import, val) for val in value]
-                setattr(self, key, class_objects)
             elif key == 'exploiter_classes':
                 class_objects = [getattr(exploit_import, val) for val in value]
                 setattr(self, key, class_objects)

monkey/infection_monkey/exploit/smbexec.py

@@ -7,7 +7,7 @@ from infection_monkey.exploit import HostExploiter
 from infection_monkey.exploit.tools.helpers import get_target_monkey, get_monkey_depth, build_monkey_commandline
 from infection_monkey.exploit.tools.smb_tools import SmbTools
 from infection_monkey.model import MONKEY_CMDLINE_DETACHED_WINDOWS, DROPPER_CMDLINE_DETACHED_WINDOWS
-from infection_monkey.network import SMBFinger
+from infection_monkey.network.smbfinger import SMBFinger
 from infection_monkey.network.tools import check_tcp_port
 from common.utils.exploit_enum import ExploitType
 from infection_monkey.telemetry.attack.t1035_telem import T1035Telem

monkey/infection_monkey/exploit/win_ms08_067.py

@@ -17,7 +17,7 @@ from impacket.dcerpc.v5 import transport
 from infection_monkey.exploit.tools.helpers import get_target_monkey, get_monkey_depth, build_monkey_commandline
 from infection_monkey.exploit.tools.smb_tools import SmbTools
 from infection_monkey.model import DROPPER_CMDLINE_WINDOWS, MONKEY_CMDLINE_WINDOWS
-from infection_monkey.network import SMBFinger
+from infection_monkey.network.smbfinger import SMBFinger
 from infection_monkey.network.tools import check_tcp_port
 from . import HostExploiter
 
@@ -234,7 +234,8 @@ class Ms08_067_Exploiter(HostExploiter):
         # execute the remote dropper in case the path isn't final
         if remote_full_path.lower() != self._config.dropper_target_path_win_32.lower():
             cmdline = DROPPER_CMDLINE_WINDOWS % {'dropper_path': remote_full_path} + \
-                      build_monkey_commandline(self.host, get_monkey_depth() - 1, self._config.dropper_target_path_win_32)
+                      build_monkey_commandline(self.host, get_monkey_depth() - 1,
+                                               self._config.dropper_target_path_win_32)
         else:
             cmdline = MONKEY_CMDLINE_WINDOWS % {'monkey_path': remote_full_path} + \
                       build_monkey_commandline(self.host, get_monkey_depth() - 1)

monkey/infection_monkey/monkey.py

@@ -6,6 +6,7 @@ import sys
 import time
 
 import infection_monkey.tunnel as tunnel
+from infection_monkey.network.HostFinger import HostFinger
 from infection_monkey.utils.monkey_dir import create_monkey_dir, get_monkey_dir_path, remove_monkey_dir
 from infection_monkey.utils.monkey_log_path import get_monkey_log_path
 from infection_monkey.utils.environment import is_windows_os
@@ -145,7 +146,7 @@ class InfectionMonkey(object):
 
             self._exploiters = WormConfiguration.exploiter_classes
 
-            self._fingerprint = [fingerprint() for fingerprint in WormConfiguration.finger_classes]
+            self._fingerprint = HostFinger.get_instances()
 
             if not self._keep_running or not WormConfiguration.alive:
                 break
@@ -192,9 +193,7 @@ class InfectionMonkey(object):
                     self._exploiters = sorted(self._exploiters, key=lambda exploiter_: exploiter_.EXPLOIT_TYPE.value)
                     host_exploited = False
                     for exploiter in [exploiter(machine) for exploiter in self._exploiters]:
-
                         if self.try_exploiting(machine, exploiter):
-
                             host_exploited = True
                             VictimHostTelem('T1210', ScanStatus.USED, machine=machine).send()
                             break

monkey/infection_monkey/network/HostFinger.py

@@ -0,0 +1,33 @@
+from abc import abstractmethod
+
+from infection_monkey.config import WormConfiguration
+from infection_monkey.utils.plugins.plugin import Plugin
+import infection_monkey.network
+
+
+class HostFinger(Plugin):
+    @staticmethod
+    def base_package_file():
+        return infection_monkey.network.__file__
+
+    @staticmethod
+    def base_package_name():
+        return infection_monkey.network.__package__
+
+    @property
+    @abstractmethod
+    def _SCANNED_SERVICE(self):
+        pass
+
+    def init_service(self, services, service_key, port):
+        services[service_key] = {}
+        services[service_key]['display_name'] = self._SCANNED_SERVICE
+        services[service_key]['port'] = port
+
+    @abstractmethod
+    def get_host_fingerprint(self, host):
+        raise NotImplementedError()
+
+    @staticmethod
+    def should_run(class_name: str) -> bool:
+        return class_name in WormConfiguration.finger_classes

monkey/infection_monkey/network/HostScanner.py

@@ -0,0 +1,8 @@
+from abc import ABCMeta, abstractmethod
+
+
+class HostScanner(metaclass=ABCMeta):
+    @property
+    @abstractmethod
+    def is_host_alive(self, host):
+        raise NotImplementedError()

monkey/infection_monkey/network/__init__.py

@@ -1,36 +1 @@
-from abc import ABCMeta, abstractmethod
-
 __author__ = 'itamar'
-
-
-class HostScanner(object, metaclass=ABCMeta):
-    @abstractmethod
-    def is_host_alive(self, host):
-        raise NotImplementedError()
-
-
-class HostFinger(object, metaclass=ABCMeta):
-    @property
-    @abstractmethod
-    def _SCANNED_SERVICE(self):
-        pass
-
-    def init_service(self, services, service_key, port):
-        services[service_key] = {}
-        services[service_key]['display_name'] = self._SCANNED_SERVICE
-        services[service_key]['port'] = port
-
-    @abstractmethod
-    def get_host_fingerprint(self, host):
-        raise NotImplementedError()
-
-
-from infection_monkey.network.ping_scanner import PingScanner
-from infection_monkey.network.tcp_scanner import TcpScanner
-from infection_monkey.network.smbfinger import SMBFinger
-from infection_monkey.network.sshfinger import SSHFinger
-from infection_monkey.network.httpfinger import HTTPFinger
-from infection_monkey.network.elasticfinger import ElasticFinger
-from infection_monkey.network.mysqlfinger import MySQLFinger
-from infection_monkey.network.info import local_ips, get_free_tcp_port
-from infection_monkey.network.mssql_fingerprint import MSSQLFinger

monkey/infection_monkey/network/elasticfinger.py

@@ -6,9 +6,8 @@ import requests
 from requests.exceptions import Timeout, ConnectionError
 
 import infection_monkey.config
+from infection_monkey.network.HostFinger import HostFinger
 from common.data.network_consts import ES_SERVICE
-from infection_monkey.model.host import VictimHost
-from infection_monkey.network import HostFinger
 
 ES_PORT = 9200
 ES_HTTP_TIMEOUT = 5
@@ -31,7 +30,6 @@ class ElasticFinger(HostFinger):
         :param host:
         :return: Success/failure, data is saved in the host struct
         """
-        assert isinstance(host, VictimHost)
         try:
             url = 'http://%s:%s/' % (host.ip_addr, ES_PORT)
             with closing(requests.get(url, timeout=ES_HTTP_TIMEOUT)) as req:

monkey/infection_monkey/network/httpfinger.py

@@ -1,6 +1,5 @@
 import infection_monkey.config
-from infection_monkey.network import HostFinger
+from infection_monkey.network.HostFinger import HostFinger
-from infection_monkey.model.host import VictimHost
 import logging
 
 LOG = logging.getLogger(__name__)
@@ -21,7 +20,6 @@ class HTTPFinger(HostFinger):
         pass
 
     def get_host_fingerprint(self, host):
-        assert isinstance(host, VictimHost)
         from requests import head
         from requests.exceptions import Timeout, ConnectionError
         from contextlib import closing

monkey/infection_monkey/network/mssql_fingerprint.py

@@ -2,8 +2,7 @@ import errno
 import logging
 import socket
 
-from infection_monkey.model.host import VictimHost
+from infection_monkey.network.HostFinger import HostFinger
-from infection_monkey.network import HostFinger
 import infection_monkey.config
 
 __author__ = 'Maor Rayzin'
@@ -30,7 +29,6 @@ class MSSQLFinger(HostFinger):
                 Discovered server information written to the Host info struct.
                 True if success, False otherwise.
         """
-        assert isinstance(host, VictimHost)
 
         # Create a UDP socket and sets a timeout
         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

monkey/infection_monkey/network/mysqlfinger.py

@@ -2,8 +2,7 @@ import logging
 import socket
 
 import infection_monkey.config
-from infection_monkey.model.host import VictimHost
+from infection_monkey.network.HostFinger import HostFinger
-from infection_monkey.network import HostFinger
 from infection_monkey.network.tools import struct_unpack_tracker, struct_unpack_tracker_string
 
 MYSQL_PORT = 3306
@@ -28,7 +27,6 @@ class MySQLFinger(HostFinger):
         :param host:
         :return: Success/failure, data is saved in the host struct
         """
-        assert isinstance(host, VictimHost)
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         s.settimeout(self.SOCKET_TIMEOUT)
 

monkey/infection_monkey/network/network_scanner.py

@@ -6,7 +6,8 @@ from common.network.network_range import NetworkRange
 from infection_monkey.config import WormConfiguration
 from infection_monkey.model.victim_host_generator import VictimHostGenerator
 from infection_monkey.network.info import local_ips, get_interfaces_ranges
-from infection_monkey.network import TcpScanner, PingScanner
+from infection_monkey.network.tcp_scanner import TcpScanner
+from infection_monkey.network.ping_scanner import PingScanner
 
 LOG = logging.getLogger(__name__)
 

monkey/infection_monkey/network/ping_scanner.py

@@ -5,8 +5,9 @@ import subprocess
 import sys
 
 import infection_monkey.config
+from infection_monkey.network.HostFinger import HostFinger
+from infection_monkey.network.HostScanner import HostScanner
 from infection_monkey.model.host import VictimHost
-from infection_monkey.network import HostScanner, HostFinger
 
 __author__ = 'itamar'
 
@@ -28,7 +29,6 @@ class PingScanner(HostScanner, HostFinger):
         self._ttl_regex = re.compile(TTL_REGEX_STR, re.IGNORECASE)
 
     def is_host_alive(self, host):
-        assert isinstance(host, VictimHost)
 
         timeout = self._config.ping_scan_timeout
         if not "win32" == sys.platform:
@@ -42,7 +42,6 @@ class PingScanner(HostScanner, HostFinger):
                                     stderr=self._devnull)
 
     def get_host_fingerprint(self, host):
-        assert isinstance(host, VictimHost)
 
         timeout = self._config.ping_scan_timeout
         if not "win32" == sys.platform:

monkey/infection_monkey/network/smbfinger.py

@@ -3,8 +3,7 @@ import struct
 import logging
 from odict import odict
 
-from infection_monkey.network import HostFinger
+from infection_monkey.network.HostFinger import HostFinger
-from infection_monkey.model.host import VictimHost
 
 SMB_PORT = 445
 SMB_SERVICE = 'tcp-445'
@@ -114,7 +113,6 @@ class SMBFinger(HostFinger):
         self._config = WormConfiguration
 
     def get_host_fingerprint(self, host):
-        assert isinstance(host, VictimHost)
 
         try:
             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

monkey/infection_monkey/network/sshfinger.py

@@ -1,8 +1,7 @@
 import re
 
 import infection_monkey.config
-from infection_monkey.model.host import VictimHost
+from infection_monkey.network.HostFinger import HostFinger
-from infection_monkey.network import HostFinger
 from infection_monkey.network.tools import check_tcp_port
 
 SSH_PORT = 22
@@ -34,7 +33,6 @@ class SSHFinger(HostFinger):
                 break
 
     def get_host_fingerprint(self, host):
-        assert isinstance(host, VictimHost)
 
         for name, data in list(host.services.items()):
             banner = data.get('banner', '')

monkey/infection_monkey/network/tcp_scanner.py

@@ -2,7 +2,8 @@ from itertools import zip_longest
 from random import shuffle
 
 import infection_monkey.config
-from infection_monkey.network import HostScanner, HostFinger
+from infection_monkey.network.HostFinger import HostFinger
+from infection_monkey.network.HostScanner import HostScanner
 from infection_monkey.network.tools import check_tcp_ports, tcp_port_to_service
 
 __author__ = 'itamar'

monkey/infection_monkey/post_breach/actions/__init__.py

@@ -1,11 +0,0 @@
-from os.path import dirname, basename, isfile, join
-import glob
-
-
-def get_pba_files():
-    """
-    Gets all files under current directory(/actions)
-    :return: list of all files without .py ending
-    """
-    files = glob.glob(join(dirname(__file__), "*.py"))
-    return [basename(f)[:-3] for f in files if isfile(f) and not f.endswith('__init__.py')]

monkey/infection_monkey/post_breach/pba.py

@@ -6,7 +6,8 @@ from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.config import WormConfiguration
 from infection_monkey.telemetry.attack.t1064_telem import T1064Telem
-
+from infection_monkey.utils.plugins.plugin import Plugin
+import infection_monkey.post_breach.actions
 LOG = logging.getLogger(__name__)
 
 __author__ = 'VakarisZ'
@@ -14,11 +15,19 @@ __author__ = 'VakarisZ'
 EXECUTION_WITHOUT_OUTPUT = "(PBA execution produced no output)"
 
 
-class PBA(object):
+class PBA(Plugin):
     """
     Post breach action object. Can be extended to support more than command execution on target machine.
     """
 
+    @staticmethod
+    def base_package_name():
+        return infection_monkey.post_breach.actions.__package__
+
+    @staticmethod
+    def base_package_file():
+        return infection_monkey.post_breach.actions.__file__
+
     def __init__(self, name="unknown", linux_cmd="", windows_cmd=""):
         """
         :param name: Name of post breach action.

monkey/infection_monkey/post_breach/post_breach_handler.py

@@ -1,9 +1,8 @@
 import logging
-import inspect
+from typing import Sequence
-import importlib
+
-from infection_monkey.post_breach.pba import PBA
-from infection_monkey.post_breach.actions import get_pba_files
 from infection_monkey.utils.environment import is_windows_os
+from infection_monkey.post_breach.pba import PBA
 
 LOG = logging.getLogger(__name__)
 
@@ -34,25 +33,8 @@ class PostBreach(object):
         LOG.info("All PBAs executed. Total {} executed.".format(len(self.pba_list)))
 
     @staticmethod
-    def config_to_pba_list():
+    def config_to_pba_list() -> Sequence[PBA]:
         """
-        Passes config to each post breach action class and aggregates results into a list.
         :return: A list of PBA objects.
         """
-        pba_list = []
+        return PBA.get_instances()
-        pba_files = get_pba_files()
-        # Go through all of files in ./actions
-        for pba_file in pba_files:
-            # Import module from that file
-            module = importlib.import_module(PATH_TO_ACTIONS + pba_file)
-            # Get all classes in a module
-            pba_classes = [m[1] for m in inspect.getmembers(module, inspect.isclass)
-                           if ((m[1].__module__ == module.__name__) and issubclass(m[1], PBA))]
-            # Get post breach action object from class
-            for pba_class in pba_classes:
-                LOG.debug("Checking if should run PBA {}".format(pba_class.__name__))
-                if pba_class.should_run(pba_class.__name__):
-                    pba = pba_class()
-                    pba_list.append(pba)
-                    LOG.debug("Added PBA {} to PBA list".format(pba_class.__name__))
-        return pba_list

monkey/infection_monkey/pyinstaller_hooks/hook-infection_monkey.network.py

@@ -0,0 +1,4 @@
+from PyInstaller.utils.hooks import collect_submodules, collect_data_files
+
+hiddenimports = collect_submodules('infection_monkey.network')
+datas = (collect_data_files('infection_monkey.network', include_py_files=True))

monkey/infection_monkey/utils/plugins/plugin.py

@@ -0,0 +1,65 @@
+import importlib
+import inspect
+import logging
+from abc import ABCMeta, abstractmethod
+from os.path import dirname, basename, isfile, join
+import glob
+from typing import Sequence, TypeVar, Type
+
+LOG = logging.getLogger(__name__)
+
+
+def _get_candidate_files(base_package_file):
+    files = glob.glob(join(dirname(base_package_file), "*.py"))
+    return [basename(f)[:-3] for f in files if isfile(f) and not f.endswith('__init__.py')]
+
+
+Plugin_type = TypeVar('Plugin_type', bound='Plugin')
+
+
+class Plugin(metaclass=ABCMeta):
+
+    @staticmethod
+    @abstractmethod
+    def should_run(class_name: str) -> bool:
+        raise NotImplementedError()
+
+    @classmethod
+    def get_instances(cls) -> Sequence[Type[Plugin_type]]:
+        """
+        Returns the type objects from base_package_spec.
+        base_package name and file must refer to the same package otherwise bad results
+        :return: A list of parent_class objects.
+        """
+        objects = []
+        candidate_files = _get_candidate_files(cls.base_package_file())
+        LOG.info("looking for classes of type {} in {}".format(cls.__name__, cls.base_package_name()))
+        # Go through all of files
+        for file in candidate_files:
+            # Import module from that file
+            module = importlib.import_module('.' + file, cls.base_package_name())
+            # Get all classes in a module
+            # m[1] because return object is (name,class)
+            classes = [m[1] for m in inspect.getmembers(module, inspect.isclass)
+                       if ((m[1].__module__ == module.__name__) and issubclass(m[1], cls))]
+            # Get object from class
+            for class_object in classes:
+                LOG.debug("Checking if should run object {}".format(class_object.__name__))
+                try:
+                    if class_object.should_run(class_object.__name__):
+                        instance = class_object()
+                        objects.append(instance)
+                        LOG.debug("Added {} to list".format(class_object.__name__))
+                except Exception as e:
+                    LOG.warning("Exception {} when checking if {} should run".format(str(e), class_object.__name__))
+        return objects
+
+    @staticmethod
+    @abstractmethod
+    def base_package_file():
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def base_package_name():
+        pass