oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#1355 - Renamed the NodeType custom type in custom_typing.py as it clased with the NodeType enum in enums.py

Browse Source
This commit is contained in:
Chris McCarthy
2023-05-26 09:43:37 +01:00
parent 6245ad9298
commit 05ebd15053
4 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/common/custom_typing.py
View File

@@ -4,5 +4,5 @@ from primaite.nodes.active_node import ActiveNode
from primaite.nodes.passive_node import PassiveNode
from primaite.nodes.service_node import ServiceNode
NodeType: Type = Union[ActiveNode, PassiveNode, ServiceNode]
NodeUnion: Type = Union[ActiveNode, PassiveNode, ServiceNode]
"""A Union of ActiveNode, PassiveNode, and ServiceNode."""

8
src/primaite/environment/primaite_env.py
View File

@@ -15,7 +15,7 @@ from gym import Env, spaces
from matplotlib import pyplot as plt
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeType
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import (
ActionType,
FileSystemState,
@@ -80,10 +80,10 @@ class Primaite(Env):
self.agent_identifier = self.config_values.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes: Dict[str, NodeType] = {}
self.nodes: Dict[str, NodeUnion] = {}
# Create a dictionary to hold a reference set of nodes
self.nodes_reference: Dict[str, NodeType] = {}
self.nodes_reference: Dict[str, NodeUnion] = {}
# Create a dictionary to hold all the links
self.links: Dict[str, Link] = {}
@@ -1128,7 +1128,7 @@ class Primaite(Env):
node_class = item["node_class"]
node_hardware_state: HardwareState = HardwareState[item["hardware_state"]]
node: NodeType = self.nodes[node_id]
node: NodeUnion = self.nodes[node_id]
node_ref = self.nodes_reference[node_id]
# Reset the hardware state (common for all node types)

6
src/primaite/pol/green_pol.py
View File

@@ -5,7 +5,7 @@ from typing import Dict, Union
from networkx import MultiGraph, shortest_path
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeType
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import HardwareState, NodePOLType, NodeType, SoftwareState
from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode
@@ -19,7 +19,7 @@ _VERBOSE = False
def apply_iers(
network: MultiGraph,
nodes: Dict[str, NodeType],
nodes: Dict[str, NodeUnion],
links: Dict[str, Link],
iers: Dict[str, IER],
acl: AccessControlList,
@@ -230,7 +230,7 @@ def apply_iers(
def apply_node_pol(
nodes: Dict[str, NodeType],
nodes: Dict[str, NodeUnion],
node_pol: Dict[any, Union[NodeStateInstructionGreen, NodeStateInstructionRed]],
step: int,
):

8
src/primaite/pol/red_agent_pol.py
View File

@@ -5,7 +5,7 @@ from typing import Dict
from networkx import MultiGraph, shortest_path
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeType
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import (
HardwareState,
NodePOLInitiator,
@@ -24,7 +24,7 @@ _VERBOSE = False
def apply_red_agent_iers(
network: MultiGraph,
nodes: Dict[str, NodeType],
nodes: Dict[str, NodeUnion],
links: Dict[str, Link],
iers: Dict[str, IER],
acl: AccessControlList,
@@ -221,7 +221,7 @@ def apply_red_agent_iers(
def apply_red_agent_node_pol(
nodes: NodeType,
nodes: Dict[str, NodeUnion],
iers: Dict[str, IER],
node_pol: Dict[str, NodeStateInstructionRed],
step: int,
@@ -256,7 +256,7 @@ def apply_red_agent_node_pol(
if step >= start_step and step <= stop_step:
# continue --------------------------
target_node: NodeType = nodes[target_node_id]
target_node: NodeUnion = nodes[target_node_id]
# Based the action taken on the initiator type
if initiator == NodePOLInitiator.DIRECT: