p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Agent: Modify scan_tcp_ports() to accept an IAgentEventQueue

This commit is contained in:
Mike Salvatore 2022-09-29 14:45:24 -04:00
parent 4b3402f7a8
commit 052c31e8ff
3 changed files with 17 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
monkey/infection_monkey/network_scanning/tcp_scanner.py
View File

@ -5,6 +5,7 @@ import time
from pprint import pformat
from typing import Collection, Dict, Iterable, Mapping, Tuple
from common.event_queue import IAgentEventQueue
from common.types import PortStatus
from common.utils import Timer
from infection_monkey.i_puppet import PortScanData
@ -17,7 +18,7 @@ EMPTY_PORT_SCAN = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)}
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]:
try:
return _scan_tcp_ports(host, ports_to_scan, timeout)

2
monkey/infection_monkey/puppet/puppet.py
View File

@ -48,7 +48,7 @@ class Puppet(IPuppet):
def scan_tcp_ports(
self, host: str, ports: Sequence[int], timeout: float = CONNECTION_TIMEOUT
) -> Dict[int, PortScanData]:
return network_scanning.scan_tcp_ports(host, ports, timeout)
return network_scanning.scan_tcp_ports(host, ports, timeout, self._agent_event_queue)
def fingerprint(
self,

22
monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py
View File

@ -20,10 +20,12 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data):
@pytest.mark.parametrize("open_ports_data", [OPEN_PORTS_DATA])
def test_tcp_successful(monkeypatch, patch_check_tcp_ports, 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)
port_scan_data = scan_tcp_ports("127.0.0.1", PORTS_TO_SCAN, 0, mock_agent_event_queue)
assert len(port_scan_data) == 6
for port in open_ports_data.keys():
@ -38,8 +40,10 @@ def test_tcp_successful(monkeypatch, patch_check_tcp_ports, open_ports_data):
@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)
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)
assert len(port_scan_data) == 6
for port in open_ports_data:
@ -49,15 +53,17 @@ def test_tcp_empty_response(monkeypatch, patch_check_tcp_ports, open_ports_data)
@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)
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)
assert len(port_scan_data) == 0
def test_exception_handling(monkeypatch):
def test_exception_handling(monkeypatch, mock_agent_event_queue):
monkeypatch.setattr(
"infection_monkey.network_scanning.tcp_scanner._scan_tcp_ports",
MagicMock(side_effect=Exception),
)
assert scan_tcp_ports("abc", [123], 123) == EMPTY_PORT_SCAN
assert scan_tcp_ports("abc", [123], 123, mock_agent_event_queue) == EMPTY_PORT_SCAN