UT: Use operating systems consts

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": "HadoopExploiter",
-            {"name": "WmiExploiter", "supported_os": ["windows"], "options": {"timeout": 10}},
+                "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": {