From f9b803b1ae9a537e8dc498f7fc2bd23e7aadfdde Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 10 Feb 2022 07:08:02 -0500 Subject: [PATCH] Agent: Minor code quality improvements to SSHFingerprinter --- .../network/ssh_fingerprinter.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/monkey/infection_monkey/network/ssh_fingerprinter.py b/monkey/infection_monkey/network/ssh_fingerprinter.py index 14c9ae70e..32aa20ad9 100644 --- a/monkey/infection_monkey/network/ssh_fingerprinter.py +++ b/monkey/infection_monkey/network/ssh_fingerprinter.py @@ -1,38 +1,40 @@ import re -from typing import Dict, Tuple, Union +from typing import Dict, Optional, Tuple -from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData +from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PingScanData, PortScanData -SSH_PORT = 22 SSH_REGEX = r"SSH-\d\.\d-OpenSSH" LINUX_DIST_SSH = ["ubuntu", "debian"] +DISPLAY_NAME = "SSH" class SSHFingerprinter(IFingerprinter): - _SCANNED_SERVICE = "SSH" - def __init__(self): self._banner_regex = re.compile(SSH_REGEX, re.IGNORECASE) def get_host_fingerprint( - self, host: str, _ping_scan_data, port_scan_data: Dict[int, PortScanData], _options + self, + host: str, + _ping_scan_data: PingScanData, + port_scan_data: Dict[int, PortScanData], + _options: Dict, ) -> FingerprintData: os_type = None os_version = None services = {} - for ps_data in list(port_scan_data.values()): + for ps_data in port_scan_data.values(): if ps_data.banner and self._banner_regex.search(ps_data.banner): os_type, os_version = self._get_host_os(ps_data.banner) services[f"tcp-{ps_data.port}"] = { - "display_name": SSHFingerprinter._SCANNED_SERVICE, + "display_name": DISPLAY_NAME, "port": ps_data.port, "name": "ssh", } return FingerprintData(os_type, os_version, services) @staticmethod - def _get_host_os(banner) -> Tuple[Union[str, None], Union[str, None]]: + def _get_host_os(banner) -> Tuple[Optional[str], Optional[str]]: os = None os_version = None for dist in LINUX_DIST_SSH: