From dd5b796bfe7ed845533e5dcca12d4a030d65ec29 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Mon, 26 Sep 2022 17:48:27 +0000
Subject: [PATCH] Agent: Convert PingScanData to dataclass
---
monkey/infection_monkey/i_puppet/i_puppet.py | 8 +++++++-
monkey/infection_monkey/network_scanning/ping_scanner.py | 7 ++-----
2 files changed, 9 insertions(+), 6 deletions(-)
diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py
index 88e790520..73afa1227 100644
--- a/monkey/infection_monkey/i_puppet/i_puppet.py
+++ b/monkey/infection_monkey/i_puppet/i_puppet.py
@@ -5,6 +5,7 @@ from dataclasses import dataclass
from enum import Enum
from typing import Dict, Iterable, Mapping, Optional, Sequence
+from common import OperatingSystem
from common.credentials import Credentials
from infection_monkey.model import VictimHost
@@ -31,7 +32,12 @@ class ExploiterResultData:
error_message: str = ""
-PingScanData = namedtuple("PingScanData", ["response_received", "os"])
+@dataclass
+class PingScanData:
+ response_received: bool
+ os: Optional[OperatingSystem]
+
+
PortScanData = namedtuple("PortScanData", ["port", "status", "banner", "service"])
FingerprintData = namedtuple("FingerprintData", ["os_type", "os_version", "services"])
PostBreachData = namedtuple("PostBreachData", ["display_name", "command", "result"])
diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py
index 8b30191b4..c181d784b 100644
--- a/monkey/infection_monkey/network_scanning/ping_scanner.py
+++ b/monkey/infection_monkey/network_scanning/ping_scanner.py
@@ -78,11 +78,8 @@ def _process_ping_command_output(ping_command_output: str) -> PingScanData:
# match at all if the group isn't found or the contents of the group are not only digits.
ttl = int(ttl_match.group(1))
- operating_system = None
- if ttl <= LINUX_TTL:
- operating_system = OperatingSystem.LINUX
- else: # as far we we know, could also be OSX/BSD, but lets handle that when it comes up.
- operating_system = OperatingSystem.WINDOWS
+ # could also be OSX/BSD, but lets handle that when it comes up.
+ operating_system = OperatingSystem.LINUX if ttl <= LINUX_TTL else OperatingSystem.WINDOWS
return PingScanData(True, operating_system)