From 77f8be523cfb1270591d1f2cf57453059268ca49 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 23 Jun 2022 11:34:11 +0300 Subject: [PATCH] UT: Use operating systems consts --- .../infection_monkey/master/mock_puppet.py | 35 ++++++++++++++----- .../infection_monkey/master/test_exploiter.py | 25 +++++++++---- .../network_scanning/test_ping_scanner.py | 5 +-- .../test_ssh_fingerprinter.py | 5 +-- 4 files changed, 51 insertions(+), 19 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py index 295c68ceb..4d77bd0b3 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py +++ b/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): diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py index 3c76c903f..c804b60e5 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py +++ b/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) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index 88c9dbeca..9e2891e3f 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/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") diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ssh_fingerprinter.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ssh_fingerprinter.py index 69c8eb580..09d0705ef 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ssh_fingerprinter.py +++ b/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": {