Add More Typehint

src/primaite/nodes/active_node.py

@@ -24,7 +24,7 @@ class ActiveNode(Node):
         software_state: SoftwareState,
         file_system_state: FileSystemState,
         config_values: TrainingConfig,
-    ):
+    ) -> None:
         """
         Initialise an active node.
 
@@ -60,7 +60,7 @@ class ActiveNode(Node):
         return self._software_state
 
     @software_state.setter
-    def software_state(self, software_state: SoftwareState):
+    def software_state(self, software_state: SoftwareState) -> None:
         """
         Get the software_state.
 
@@ -79,7 +79,7 @@ class ActiveNode(Node):
                 f"Node.software_state:{self._software_state}"
             )
 
-    def set_software_state_if_not_compromised(self, software_state: SoftwareState):
+    def set_software_state_if_not_compromised(self, software_state: SoftwareState) -> None:
         """
         Sets Software State if the node is not compromised.
 
@@ -99,14 +99,14 @@ class ActiveNode(Node):
                 f"Node.software_state:{self._software_state}"
             )
 
-    def update_os_patching_status(self):
+    def update_os_patching_status(self) -> None:
         """Updates operating system status based on patching cycle."""
         self.patching_count -= 1
         if self.patching_count <= 0:
             self.patching_count = 0
             self._software_state = SoftwareState.GOOD
 
-    def set_file_system_state(self, file_system_state: FileSystemState):
+    def set_file_system_state(self, file_system_state: FileSystemState) -> None:
         """
         Sets the file system state (actual and observed).
 
@@ -133,7 +133,7 @@ class ActiveNode(Node):
                 f"Node.file_system_state.actual:{self.file_system_state_actual}"
             )
 
-    def set_file_system_state_if_not_compromised(self, file_system_state: FileSystemState):
+    def set_file_system_state_if_not_compromised(self, file_system_state: FileSystemState) -> None:
         """
         Sets the file system state (actual and observed) if not in a compromised state.
 
@@ -166,12 +166,12 @@ class ActiveNode(Node):
                 f"Node.file_system_state.actual:{self.file_system_state_actual}"
             )
 
-    def start_file_system_scan(self):
+    def start_file_system_scan(self) -> None:
         """Starts a file system scan."""
         self.file_system_scanning = True
         self.file_system_scanning_count = self.config_values.file_system_scanning_limit
 
-    def update_file_system_state(self):
+    def update_file_system_state(self) -> None:
         """Updates file system status based on scanning/restore/repair cycle."""
         # Deprecate both the action count (for restoring or reparing) and the scanning count
         self.file_system_action_count -= 1
@@ -193,14 +193,14 @@ class ActiveNode(Node):
             self.file_system_scanning = False
             self.file_system_scanning_count = 0
 
-    def update_resetting_status(self):
+    def update_resetting_status(self) -> None:
         """Updates the reset count & makes software and file state to GOOD."""
         super().update_resetting_status()
         if self.resetting_count <= 0:
             self.file_system_state_actual = FileSystemState.GOOD
             self.software_state = SoftwareState.GOOD
 
-    def update_booting_status(self):
+    def update_booting_status(self) -> None:
         """Updates the booting software and file state to GOOD."""
         super().update_booting_status()
         if self.booting_count <= 0:

src/primaite/nodes/node.py

@@ -38,40 +38,40 @@ class Node:
         self.booting_count: int = 0
         self.shutting_down_count: int = 0
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Returns the name of the node."""
         return self.name
 
-    def turn_on(self):
+    def turn_on(self) -> None:
         """Sets the node state to ON."""
         self.hardware_state = HardwareState.BOOTING
         self.booting_count = self.config_values.node_booting_duration
 
-    def turn_off(self):
+    def turn_off(self) -> None:
         """Sets the node state to OFF."""
         self.hardware_state = HardwareState.OFF
         self.shutting_down_count = self.config_values.node_shutdown_duration
 
-    def reset(self):
+    def reset(self) -> None:
         """Sets the node state to Resetting and starts the reset count."""
         self.hardware_state = HardwareState.RESETTING
         self.resetting_count = self.config_values.node_reset_duration
 
-    def update_resetting_status(self):
+    def update_resetting_status(self) -> None:
         """Updates the resetting count."""
         self.resetting_count -= 1
         if self.resetting_count <= 0:
             self.resetting_count = 0
             self.hardware_state = HardwareState.ON
 
-    def update_booting_status(self):
+    def update_booting_status(self) -> None:
         """Updates the booting count."""
         self.booting_count -= 1
         if self.booting_count <= 0:
             self.booting_count = 0
             self.hardware_state = HardwareState.ON
 
-    def update_shutdown_status(self):
+    def update_shutdown_status(self) -> None:
         """Updates the shutdown count."""
         self.shutting_down_count -= 1
         if self.shutting_down_count <= 0:

src/primaite/nodes/node_state_instruction_green.py

@@ -1,5 +1,9 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
 """Defines node behaviour for Green PoL."""
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from primaite.common.enums import HardwareState, NodePOLType, SoftwareState
 
 
 class NodeStateInstructionGreen(object):
@@ -7,10 +11,10 @@ class NodeStateInstructionGreen(object):
 
     def __init__(
         self,
-        _id,
-        _start_step,
-        _end_step,
-        _node_id,
+        _id: str,
+        _start_step: int,
+        _end_step: int,
+        _node_id: str,
         _node_pol_type,
         _service_name,
         _state,
@@ -30,9 +34,10 @@ class NodeStateInstructionGreen(object):
         self.start_step = _start_step
         self.end_step = _end_step
         self.node_id = _node_id
-        self.node_pol_type = _node_pol_type
-        self.service_name = _service_name  # Not used when not a service instruction
-        self.state = _state
+        self.node_pol_type: "NodePOLType" = _node_pol_type
+        self.service_name: str = _service_name  # Not used when not a service instruction
+        # TODO: confirm type of state
+        self.state: Union["HardwareState", "SoftwareState"] = _state
 
     def get_start_step(self):
         """

src/primaite/nodes/passive_node.py

@@ -16,7 +16,7 @@ class PassiveNode(Node):
         priority: Priority,
         hardware_state: HardwareState,
         config_values: TrainingConfig,
-    ):
+    ) -> None:
         """
         Initialise a passive node.
 

src/primaite/nodes/service_node.py

@@ -25,7 +25,7 @@ class ServiceNode(ActiveNode):
         software_state: SoftwareState,
         file_system_state: FileSystemState,
         config_values: TrainingConfig,
-    ):
+    ) -> None:
         """
         Initialise a Service Node.
 
@@ -52,7 +52,7 @@ class ServiceNode(ActiveNode):
         )
         self.services: Dict[str, Service] = {}
 
-    def add_service(self, service: Service):
+    def add_service(self, service: Service) -> None:
         """
         Adds a service to the node.
 
@@ -102,7 +102,7 @@ class ServiceNode(ActiveNode):
                     return False
         return False
 
-    def set_service_state(self, protocol_name: str, software_state: SoftwareState):
+    def set_service_state(self, protocol_name: str, software_state: SoftwareState) -> None:
         """
         Sets the software_state of a service (protocol) on the node.
 
@@ -131,7 +131,7 @@ class ServiceNode(ActiveNode):
                 f"Node.services[<key>].software_state:{software_state}"
             )
 
-    def set_service_state_if_not_compromised(self, protocol_name: str, software_state: SoftwareState):
+    def set_service_state_if_not_compromised(self, protocol_name: str, software_state: SoftwareState) -> None:
         """
         Sets the software_state of a service (protocol) on the node.
 
@@ -158,7 +158,7 @@ class ServiceNode(ActiveNode):
                 f"Node.services[<key>].software_state:{software_state}"
             )
 
-    def get_service_state(self, protocol_name):
+    def get_service_state(self, protocol_name: str) -> SoftwareState:
         """
         Gets the state of a service.
 
@@ -169,20 +169,20 @@ class ServiceNode(ActiveNode):
         if service_value:
             return service_value.software_state
 
-    def update_services_patching_status(self):
+    def update_services_patching_status(self) -> None:
         """Updates the patching counter for any service that are patching."""
         for service_key, service_value in self.services.items():
             if service_value.software_state == SoftwareState.PATCHING:
                 service_value.reduce_patching_count()
 
-    def update_resetting_status(self):
+    def update_resetting_status(self) -> None:
         """Update resetting counter and set software state if it reached 0."""
         super().update_resetting_status()
         if self.resetting_count <= 0:
             for service in self.services.values():
                 service.software_state = SoftwareState.GOOD
 
-    def update_booting_status(self):
+    def update_booting_status(self) -> None:
         """Update booting counter and set software to good if it reached 0."""
         super().update_booting_status()
         if self.booting_count <= 0: