From 45c6cac60c9e671b40415857b7d6aec40fdd4ddf Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 8 Apr 2022 10:36:14 +0300 Subject: [PATCH] Agent: Improve method naming and exception handling --- monkey/infection_monkey/master/automated_master.py | 2 +- monkey/infection_monkey/network_scanning/__init__.py | 4 ++-- monkey/infection_monkey/network_scanning/ping_scanner.py | 4 ++-- monkey/infection_monkey/network_scanning/tcp_scanner.py | 4 ++-- monkey/infection_monkey/puppet/puppet.py | 6 +++--- .../network_scanning/test_tcp_scanning.py | 8 ++++---- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 99f63522a..af263af6e 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -228,7 +228,7 @@ class AutomatedMaster(IMaster): try: callback(p) except Exception: - logging.exception( + logger.exception( f"Got unhandled exception when running {plugin_type} plugin {p}. " f"Plugin was passed to {callback}" ) diff --git a/monkey/infection_monkey/network_scanning/__init__.py b/monkey/infection_monkey/network_scanning/__init__.py index a1ec7240d..8e97b0ec4 100644 --- a/monkey/infection_monkey/network_scanning/__init__.py +++ b/monkey/infection_monkey/network_scanning/__init__.py @@ -1,2 +1,2 @@ -from .ping_scanner import try_ping -from .tcp_scanner import try_scan_tcp_ports +from .ping_scanner import ping +from .tcp_scanner import scan_tcp_ports diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index c2417d47e..16fb2df96 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -16,11 +16,11 @@ EMPTY_PING_SCAN = PingScanData(False, None) logger = logging.getLogger(__name__) -def try_ping(host: str, timeout: float) -> PingScanData: +def ping(host: str, timeout: float) -> PingScanData: try: return _ping(host, timeout) except Exception: - logging.exception("Unhandled exception occurred while running ping") + logger.exception("Unhandled exception occurred while running ping") return EMPTY_PING_SCAN diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index bdada64e8..d0c6e3e7a 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -14,13 +14,13 @@ POLL_INTERVAL = 0.5 EMPTY_PORT_SCAN = {-1: PortScanData(-1, PortStatus.CLOSED, None, None)} -def try_scan_tcp_ports( +def scan_tcp_ports( host: str, ports_to_scan: Iterable[int], timeout: float ) -> Mapping[int, PortScanData]: try: return _scan_tcp_ports(host, ports_to_scan, timeout) except Exception: - logging.exception("Unhandled exception occurred while trying to scan tcp ports") + logger.exception("Unhandled exception occurred while trying to scan tcp ports") return EMPTY_PORT_SCAN diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 853e63a95..c7953d83c 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -41,12 +41,12 @@ class Puppet(IPuppet): return pba.run(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: - return network_scanning.try_ping(host, timeout) + return network_scanning.ping(host, timeout) def scan_tcp_ports( self, host: str, ports: List[int], timeout: float = CONNECTION_TIMEOUT ) -> Dict[int, PortScanData]: - return network_scanning.try_scan_tcp_ports(host, ports, timeout) + return network_scanning.scan_tcp_ports(host, ports, timeout) def fingerprint( self, @@ -60,7 +60,7 @@ class Puppet(IPuppet): fingerprinter = self._plugin_registry.get_plugin(name, PluginType.FINGERPRINTER) return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options) except Exception: - logging.exception( + logger.exception( f"Unhandled exception occurred " f"while trying to run {name} fingerprinter" ) return EMPTY_FINGERPRINT diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py index 5079c1ac9..725a3aaa0 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanning.py @@ -1,7 +1,7 @@ import pytest from infection_monkey.i_puppet import PortStatus -from infection_monkey.network_scanning import try_scan_tcp_ports +from infection_monkey.network_scanning import scan_tcp_ports PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] @@ -20,7 +20,7 @@ def patch_check_tcp_ports(monkeypatch, open_ports_data): def test_tcp_successful(monkeypatch, patch_check_tcp_ports, open_ports_data): closed_ports = [8080, 143, 445] - port_scan_data = try_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) assert len(port_scan_data) == 6 for port in open_ports_data.keys(): @@ -37,7 +37,7 @@ 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 = try_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) assert len(port_scan_data) == 6 for port in open_ports_data: @@ -49,6 +49,6 @@ 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 = try_scan_tcp_ports("127.0.0.1", [], 0) + port_scan_data = scan_tcp_ports("127.0.0.1", [], 0) assert len(port_scan_data) == 0