Agent: Change return type of IPuppet.fingerprint()

monkey/infection_monkey/i_puppet.py

@@ -13,6 +13,7 @@ class PortStatus(Enum):
 ExploiterResultData = namedtuple("ExploiterResultData", ["result", "info", "attempts"])
 PingScanData = namedtuple("PingScanData", ["response_received", "os"])
 PortScanData = namedtuple("PortScanData", ["port", "status", "banner", "service"])
+FingerprintData = namedtuple("FingerprintData", ["os_type", "os_version", "services"])
 PostBreachData = namedtuple("PostBreachData", ["command", "result"])
 
 
@@ -57,13 +58,13 @@ class IPuppet(metaclass=abc.ABCMeta):
         """
 
     @abc.abstractmethod
-    def fingerprint(self, name: str, host: str) -> Dict:
+    def fingerprint(self, name: str, host: str) -> FingerprintData:
         """
         Runs a fingerprinter against a remote host
         :param str name: The name of the fingerprinter to run
         :param str host: The domain name or IP address of a host
-        :return: A dictionary containing the information collected by the fingerprinter
-        :rtype: Dict
+        :return: The data collected by running the fingerprinter on the specified host
+        :rtype: FingerprintData
         """
 
     @abc.abstractmethod

monkey/infection_monkey/puppet/mock_puppet.py

@@ -4,6 +4,7 @@ from typing import Dict, Tuple
 
 from infection_monkey.i_puppet import (
     ExploiterResultData,
+    FingerprintData,
     IPuppet,
     PingScanData,
     PortScanData,
@@ -193,29 +194,37 @@ class MockPuppet(IPuppet):
 
         return _get_empty_results(port)
 
-    def fingerprint(self, name: str, host: str) -> Dict:
+    def fingerprint(self, name: str, host: str) -> FingerprintData:
         logger.debug(f"fingerprint({name}, {host})")
+        empty_fingerprint_data = FingerprintData(None, None, {})
+
         dot_1_results = {
-            "SMBFinger": {
-                "os": {"type": "windows", "version": "vista"},
-                "services": {"tcp-445": {"name": "SSH", "os": "linux"}},
-            }
+            "SMBFinger": FingerprintData(
+                "windows", "vista", {"tcp-445": {"name": "smb_service_name"}}
+            )
         }
 
         dot_3_results = {
-            "SSHFinger": {"os": "linux", "services": {"tcp-22": {"name": "SSH"}}},
-            "HTTPFinger": {
-                "services": {"tcp-https": {"name": "http", "data": ("SERVER_HEADERS", DOT_3)}}
-            },
+            "SSHFinger": FingerprintData(
+                "linux", "ubuntu", {"tcp-22": {"name": "SSH", "banner": "SSH BANNER"}}
+            ),
+            "HTTPFinger": FingerprintData(
+                None,
+                None,
+                {
+                    "tcp-80": {"name": "http", "data": ("SERVER_HEADERS", False)},
+                    "tcp-443": {"name": "http", "data": ("SERVER_HEADERS_2", True)},
+                },
+            ),
         }
 
         if host == DOT_1:
-            return dot_1_results.get(name, {})
+            return dot_1_results.get(name, empty_fingerprint_data)
 
         if host == DOT_3:
-            return dot_3_results.get(name, {})
+            return dot_3_results.get(name, empty_fingerprint_data)
 
-        return {}
+        return empty_fingerprint_data
 
     def exploit_host(
         self, name: str, host: str, options: Dict, interrupt: threading.Event