UT: Use operating systems consts

This commit is contained in:
vakarisz 2022-06-23 11:34:11 +03:00 committed by vakaris_zilius
parent 4c1c8044cd
commit 77f8be523c
4 changed files with 51 additions and 19 deletions

View File

@ -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):

View File

@ -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)

View File

@ -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")

View File

@ -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": {