From 543ff24ac38b94c703ca59ea952931aed890325c Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 10 Feb 2022 17:56:38 +0100 Subject: [PATCH] UT: Add tests for tcp scanning --- .../network/test_tcp_scanning.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 monkey/tests/unit_tests/infection_monkey/network/test_tcp_scanning.py diff --git a/monkey/tests/unit_tests/infection_monkey/network/test_tcp_scanning.py b/monkey/tests/unit_tests/infection_monkey/network/test_tcp_scanning.py new file mode 100644 index 000000000..e383e1004 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/network/test_tcp_scanning.py @@ -0,0 +1,54 @@ +import pytest + +from infection_monkey.i_puppet import PortStatus +from infection_monkey.network import scan_tcp_ports + +PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] + +OPEN_PORTS_DATA = {22: "SSH-banner", 80: "", 2222: "SSH2-banner"} + + +@pytest.fixture +def patch_check_tcp_ports(monkeypatch, open_ports_data): + monkeypatch.setattr( + "infection_monkey.network.tcp_scanner._check_tcp_ports", + lambda *_: open_ports_data, + ) + + +@pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) +def test_tcp_successful(monkeypatch, patch_check_tcp_ports, open_ports_data): + closed_ports = [8080, 143, 445] + + port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + + assert len(port_scan_data) == 6 + for port in open_ports_data.keys(): + assert port_scan_data[port].port == port + assert port_scan_data[port].status == PortStatus.OPEN + assert port_scan_data[port].banner == open_ports_data.get(port) + + for port in closed_ports: + assert port_scan_data[port].port == port + assert port_scan_data[port].status == PortStatus.CLOSED + assert port_scan_data[port].banner is None + + +@pytest.mark.parametrize("open_ports_data", [{}]) +def test_tcp_empty_response(monkeypatch, patch_check_tcp_ports, open_ports_data): + + port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0) + + assert len(port_scan_data) == 6 + for port in open_ports_data: + assert port_scan_data[port].port == port + assert port_scan_data[port].status == PortStatus.CLOSED + assert port_scan_data[port].banner is None + + +@pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) +def test_tcp_no_ports_to_scan(monkeypatch, patch_check_tcp_ports, open_ports_data): + + port_scan_data = scan_tcp_ports("127.0.0.1", [], 0) + + assert len(port_scan_data) == 0