Merge pull request #2040 from guardicore/1960-use-operating-systems

1960 use operating systems

monkey/infection_monkey/exploit/hadoop.py

@@ -105,10 +105,10 @@ class HadoopExploiter(WebRCE):
     def _build_command(self, path, http_path):
         # Build command to execute
         monkey_cmd = build_monkey_commandline(self.host, self.current_depth + 1)
-        if "linux" in self.host.os["type"]:
-            base_command = HADOOP_LINUX_COMMAND
-        else:
+        if self.host.is_windows():
             base_command = HADOOP_WINDOWS_COMMAND
+        else:
+            base_command = HADOOP_LINUX_COMMAND
 
         return base_command % {
             "monkey_path": path,

monkey/infection_monkey/exploit/log4shell.py

@@ -2,6 +2,7 @@ import logging
 import time
 from pathlib import PurePath
 
+from common import OperatingSystems
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT
 from common.utils import Timer
 from infection_monkey.exploit.log4shell_utils import (
@@ -115,10 +116,10 @@ class Log4ShellExploiter(WebRCE):
     def _build_command(self, path: PurePath, http_path) -> str:
         # Build command to execute
         monkey_cmd = build_monkey_commandline(self.host, self.current_depth + 1, location=path)
-        if "linux" in self.host.os["type"]:
-            base_command = LOG4SHELL_LINUX_COMMAND
-        else:
+        if self.host.is_windows():
             base_command = LOG4SHELL_WINDOWS_COMMAND
+        else:
+            base_command = LOG4SHELL_LINUX_COMMAND
 
         return base_command % {
             "monkey_path": path,
@@ -128,7 +129,7 @@ class Log4ShellExploiter(WebRCE):
         }
 
     def _build_java_class(self, exploit_command: str) -> bytes:
-        if "linux" in self.host.os["type"]:
+        if OperatingSystems.LINUX in self.host.os["type"]:
             return build_exploit_bytecode(exploit_command, LINUX_EXPLOIT_TEMPLATE_PATH)
         else:
             return build_exploit_bytecode(exploit_command, WINDOWS_EXPLOIT_TEMPLATE_PATH)

monkey/infection_monkey/exploit/tools/helpers.py

@@ -15,7 +15,7 @@ AGENT_BINARY_PATH_WIN64 = PureWindowsPath(r"C:\Windows\temp\monkey64.exe")
 
 
 def get_agent_dst_path(host: VictimHost) -> PurePath:
-    if host.os["type"] == "windows":
+    if host.is_windows():
         path = PureWindowsPath(AGENT_BINARY_PATH_WIN64)
     else:
         path = PurePosixPath(AGENT_BINARY_PATH_LINUX)

monkey/infection_monkey/exploit/web_rce.py

@@ -3,6 +3,7 @@ from abc import abstractmethod
 from posixpath import join
 from typing import List, Tuple
 
+from common import OperatingSystems
 from common.utils.attack_utils import BITS_UPLOAD_STRING, ScanStatus
 from infection_monkey.exploit.HostExploiter import HostExploiter
 from infection_monkey.exploit.tools.http_tools import HTTPTools
@@ -162,10 +163,10 @@ class WebRCE(HostExploiter):
 
     def get_command(self, path, http_path, commands):
         try:
-            if "linux" in self.host.os["type"]:
-                command = commands["linux"]
-            else:
+            if self.host.is_windows():
                 command = commands["windows"]
+            else:
+                command = commands["linux"]
             # Format command
             command = command % {"monkey_path": path, "http_path": http_path}
         except KeyError:
@@ -326,7 +327,7 @@ class WebRCE(HostExploiter):
         :return: response, False if failed and True if permission change is not needed
         """
         logger.info("Changing monkey's permissions")
-        if "windows" in self.host.os["type"]:
+        if self.host.is_windows():
             logger.info("Permission change not required for windows")
             return True
         if not command:
@@ -411,13 +412,14 @@ class WebRCE(HostExploiter):
         :return: Default monkey's destination path for corresponding host or False if failed.
         """
         if not self.host.os.get("type") or (
-            self.host.os["type"] != "linux" and self.host.os["type"] != "windows"
+            self.host.os["type"] != OperatingSystems.LINUX
+            and self.host.os["type"] != OperatingSystems.WINDOWS
         ):
             logger.error("Target's OS was either unidentified or not supported. Aborting")
             return False
-        if self.host.os["type"] == "linux":
+        if self.host.os["type"] == OperatingSystems.LINUX:
             return DROPPER_TARGET_PATH_LINUX
-        if self.host.os["type"] == "windows":
+        if self.host.os["type"] == OperatingSystems.WINDOWS:
             return DROPPER_TARGET_PATH_WIN64
 
     def get_target_url(self):

monkey/infection_monkey/model/host.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from common import OperatingSystems
+
 
 class VictimHost(object):
     def __init__(self, ip_addr: str, domain_name: str = ""):
@@ -14,6 +16,9 @@ class VictimHost(object):
     def as_dict(self):
         return self.__dict__
 
+    def is_windows(self) -> bool:
+        return OperatingSystems.WINDOWS in self.os["type"]
+
     def __hash__(self):
         return hash(self.ip_addr)
 

monkey/infection_monkey/network_scanning/ping_scanner.py

@@ -5,6 +5,7 @@ import re
 import subprocess
 import sys
 
+from common import OperatingSystems
 from infection_monkey.i_puppet import PingScanData
 from infection_monkey.utils.environment import is_windows_os
 
@@ -79,9 +80,9 @@ def _process_ping_command_output(ping_command_output: str) -> PingScanData:
 
     operating_system = None
     if ttl <= LINUX_TTL:
-        operating_system = "linux"
+        operating_system = OperatingSystems.LINUX
     else:  # as far we we know, could also be OSX/BSD, but lets handle that when it comes up.
-        operating_system = "windows"
+        operating_system = OperatingSystems.WINDOWS
 
     return PingScanData(True, operating_system)
 

monkey/infection_monkey/network_scanning/smb_fingerprinter.py

@@ -5,6 +5,7 @@ from typing import Dict
 
 from odict import odict
 
+from common import OperatingSystems
 from infection_monkey.i_puppet import (
     FingerprintData,
     IFingerprinter,
@@ -193,9 +194,9 @@ class SMBFingerprinter(IFingerprinter):
                 logger.debug(f'os_version: "{os_version}", service_client: "{service_client}"')
 
                 if os_version.lower() != "unix":
-                    os_type = "windows"
+                    os_type = OperatingSystems.WINDOWS
                 else:
-                    os_type = "linux"
+                    os_type = OperatingSystems.LINUX
 
                 smb_service["name"] = service_client
 

monkey/infection_monkey/network_scanning/ssh_fingerprinter.py

@@ -1,6 +1,7 @@
 import re
 from typing import Dict, Optional, Tuple
 
+from common import OperatingSystems
 from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PingScanData, PortScanData
 
 SSH_REGEX = r"SSH-\d\.\d-OpenSSH"
@@ -40,6 +41,6 @@ class SSHFingerprinter(IFingerprinter):
         for dist in LINUX_DIST_SSH:
             if banner.lower().find(dist) != -1:
                 os_version = banner.split(" ").pop().strip()
-                os = "linux"
+                os = OperatingSystems.LINUX
 
         return os, os_version

monkey/infection_monkey/telemetry/base_telem.py

@@ -4,6 +4,7 @@ import logging
 
 from infection_monkey.control import ControlClient
 from infection_monkey.telemetry.i_telem import ITelem
+from infection_monkey.telemetry.telem_encoder import TelemetryJSONEncoder
 
 logger = logging.getLogger(__name__)
 LOGGED_DATA_LENGTH = 300  # How many characters of telemetry data will be logged
@@ -39,7 +40,7 @@ class BaseTelem(ITelem, metaclass=abc.ABCMeta):
 
     @property
     def json_encoder(self):
-        return json.JSONEncoder
+        return TelemetryJSONEncoder
 
     def _log_telem_sending(self, serialized_data: str, log_data=True):
         logger.debug(f"Sending {self.telem_category} telemetry.")

monkey/infection_monkey/telemetry/telem_encoder.py

@@ -0,0 +1,10 @@
+import json
+
+from common import OperatingSystems
+
+
+class TelemetryJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, OperatingSystems):
+            return obj.name
+        return json.JSONEncoder.default(self, obj)

monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py

@@ -2,6 +2,7 @@ import logging
 import threading
 from typing import Dict, Iterable, List, Sequence
 
+from common import OperatingSystems
 from infection_monkey.credential_collectors import LMHash, Password, SSHKeypair, Username
 from infection_monkey.i_puppet import (
     Credentials,
@@ -182,19 +183,23 @@ class MockPuppet(IPuppet):
             "vulnerable_ports": [22],
             "executed_cmds": [],
         }
-        os_windows = "windows"
-        os_linux = "linux"
 
         successful_exploiters = {
             DOT_1: {
                 "ZerologonExploiter": ExploiterResultData(
-                    False, False, False, os_windows, {}, [], "Zerologon failed"
+                    False, False, False, OperatingSystems.WINDOWS, {}, [], "Zerologon failed"
                 ),
                 "SSHExploiter": ExploiterResultData(
-                    False, False, False, os_linux, info_ssh, attempts, "Failed exploiting"
+                    False,
+                    False,
+                    False,
+                    OperatingSystems.LINUX,
+                    info_ssh,
+                    attempts,
+                    "Failed exploiting",
                 ),
                 "WmiExploiter": ExploiterResultData(
-                    True, True, False, os_windows, info_wmi, attempts, None
+                    True, True, False, OperatingSystems.WINDOWS, info_wmi, attempts, None
                 ),
             },
             DOT_3: {
@@ -202,16 +207,22 @@ class MockPuppet(IPuppet):
                     False,
                     False,
                     False,
-                    os_windows,
+                    OperatingSystems.WINDOWS,
                     info_wmi,
                     attempts,
                     "PowerShell Exploiter Failed",
                 ),
                 "SSHExploiter": ExploiterResultData(
-                    False, False, False, os_linux, info_ssh, attempts, "Failed exploiting"
+                    False,
+                    False,
+                    False,
+                    OperatingSystems.LINUX,
+                    info_ssh,
+                    attempts,
+                    "Failed exploiting",
                 ),
                 "ZerologonExploiter": ExploiterResultData(
-                    True, False, False, os_windows, {}, [], None
+                    True, False, False, OperatingSystems.WINDOWS, {}, [], None
                 ),
             },
         }
@@ -220,7 +231,13 @@ class MockPuppet(IPuppet):
             return successful_exploiters[host.ip_addr][name]
         except KeyError:
             return ExploiterResultData(
-                False, False, False, os_linux, {}, [], f"{name} failed for host {host}"
+                False,
+                False,
+                False,
+                OperatingSystems.LINUX,
+                {},
+                [],
+                f"{name} failed for host {host}",
             )
 
     def run_payload(self, name: str, options: Dict, interrupt: threading.Event):

monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py

@@ -7,6 +7,7 @@ from unittest.mock import MagicMock
 import pytest
 from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet
 
+from common import OperatingSystems
 from infection_monkey.master import Exploiter
 from infection_monkey.model import VictimHost
 
@@ -38,12 +39,24 @@ def exploiter_config():
     return {
         "options": {"dropper_path_linux": "/tmp/monkey"},
         "brute_force": [
-            {"name": "HadoopExploiter", "supported_os": ["windows"], "options": {"timeout": 10}},
-            {"name": "SSHExploiter", "supported_os": ["linux"], "options": {}},
-            {"name": "WmiExploiter", "supported_os": ["windows"], "options": {"timeout": 10}},
+            {
+                "name": "HadoopExploiter",
+                "supported_os": [OperatingSystems.WINDOWS],
+                "options": {"timeout": 10},
+            },
+            {"name": "SSHExploiter", "supported_os": [OperatingSystems.LINUX], "options": {}},
+            {
+                "name": "WmiExploiter",
+                "supported_os": [OperatingSystems.WINDOWS],
+                "options": {"timeout": 10},
+            },
         ],
         "vulnerability": [
-            {"name": "ZerologonExploiter", "supported_os": ["windows"], "options": {}},
+            {
+                "name": "ZerologonExploiter",
+                "supported_os": [OperatingSystems.WINDOWS],
+                "options": {},
+            },
         ],
     }
 
@@ -160,7 +173,7 @@ def test_exploiter_raises_exception(callback, hosts, hosts_to_exploit, run_explo
 
 def test_windows_exploiters_run_on_windows_host(callback, hosts, hosts_to_exploit, run_exploiters):
     host = VictimHost("10.0.0.1")
-    host.os["type"] = "windows"
+    host.os["type"] = OperatingSystems.WINDOWS
     q = enqueue_hosts([host])
     run_exploiters(MockPuppet(), 1, q)
 
@@ -172,7 +185,7 @@ def test_windows_exploiters_run_on_windows_host(callback, hosts, hosts_to_exploi
 
 def test_linux_exploiters_run_on_linux_host(callback, hosts, hosts_to_exploit, run_exploiters):
     host = VictimHost("10.0.0.1")
-    host.os["type"] = "linux"
+    host.os["type"] = OperatingSystems.LINUX
     q = enqueue_hosts([host])
     run_exploiters(MockPuppet(), 1, q)
 

monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py

@@ -4,6 +4,7 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from common import OperatingSystems
 from infection_monkey.network_scanning import ping
 from infection_monkey.network_scanning.ping_scanner import EMPTY_PING_SCAN
 
@@ -91,7 +92,7 @@ def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output):
     result = ping("192.168.1.1", 1.0)
 
     assert result.response_received
-    assert result.os == "linux"
+    assert result.os == OperatingSystems.LINUX
 
 
 @pytest.mark.usefixtures("set_os_linux")
@@ -109,7 +110,7 @@ def test_windows_ping_success(patch_subprocess_running_ping_with_ping_output):
     result = ping("192.168.1.1", 1.0)
 
     assert result.response_received
-    assert result.os == "windows"
+    assert result.os == OperatingSystems.WINDOWS
 
 
 @pytest.mark.usefixtures("set_os_windows")

monkey/tests/unit_tests/infection_monkey/network_scanning/test_ssh_fingerprinter.py

@@ -1,5 +1,6 @@
 import pytest
 
+from common import OperatingSystems
 from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus
 from infection_monkey.network_scanning.ssh_fingerprinter import SSHFingerprinter
 
@@ -56,7 +57,7 @@ def test_ssh_os(ssh_fingerprinter):
     results = ssh_fingerprinter.get_host_fingerprint("127.0.0.1", None, port_scan_data, None)
 
     assert results == FingerprintData(
-        "linux",
+        OperatingSystems.LINUX,
         "Ubuntu-4ubuntu0.2",
         {
             "tcp-22": {
@@ -78,7 +79,7 @@ def test_multiple_os(ssh_fingerprinter):
     results = ssh_fingerprinter.get_host_fingerprint("127.0.0.1", None, port_scan_data, None)
 
     assert results == FingerprintData(
-        "linux",
+        OperatingSystems.LINUX,
         "Debian",
         {
             "tcp-22": {