UT: Add tests for all_exploitation_ports_are_closed()

monkey/tests/unit_tests/infection_monkey/exploit/tools/test_exploit_tools_utils.py

@@ -0,0 +1,64 @@
+from ipaddress import IPv4Address
+
+import pytest
+
+from common import OperatingSystem
+from common.types import PortStatus
+from infection_monkey.exploit.tools import all_exploitation_ports_are_closed
+from infection_monkey.i_puppet import PortScanData, PortScanDataDict, TargetHost, TargetHostPorts
+
+TARGET_IP = IPv4Address("127.0.0.1")
+PORTS_STATUS = TargetHostPorts(
+    tcp_ports=PortScanDataDict(
+        {
+            5000: PortScanData(port=5000, status=PortStatus.OPEN),
+            5001: PortScanData(port=5001, status=PortStatus.CLOSED),
+            6000: PortScanData(port=6000, status=PortStatus.OPEN),
+            6001: PortScanData(port=6001, status=PortStatus.CLOSED),
+        }
+    )
+)
+
+
+@pytest.fixture
+def target_host() -> TargetHost:
+    return TargetHost(
+        ip=TARGET_IP,
+        operating_system=OperatingSystem.WINDOWS,
+        ports_status=PORTS_STATUS,
+    )
+
+
+def test_all_exploitation_ports_open(target_host):
+    exploitation_ports = [5000, 6000]
+
+    assert not all_exploitation_ports_are_closed(target_host, exploitation_ports)
+
+
+def test_some_exploitation_ports_open(target_host):
+    exploitation_ports = [5000, 5001, 6000, 6001]
+
+    assert not all_exploitation_ports_are_closed(target_host, exploitation_ports)
+
+
+def test_all_exploitation_ports_closed(target_host):
+    exploitation_ports = [5001, 6001]
+
+    assert all_exploitation_ports_are_closed(target_host, exploitation_ports)
+
+
+def test_no_exploitation_ports_exist(target_host):
+    exploitation_ports = []
+
+    assert all_exploitation_ports_are_closed(target_host, exploitation_ports)
+
+
+def test_host_has_no_ports_status():
+    target_host = TargetHost(
+        ip=TARGET_IP,
+        operating_system=OperatingSystem.WINDOWS,
+        ports_status=PortScanDataDict({}),
+    )
+    exploitation_ports = [5000, 5001]
+
+    assert not all_exploitation_ports_are_closed(target_host, exploitation_ports)