From a4a9de6a8d73fe077fde0a4047220e2a85c43abb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 23 Nov 2021 07:16:06 -0500 Subject: [PATCH] Agent: Add a timeout parameter to scan_tcp_port() --- monkey/infection_monkey/i_puppet.py | 3 ++- monkey/infection_monkey/puppet/mock_puppet.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index d44f827eb..c10731d8f 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -42,11 +42,12 @@ class IPuppet(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def scan_tcp_port(self, host: str, port: int) -> PortScanData: + def scan_tcp_port(self, host: str, port: int, timeout: int) -> PortScanData: """ Scans a TCP port on a remote host :param str host: The domain name or IP address of a host :param int port: A TCP port number to scan + :param int timeout: The maximum amount of time (in seconds) to wait for a response :return: The data collected by scanning the provided host:port combination :rtype: PortScanData """ diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 910c3a636..0652c109b 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/monkey/infection_monkey/puppet/mock_puppet.py @@ -161,8 +161,8 @@ class MockPuppet(IPuppet): return (False, None) - def scan_tcp_port(self, host: str, port: int) -> PortScanData: - logger.debug(f"run_scan_tcp_port({host}, {port})") + def scan_tcp_port(self, host: str, port: int, timeout: int = 3) -> PortScanData: + logger.debug(f"run_scan_tcp_port({host}, {port}, {timeout})") dot_1_results = { 22: PortScanData(22, PortStatus.CLOSED, None, None), 445: PortScanData(445, PortStatus.OPEN, "SMB BANNER", "tcp-445"),