Merge branch 'master' of https://github.com/Limmen/csle

Limmen · Aug 16, 2023 · f9c2833 · f9c2833
2 parents bb15625 + 6e435fc
commit f9c2833
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/...common/dao/emulation_observation/attacker/emulation_attacker_machine_observation_state.py b/...common/dao/emulation_observation/attacker/emulation_attacker_machine_observation_state.py
@@ -1,4 +1,4 @@
-from typing import List, Set, Dict, Any
+from typing import List, Set, Dict, Any, Optional
 import copy
 from csle_common.dao.emulation_config.flag import Flag
 from csle_common.dao.emulation_config.credential import Credential
@@ -43,7 +43,7 @@ def __init__(self, ips: List[str]):
         self.logged_in_services: List[str] = []
         self.root_services: List[str] = []
         self.hostnames: List[str] = []
-        self.trace: NmapTrace = None
+        self.trace: Optional[NmapTrace] = None
         self.telnet_brute_tried = False
         self.ssh_brute_tried = False
         self.ftp_brute_tried = False
@@ -126,7 +126,7 @@ def sort_shell_access(self, service_lookup) -> None:
             key=lambda x: service_lookup[x.service] if x.service is not None else x.username,
             reverse=False)
 
-    def _vuln_lookup(self, name, lookup_table) -> int:
+    def _vuln_lookup(self, name: str, lookup_table: Dict[str, int]) -> int:
         """
         Looks up the id of a vulnerability in a lookup table
 
@@ -216,6 +216,12 @@ def from_nmap_result(nmap_host_result: NmapHostResult) -> "EmulationAttackerMach
 
         :return: the created AttackerMachineObservationState
         """
+        if nmap_host_result.ips is None or \
+            nmap_host_result.ports is None or \
+                nmap_host_result.vulnerabilities is None or \
+                    nmap_host_result.credentials is None or \
+                        nmap_host_result.hostnames is None:
+            raise ValueError("NMapHostResult is incomplete, at least one attribute is None")
         m_obs = EmulationAttackerMachineObservationState(ips=nmap_host_result.ips)
         ports = list(map(lambda x: x.to_obs(), nmap_host_result.ports))
         m_obs.ports = ports
@@ -301,7 +307,7 @@ def to_dict(self) -> Dict[str, Any]:
 
         :return: a dict representation of the object
         """
-        d = {}
+        d: Dict[str, Any] = {}
         d["ips"] = self.ips
         d["os"] = self.os
         d["ports"] = list(map(lambda x: x.to_dict(), self.ports))
@@ -405,7 +411,7 @@ def schema() -> "EmulationAttackerMachineObservationState":
         dto.osvdb_vulns = [EmulationVulnerabilityObservationState.schema()]
         dto.shell_access_credentials = [Credential.schema()]
         dto.backdoor_credentials = [Credential.schema()]
-        dto.flags_found.add("")
+        dto.flags_found.add(Flag.schema())
         dto.ssh_connections = [EmulationConnectionObservationState.schema()]
         dto.ftp_connections = [EmulationConnectionObservationState.schema()]
         dto.telnet_connections = [EmulationConnectionObservationState.schema()]