Merge pull request #217 from Limmen/dev1

emulation_vulnerability_observation_state.py
Limmen · Aug 17, 2023 · b764cb7 · b764cb7
2 parents fecdfd3 + 4e53dbe
commit b764cb7
Showing 1 changed file with 20 additions and 8 deletions.
diff --git a/...csle_common/dao/emulation_observation/common/emulation_vulnerability_observation_state.py b/...csle_common/dao/emulation_observation/common/emulation_vulnerability_observation_state.py
@@ -37,23 +37,27 @@ def __init__(self, name: str, port: int,
         self.credentials = credentials
         self.service = service
         if service is None or service == "":
-            for cr in self.credentials:
-                if cr.service is not None and cr.service != "":
-                    self.service = cr.service
+            if self.credentials is not None:
+                for cr in self.credentials:
+                    if cr.service is not None and cr.service != "":
+                        self.service = cr.service
+            else:
+                raise ValueError("Credentials is None")
 
     def to_dict(self) -> Dict[str, Any]:
         """
         Converts the object to a dict representation
         
         :return: a dict representation of the object
         """
-        d = {}
+        d: Dict[str, Any] = {}
         d["name"] = self.name
         d["protocol"] = self.protocol
         d["cvss"] = self.cvss
         d["osvdb_id"] = self.osvdb_id
         d["description"] = self.description
-        d["credentials"] = list(map(lambda x: x.to_dict(), self.credentials))
+        d["credentials"] = list(map(lambda x: x.to_dict(), self.credentials)) if self.credentials \
+            is not None else self.credentials
         d["service"] = self.service
         d["port"] = self.port
         return d
@@ -84,7 +88,8 @@ def __str__(self):
         """
         return "name:{}, port:{}, protocol:{}, cvss:{}, osvdb_id:{}, desc:{}, service:{}, credentials:{}".format(
             self.name, self.port, self.protocol, self.cvss, self.osvdb_id, self.description, self.service,
-            list(map(lambda x: str(x), self.credentials)))
+            list(map(lambda x: str(x), self.credentials)) if self.credentials is not None \
+                else self.credentials)
 
     def to_vulnerability(self) -> NodeVulnerabilityConfig:
         """
@@ -103,8 +108,11 @@ def to_network_services(self) -> List[NetworkService]:
 
         :return: the network service representation
         """
-        services = [NetworkService(protocol=self.protocol, port=self.port, name=self.service,
+        services = [NetworkService(protocol=self.protocol if self.protocol is not None else "",
+                                   port=self.port, name=self.service if self.service is not None else "",
                                    credentials=self.credentials)]
+        if self.credentials is None:
+            raise ValueError("credentials not iterable")
         for cr in self.credentials:
             new_service = True
             for s in services:
@@ -132,8 +140,12 @@ def num_attributes(self) -> int:
         """
         :return: The number of attribute of the DTO
         """
-        if len(self.credentials) > 0:
+        if self.credentials is None:
+            return 0
+        elif len(self.credentials) > 0:
             return 7 + len(self.credentials) * self.credentials[0].num_attributes()
+        else:
+            return 0
 
     @staticmethod
     def schema() -> "EmulationVulnerabilityObservationState":