Agent: Improve method naming and exception handling

This commit is contained in:
vakarisz 2022-04-08 10:36:14 +03:00
parent 6157ffee76
commit 45c6cac60c
6 changed files with 14 additions and 14 deletions

View File

@ -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}"
)

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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