From a3ca21481ec7602e7ecd1786f60410e5881204dc Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 11:27:07 +0530 Subject: [PATCH 01/11] Agent: Add missing/fix existing type hints in TCP scanner --- monkey/infection_monkey/network_scanning/tcp_scanner.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index af1efb182..efff5ceaf 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -27,7 +27,9 @@ def scan_tcp_ports( return EMPTY_PORT_SCAN -def _scan_tcp_ports(host: str, ports_to_scan: Collection[int], timeout: float): +def _scan_tcp_ports( + host: str, ports_to_scan: Collection[int], timeout: float +) -> Dict[int, PortScanData]: open_ports = _check_tcp_ports(host, ports_to_scan, timeout) return _build_port_scan_data(ports_to_scan, open_ports) @@ -35,7 +37,7 @@ def _scan_tcp_ports(host: str, ports_to_scan: Collection[int], timeout: float): def _build_port_scan_data( ports_to_scan: Iterable[int], open_ports: Mapping[int, str] -) -> Mapping[int, PortScanData]: +) -> Dict[int, PortScanData]: port_scan_data = {} for port in ports_to_scan: if port in open_ports: @@ -55,7 +57,7 @@ def _get_closed_port_data(port: int) -> PortScanData: def _check_tcp_ports( ip: str, ports_to_scan: Collection[int], timeout: float = DEFAULT_TIMEOUT -) -> Mapping[int, str]: +) -> Dict[int, str]: """ Checks whether any of the given ports are open on a target IP. :param ip: IP of host to attack From 58ddd6e47d72b36b7a95d8a5d80d22e9c64c8a77 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 11:43:51 +0530 Subject: [PATCH 02/11] Agent: Partially implement publishing TCPScanEvent in TCP scanner --- .../network_scanning/tcp_scanner.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index efff5ceaf..0cbe69494 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -21,18 +21,27 @@ def scan_tcp_ports( host: str, ports_to_scan: Collection[int], timeout: float, agent_event_queue: IAgentEventQueue ) -> Dict[int, PortScanData]: try: - return _scan_tcp_ports(host, ports_to_scan, timeout) + return _scan_tcp_ports(host, ports_to_scan, timeout, agent_event_queue) except Exception: logger.exception("Unhandled exception occurred while trying to scan tcp ports") return EMPTY_PORT_SCAN def _scan_tcp_ports( - host: str, ports_to_scan: Collection[int], timeout: float + host: str, ports_to_scan: Collection[int], timeout: float, agent_event_queue: IAgentEventQueue ) -> Dict[int, PortScanData]: open_ports = _check_tcp_ports(host, ports_to_scan, timeout) - return _build_port_scan_data(ports_to_scan, open_ports) + tcp_scan_data = _build_port_scan_data(ports_to_scan, open_ports) + + tcp_scan_event = _generate_tcp_scan_event(host, tcp_scan_data) + agent_event_queue.publish(tcp_scan_event) + + return tcp_scan_data + + +def _generate_tcp_scan_event(host: str, tcp_scan_data: Dict[int, PortScanData]): + pass def _build_port_scan_data( From 35d3038bc841fbefeef796ff64592080b292a923 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 13:28:24 +0530 Subject: [PATCH 03/11] Agent: Create TCPScanEvent and return from _generate_tcp_scan_event() in TCP scanner --- .../network_scanning/tcp_scanner.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 0cbe69494..3f93069f3 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -5,6 +5,7 @@ import time from pprint import pformat from typing import Collection, Dict, Iterable, Mapping, Tuple +from common.agent_events import TCPScanEvent from common.event_queue import IAgentEventQueue from common.types import PortStatus from common.utils import Timer @@ -32,16 +33,20 @@ def _scan_tcp_ports( ) -> Dict[int, PortScanData]: open_ports = _check_tcp_ports(host, ports_to_scan, timeout) - tcp_scan_data = _build_port_scan_data(ports_to_scan, open_ports) + port_scan_data = _build_port_scan_data(ports_to_scan, open_ports) - tcp_scan_event = _generate_tcp_scan_event(host, tcp_scan_data) + tcp_scan_event = _generate_tcp_scan_event(host, port_scan_data) agent_event_queue.publish(tcp_scan_event) - return tcp_scan_data + return port_scan_data -def _generate_tcp_scan_event(host: str, tcp_scan_data: Dict[int, PortScanData]): - pass +def _generate_tcp_scan_event(host: str, port_scan_data: Dict[int, PortScanData]): + port_statuses = {} + for port, data in port_scan_data.items(): + port_statuses[port] = data.status + + return TCPScanEvent(target=host, ports=port_statuses) def _build_port_scan_data( From bab4ebc2bcd9ceb3ba2ebde5e71a829aeb47142c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 13:47:31 +0530 Subject: [PATCH 04/11] Agent: Add 'source' field when creating TCPScanEvent in TCP scanner --- monkey/infection_monkey/network_scanning/tcp_scanner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 3f93069f3..2aedc6ee5 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -11,6 +11,7 @@ from common.types import PortStatus from common.utils import Timer from infection_monkey.i_puppet import PortScanData from infection_monkey.network.tools import BANNER_READ, DEFAULT_TIMEOUT, tcp_port_to_service +from infection_monkey.utils.ids import get_agent_id logger = logging.getLogger(__name__) @@ -46,7 +47,7 @@ def _generate_tcp_scan_event(host: str, port_scan_data: Dict[int, PortScanData]) for port, data in port_scan_data.items(): port_statuses[port] = data.status - return TCPScanEvent(target=host, ports=port_statuses) + return TCPScanEvent(source=get_agent_id(), target=host, ports=port_statuses) def _build_port_scan_data( From 0bf9309e075f3a12c3eeb2008f38b403bee49da0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 13:48:35 +0530 Subject: [PATCH 05/11] UT: Assert mock_agent_event_queue.publish's call counts in test_tcp_scanner.py --- .../infection_monkey/network_scanning/test_tcp_scanner.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index 3e0c38d7e..578f57508 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -38,6 +38,8 @@ def test_tcp_successful( assert port_scan_data[port].status == PortStatus.CLOSED assert port_scan_data[port].banner is None + assert mock_agent_event_queue.publish.call_count == 1 + @pytest.mark.parametrize("open_ports_data", [{}]) def test_tcp_empty_response( @@ -51,6 +53,8 @@ def test_tcp_empty_response( assert port_scan_data[port].status == PortStatus.CLOSED assert port_scan_data[port].banner is None + assert mock_agent_event_queue.publish.call_count == 1 + @pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) def test_tcp_no_ports_to_scan( @@ -59,6 +63,7 @@ def test_tcp_no_ports_to_scan( port_scan_data = scan_tcp_ports("127.0.0.1", [], 0, mock_agent_event_queue) assert len(port_scan_data) == 0 + assert mock_agent_event_queue.publish.call_count == 1 def test_exception_handling(monkeypatch, mock_agent_event_queue): @@ -67,3 +72,4 @@ def test_exception_handling(monkeypatch, mock_agent_event_queue): MagicMock(side_effect=Exception), ) assert scan_tcp_ports("abc", [123], 123, mock_agent_event_queue) == EMPTY_PORT_SCAN + assert mock_agent_event_queue.publish.call_count == 0 From 0a11d34fb7c3597c3678e497aaca02f4487fce6e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 14:20:22 +0530 Subject: [PATCH 06/11] UT: Assert mock_agent_event_queue.publish's call args in test_tcp_scanner.py --- .../network_scanning/test_tcp_scanner.py | 34 +++++++++++++++++-- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index 578f57508..c96728658 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -2,9 +2,12 @@ from unittest.mock import MagicMock import pytest +from common.agent_events import TCPScanEvent from common.types import PortStatus +from infection_monkey.i_puppet import PortScanData from infection_monkey.network_scanning import scan_tcp_ports from infection_monkey.network_scanning.tcp_scanner import EMPTY_PORT_SCAN +from infection_monkey.utils.ids import get_agent_id PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] @@ -19,13 +22,28 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data): ) +HOST_IP = "127.0.0.1" + + +def _get_tcp_scan_event(port_scan_data: PortScanData): + port_statuses = {} + for port, data in port_scan_data.items(): + port_statuses[port] = data.status + + return TCPScanEvent( + source=get_agent_id(), + target=HOST_IP, + ports=port_statuses, + ) + + @pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) def test_tcp_successful( monkeypatch, patch_check_tcp_ports, open_ports_data, mock_agent_event_queue ): closed_ports = [8080, 143, 445] - port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue) assert len(port_scan_data) == 6 for port in open_ports_data.keys(): @@ -38,14 +56,17 @@ def test_tcp_successful( assert port_scan_data[port].status == PortStatus.CLOSED assert port_scan_data[port].banner is None + event = _get_tcp_scan_event(port_scan_data) + assert mock_agent_event_queue.publish.call_count == 1 + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.parametrize("open_ports_data", [{}]) def test_tcp_empty_response( monkeypatch, patch_check_tcp_ports, open_ports_data, mock_agent_event_queue ): - port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue) assert len(port_scan_data) == 6 for port in open_ports_data: @@ -53,17 +74,24 @@ def test_tcp_empty_response( assert port_scan_data[port].status == PortStatus.CLOSED assert port_scan_data[port].banner is None + event = _get_tcp_scan_event(port_scan_data) + assert mock_agent_event_queue.publish.call_count == 1 + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA]) def test_tcp_no_ports_to_scan( monkeypatch, patch_check_tcp_ports, open_ports_data, mock_agent_event_queue ): - port_scan_data = scan_tcp_ports("127.0.0.1", [], 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, [], 0, mock_agent_event_queue) assert len(port_scan_data) == 0 + + event = _get_tcp_scan_event(port_scan_data) + assert mock_agent_event_queue.publish.call_count == 1 + mock_agent_event_queue.publish.assert_called_with(event) def test_exception_handling(monkeypatch, mock_agent_event_queue): From 9754b4731ce62f2e79d2ded754216ea05722d428 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 17:27:17 +0530 Subject: [PATCH 07/11] UT: Mock AbstractAgentEvent's timestamp (time.time()) in test_tcp_scanner.py --- .../network_scanning/test_tcp_scanner.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index c96728658..e47febc36 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -1,3 +1,4 @@ +import time from unittest.mock import MagicMock import pytest @@ -13,6 +14,17 @@ PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] OPEN_PORTS_DATA = {22: "SSH-banner", 80: "", 2222: "SSH2-banner"} +TIMESTAMP = 123.321 + + +@pytest.fixture(scope="module") +def patch_timestamp(monkeypatch): + monkeypatch.setattr( + time, + "time", + TIMESTAMP, + ) + @pytest.fixture def patch_check_tcp_ports(monkeypatch, open_ports_data): From 96af86f766de2a536ed2c3b240af76ea8c8d9478 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 30 Sep 2022 17:29:52 +0530 Subject: [PATCH 08/11] UT: Move variable to above fixtures in test_tcp_scanner.py --- .../infection_monkey/network_scanning/test_tcp_scanner.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index e47febc36..17acd421a 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -16,6 +16,8 @@ OPEN_PORTS_DATA = {22: "SSH-banner", 80: "", 2222: "SSH2-banner"} TIMESTAMP = 123.321 +HOST_IP = "127.0.0.1" + @pytest.fixture(scope="module") def patch_timestamp(monkeypatch): @@ -34,9 +36,6 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data): ) -HOST_IP = "127.0.0.1" - - def _get_tcp_scan_event(port_scan_data: PortScanData): port_statuses = {} for port, data in port_scan_data.items(): From 9154f6f9dcc021190eb11ff117659ffe1f2d93f7 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:48:16 +0200 Subject: [PATCH 09/11] Agent: Generate timestamp when checking for tcp ports --- .../network_scanning/tcp_scanner.py | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 2aedc6ee5..e19b2e0eb 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -1,8 +1,9 @@ import logging import select import socket -import time +from ipaddress import IPv4Address from pprint import pformat +from time import sleep, time from typing import Collection, Dict, Iterable, Mapping, Tuple from common.agent_events import TCPScanEvent @@ -32,22 +33,29 @@ def scan_tcp_ports( def _scan_tcp_ports( host: str, ports_to_scan: Collection[int], timeout: float, agent_event_queue: IAgentEventQueue ) -> Dict[int, PortScanData]: - open_ports = _check_tcp_ports(host, ports_to_scan, timeout) + event_timestamp, open_ports = _check_tcp_ports(host, ports_to_scan, timeout) port_scan_data = _build_port_scan_data(ports_to_scan, open_ports) - tcp_scan_event = _generate_tcp_scan_event(host, port_scan_data) + tcp_scan_event = _generate_tcp_scan_event(host, port_scan_data, event_timestamp) agent_event_queue.publish(tcp_scan_event) return port_scan_data -def _generate_tcp_scan_event(host: str, port_scan_data: Dict[int, PortScanData]): +def _generate_tcp_scan_event( + host: str, port_scan_data: Dict[int, PortScanData], event_timestamp: float +): port_statuses = {} for port, data in port_scan_data.items(): port_statuses[port] = data.status - return TCPScanEvent(source=get_agent_id(), target=host, ports=port_statuses) + return TCPScanEvent( + source=get_agent_id(), + target=IPv4Address(host), + timestamp=event_timestamp, + ports=port_statuses, + ) def _build_port_scan_data( @@ -72,7 +80,7 @@ def _get_closed_port_data(port: int) -> PortScanData: def _check_tcp_ports( ip: str, ports_to_scan: Collection[int], timeout: float = DEFAULT_TIMEOUT -) -> Dict[int, str]: +) -> Tuple[float, Dict[int, str]]: """ Checks whether any of the given ports are open on a target IP. :param ip: IP of host to attack @@ -89,6 +97,7 @@ def _check_tcp_ports( connected_ports = set() open_ports = {} + event_timestamp = time() try: logger.debug( "Connecting to the following ports %s" % ",".join((str(x) for x in ports_to_scan)) @@ -117,7 +126,7 @@ def _check_tcp_ports( while (not timer.is_expired()) and sockets_to_try: # The call to select() may return sockets that are writeable but not actually # connected. Adding this sleep prevents excessive looping. - time.sleep(min(POLL_INTERVAL, timer.time_remaining)) + sleep(min(POLL_INTERVAL, timer.time_remaining)) sock_objects = [s[1] for s in sockets_to_try] @@ -153,7 +162,7 @@ def _check_tcp_ports( _clean_up_sockets(possible_ports, connected_ports) - return open_ports + return event_timestamp, open_ports def _clean_up_sockets( From 3f89e50930d9a4c18bf7c3fae37fb46d18fbd82a Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:50:37 +0200 Subject: [PATCH 10/11] UT: Fix tcp_scanner tests to patch the time function --- .../network_scanning/test_tcp_scanner.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index 17acd421a..f3c897f11 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -1,4 +1,3 @@ -import time from unittest.mock import MagicMock import pytest @@ -19,12 +18,11 @@ TIMESTAMP = 123.321 HOST_IP = "127.0.0.1" -@pytest.fixture(scope="module") +@pytest.fixture(autouse=True) def patch_timestamp(monkeypatch): monkeypatch.setattr( - time, - "time", - TIMESTAMP, + "infection_monkey.network_scanning.tcp_scanner.time", + lambda: TIMESTAMP, ) @@ -32,7 +30,7 @@ def patch_timestamp(monkeypatch): def patch_check_tcp_ports(monkeypatch, open_ports_data): monkeypatch.setattr( "infection_monkey.network_scanning.tcp_scanner._check_tcp_ports", - lambda *_: open_ports_data, + lambda *_: (TIMESTAMP, open_ports_data), ) @@ -44,6 +42,7 @@ def _get_tcp_scan_event(port_scan_data: PortScanData): return TCPScanEvent( source=get_agent_id(), target=HOST_IP, + timestamp=TIMESTAMP, ports=port_statuses, ) From 4987dddc0c0bd30711d8a1fe6cd9a087bb425ee1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 11:50:18 -0400 Subject: [PATCH 11/11] Agent: Use dict comprehension instead of for --- monkey/infection_monkey/network_scanning/tcp_scanner.py | 4 +--- .../infection_monkey/network_scanning/test_tcp_scanner.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index e19b2e0eb..9b6d581e2 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -46,9 +46,7 @@ def _scan_tcp_ports( def _generate_tcp_scan_event( host: str, port_scan_data: Dict[int, PortScanData], event_timestamp: float ): - port_statuses = {} - for port, data in port_scan_data.items(): - port_statuses[port] = data.status + port_statuses = {port: psd.status for port, psd in port_scan_data.items()} return TCPScanEvent( source=get_agent_id(), diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index f3c897f11..2fd67590d 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -35,9 +35,7 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data): def _get_tcp_scan_event(port_scan_data: PortScanData): - port_statuses = {} - for port, data in port_scan_data.items(): - port_statuses[port] = data.status + port_statuses = {port: psd.status for port, psd in port_scan_data.items()} return TCPScanEvent( source=get_agent_id(),