Mypy error fixes

simulation-system/libs/csle-common/src/csle_common/dao/docker/docker_container_metadata.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, Union
 from csle_base.json_serializable import JSONSerializable
 
 
@@ -8,9 +8,9 @@ class DockerContainerMetadata(JSONSerializable):
     """
 
     def __init__(self, name: str, status: str, short_id: str, image_short_id: str, image_tags: list, id: str,
-                 created: str, ip: str, network_id: str, gateway: str, mac: str, ip_prefix_len: int,
+                 created: str, ip: str, network_id: int, gateway: str, mac: str, ip_prefix_len: str,
                  name2: str, level: str, hostname: str, image_name: str, net: str,
-                 dir: str, config_path: str, container_handle: str, emulation: str, kafka_container: str):
+                 dir: Union[str, None], config_path: Union[None, str], container_handle: str, emulation: Union[None, str], kafka_container: Union[None, str]):
         """
         Intializes the DTO
 

simulation-system/libs/csle-common/src/csle_common/dao/docker/docker_env_metadata.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Union
 from csle_common.dao.docker.docker_container_metadata import DockerContainerMetadata
 from csle_common.dao.emulation_config.emulation_env_config import EmulationEnvConfig
 from csle_common.dao.emulation_config.kafka_config import KafkaConfig
@@ -11,7 +11,7 @@ class DockerEnvMetadata(JSONSerializable):
     """
 
     def __init__(self, containers: List[DockerContainerMetadata], name: str, subnet_prefix: str,
-                 subnet_mask: str, level: str, config: EmulationEnvConfig, kafka_config: KafkaConfig):
+                 subnet_mask: str, level: str, config: Union[None, EmulationEnvConfig], kafka_config: Union[None, KafkaConfig]):
         """
         Initializes the DTO
 

simulation-system/libs/csle-common/src/csle_common/dao/emulation_config/emulation_statistics_windowed.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, List, Tuple
 from csle_common.dao.emulation_config.emulation_env_state import EmulationEnvState
 from csle_common.dao.system_identification.emulation_statistics import EmulationStatistics
 from csle_common.dao.emulation_action.attacker.emulation_attacker_action import EmulationAttackerAction
@@ -21,10 +21,11 @@ def __init__(self, window_size: int, emulation_name: str, descr: str):
         :param descr: the description
         """
         self.window_size = window_size
-        self.initial_states = []
-        self.state_transitions = []
-        self.emulation_name = emulation_name
-        self.descr = descr
+        self.initial_states: List[EmulationEnvState] = []
+        self.state_transitions: List[Tuple[EmulationEnvState, EmulationEnvState,
+                                           EmulationEnvState, EmulationEnvState]] = []
+        self.emulation_name: str = emulation_name
+        self.descr: str = descr
         self.emulation_statistics = EmulationStatistics(emulation_name=self.emulation_name, descr=self.descr)
         self.statistics_id = MetastoreFacade.save_emulation_statistic(emulation_statistics=self.emulation_statistics)
 

simulation-system/libs/csle-common/src/csle_common/metastore/metastore_facade.py

@@ -71,7 +71,7 @@ def list_emulations_ids() -> List[Tuple]:
                 return records
 
     @staticmethod
-    def get_emulation_by_name(name: str) -> Union[None, EmulationEnvConfig]:
+    def get_emulation_by_name(name: Union[None, str]) -> Union[None, EmulationEnvConfig]:
         """
         Function for extracting the metadata of an emulation with a given name
 

simulation-system/libs/csle-common/src/csle_common/util/docker_util.py

@@ -23,7 +23,7 @@ def parse_runnning_emulation_infos() -> List[DockerEnvMetadata]:
         client_1 = docker.from_env()
         client_2 = docker.APIClient(base_url=constants.DOCKER.UNIX_DOCKER_SOCK_URL)
         parsed_containers = DockerUtil.parse_running_containers(client_1=client_1, client_2=client_2)
-        emulations = list(set(list(map(lambda x: x.emulation, parsed_containers))))
+        emulations: List[Union[None, str]] = list(set(list(map(lambda x: x.emulation, parsed_containers))))
         parsed_envs = DockerUtil.parse_running_emulation_envs(emulations=emulations, containers=parsed_containers)
         return parsed_envs
 
@@ -72,7 +72,7 @@ def parse_stopped_containers(client_1, client2) -> List[DockerContainerMetadata]
         return parsed_containers
 
     @staticmethod
-    def parse_running_emulation_envs(emulations: List[str], containers: List[DockerContainerMetadata]) \
+    def parse_running_emulation_envs(emulations: List[Union[None, str]], containers: List[DockerContainerMetadata]) \
             -> List[DockerEnvMetadata]:
         """
         Queries docker to get a list of all active emulation environments
@@ -180,7 +180,7 @@ def get_docker_gw_bridge_ip(container_id: str) -> str:
         containers = docker_gw_bridge_info[constants.DOCKER.CONTAINERS_KEY]
         if container_id in containers:
             container = containers[container_id]
-            ip = container[constants.DOCKER.IPV4_KEY]
+            ip: str = container[constants.DOCKER.IPV4_KEY]
             ip = ip.split("/")[0]
             return ip
         raise ValueError(f"The container with id:{container_id} does not have an IP in the docker gw bridge network. "

simulation-system/libs/csle-common/src/csle_common/util/export_util.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Dict, Any
+from typing import Tuple, Dict, Any, Union, List
 import io
 import csv
 import json
@@ -40,13 +40,13 @@ def get_dir_size_gb(dir_path: str = '.') -> float:
         :param dir_path: the path to the directory
         :return: the size of the directory in GB
         """
-        total = 0
+        total = 0.0
         with os.scandir(dir_path) as it:
             for entry in it:
                 if entry.is_file():
                     total += entry.stat().st_size
                 elif entry.is_dir():
-                    total += ExportUtil.get_dir_size(entry.path)
+                    total += ExportUtil.get_dir_size_gb(entry.path)
         return round((float(total) / 1000000000), 2)
 
     @staticmethod
@@ -87,6 +87,8 @@ def export_emulation_traces_to_disk_json(num_traces_per_file: int, output_dir: s
             Logger.__call__().get_logger().info(f"Reading trace {i}/{len(emulation_traces_ids)} from the metastore, "
                                                 f"trace id: {id_obj[0]}")
             tr = MetastoreFacade.get_emulation_trace(id=id_obj[0])
+            if tr is None:
+                raise ValueError("Tr is none and thus has no modules")
             if num_attributes_per_time_step == -1:
                 num_attributes_per_time_step = tr.num_attributes_per_time_step()
             if schema is None:
@@ -111,7 +113,7 @@ def export_emulation_traces_to_disk_json(num_traces_per_file: int, output_dir: s
         num_traces = len(emulation_traces_ids)
         with io.open(f"{output_dir}{constants.COMMANDS.SLASH_DELIM}{constants.DATASETS.METADATA_FILE_NAME}", 'w',
                      encoding='utf-8') as f:
-            metadata_dict = {}
+            metadata_dict: Dict[str, Any] = {}
             metadata_dict[constants.DATASETS.FILE_FORMAT_PROPERTY] = file_format
             metadata_dict[constants.DATASETS.NUM_TRACES_PROPERTY] = num_traces
             metadata_dict[constants.DATASETS.NUM_ATTRIBUTES_PER_TIME_STEP_PROPERTY] = num_attributes_per_time_step
@@ -138,7 +140,7 @@ def extract_emulation_traces_dataset_metadata(dir_path: str, zip_file_path: str)
         file_format = "unknown"
         added_by = "unknown"
         num_traces = -1
-        schema = ""
+        schema = {}
         columns = ""
         num_traces_per_file = -1
         num_attributes_per_time_step = -1
@@ -173,7 +175,7 @@ def extract_emulation_traces_dataset_metadata(dir_path: str, zip_file_path: str)
             if constants.DATASETS.COLUMNS_PROPERTY in metadata_dict:
                 columns = metadata_dict[constants.DATASETS.COLUMNS_PROPERTY]
             else:
-                columns = []
+                columns = ""
 
         num_files = len([name for name in os.listdir(dir_path) if os.path.isfile(os.path.join(dir_path, name))]) - 1
         size_compressed_gb = round(float(os.path.getsize(zip_file_path)) / 1000000000, 2)
@@ -213,10 +215,12 @@ def export_emulation_traces_to_disk_csv(
         file_name = f"{file_id}.csv"
         last_export = 0
         num_attributes_per_time_step = -1
-        columns = None
+        columns: Union[List[Union[str, int, float]], None]= None
         for i, id_obj in enumerate(emulation_traces_ids):
             Logger.__call__().get_logger().info(f"Reading trace {i}/{len(emulation_traces_ids)} from the metastore")
             tr = MetastoreFacade.get_emulation_trace(id=id_obj[0])
+            if tr is None:
+                raise ValueError("tr is None and thus has no modules/attributes")
             tr_values, tr_labels = tr.to_csv_record(max_time_steps=max_time_steps, max_nodes=max_nodes,
                                                     max_ports=max_ports, max_vulns=max_vulns, null_value=null_value)
             if num_attributes_per_time_step == -1:
@@ -240,7 +244,7 @@ def export_emulation_traces_to_disk_csv(
         num_traces = len(emulation_traces_ids)
         with io.open(f"{output_dir}{constants.COMMANDS.SLASH_DELIM}{constants.DATASETS.METADATA_FILE_NAME}", 'w',
                      encoding='utf-8') as f:
-            metadata_dict = {}
+            metadata_dict: Dict[str, Any] = {}
             metadata_dict[constants.DATASETS.FILE_FORMAT_PROPERTY] = file_format
             metadata_dict[constants.DATASETS.NUM_TRACES_PROPERTY] = num_traces
             metadata_dict[constants.DATASETS.NUM_ATTRIBUTES_PER_TIME_STEP_PROPERTY] = num_attributes_per_time_step
@@ -270,6 +274,8 @@ def export_emulation_statistics_to_disk_json(output_dir: str, zip_file_output: s
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
         emulation_statistic = MetastoreFacade.get_emulation_statistic(id=statistics_id)
+        if emulation_statistic is None:
+            raise ValueError("emulation_statistic is None and thus has no modules/attributes")
         file_name = "statistics.json"
         Logger.__call__().get_logger().info(f"Exporting statistics with id {statistics_id} to file: {file_name}")
         emulation_statistic.compute_descriptive_statistics_and_distributions()
@@ -278,14 +284,14 @@ def export_emulation_statistics_to_disk_json(output_dir: str, zip_file_output: s
         metrics = ",".join(emulation_statistic.metrics)
         conditions = ",".join(emulation_statistic.conditions)
         num_conditions = emulation_statistic.num_conditions
-        statistics_dict = emulation_statistic.to_dict()
+        statistics_dict: Dict[str, Any] = emulation_statistic.to_dict()
         statistics_dict = json.dumps(statistics_dict, indent=4, sort_keys=True)
         with io.open(f"{output_dir}{constants.COMMANDS.SLASH_DELIM}{file_name}", 'w', encoding='utf-8') as f:
             f.write(statistics_dict)
         file_format = "json"
         with io.open(f"{output_dir}{constants.COMMANDS.SLASH_DELIM}{constants.DATASETS.METADATA_FILE_NAME}", 'w',
                      encoding='utf-8') as f:
-            metadata_dict = {}
+            metadata_dict: Dict[str, Any] = {}
             metadata_dict[constants.DATASETS.FILE_FORMAT_PROPERTY] = file_format
             metadata_dict[constants.DATASETS.NUM_MEASUREMENTS_PROPERTY] = num_measurements
             metadata_dict[constants.DATASETS.NUM_CONDITIONS_PROPERTY] = num_conditions

simulation-system/libs/csle-common/src/csle_common/util/import_util.py

@@ -3,15 +3,15 @@
 from csle_common.metastore.metastore_facade import MetastoreFacade
 from csle_common.dao.system_identification.emulation_statistics import EmulationStatistics
 from csle_common.dao.emulation_config.emulation_trace import EmulationTrace
-
+from typing import Union
 
 class ImportUtil:
     """
     Class with utility functions for importing data to the metastore
     """
 
     @staticmethod
-    def import_emulation_statistics_from_disk_json(input_file: str, emulation_name: str = None) -> None:
+    def import_emulation_statistics_from_disk_json(input_file: str, emulation_name: Union[str, None] = None) -> None:
         """
         Imports emulation statistics from disk to the metastore
 
@@ -30,7 +30,7 @@ def import_emulation_statistics_from_disk_json(input_file: str, emulation_name:
                                             f"input file:{input_file}")
 
     @staticmethod
-    def import_emulation_traces_from_disk_json(input_file: str, emulation_name: str = None) -> None:
+    def import_emulation_traces_from_disk_json(input_file: str, emulation_name: Union[str, None] = None) -> None:
         """
         Imports emulation traces from disk to the metastore