Fix types according to mypy

src/primaite/pol/green_pol.py

@@ -1,6 +1,6 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
 """Implements Pattern of Life on the network (nodes and links)."""
-from typing import Dict, Union
+from typing import Dict
 
 from networkx import MultiGraph, shortest_path
 
@@ -10,7 +10,6 @@ from primaite.common.enums import HardwareState, NodePOLType, NodeType, Software
 from primaite.links.link import Link
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
-from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
 from primaite.nodes.service_node import ServiceNode
 from primaite.pol.ier import IER
 
@@ -65,6 +64,8 @@ def apply_iers(
             dest_node = nodes[dest_node_id]
 
             # 1. Check the source node situation
+            # TODO: should be using isinstance rather than checking node type attribute. IE. just because it's a switch
+            # doesn't mean it has a software state? It could be a PassiveNode or ActiveNode
             if source_node.node_type == NodeType.SWITCH:
                 # It's a switch
                 if (
@@ -215,7 +216,7 @@ def apply_iers(
 
 def apply_node_pol(
     nodes: Dict[str, NodeUnion],
-    node_pol: Dict[any, Union[NodeStateInstructionGreen, NodeStateInstructionRed]],
+    node_pol: Dict[str, NodeStateInstructionGreen],
     step: int,
 ) -> None:
     """

src/primaite/pol/red_agent_pol.py

@@ -74,6 +74,9 @@ def apply_red_agent_iers(
                 pass
             else:
                 # It's not a switch or an actuator (so active node)
+                # TODO: this occurs after ruling out the possibility that the node is a switch or an actuator, but it
+                # could still be a passive/active node, therefore it won't have a hardware_state. The logic here needs
+                # to change according to duck typing.
                 if source_node.hardware_state == HardwareState.ON:
                     if source_node.has_service(protocol):
                         # Red agents IERs can only be valid if the source service is in a compromised state