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

#917 - Integrated the PrimaiteSession into all tests.

Browse Source 
- Ran a full pre-commit hook and thus encountered tons of fixes required
This commit is contained in:
Chris McCarthy
2023-06-30 09:08:13 +01:00
parent 7f912df383
commit 73015802ec
62 changed files with 1880 additions and 802 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,10 +6,17 @@ from networkx import MultiGraph, shortest_path
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import HardwareState, NodePOLType, NodeType, SoftwareState
from primaite.common.enums import (
HardwareState,
NodePOLType,
NodeType,
SoftwareState,
)
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_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
@@ -190,7 +197,9 @@ def apply_iers(
link_id = edge_dict[0].get("id")
link = links[link_id]
# Check whether the new load exceeds the bandwidth
if (link.get_current_load() + load) > link.get_bandwidth():
if (
link.get_current_load() + load
) > link.get_bandwidth():
link_capacity_exceeded = True
if _VERBOSE:
print("Link capacity exceeded")
@@ -204,7 +213,8 @@ def apply_iers(
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(
path_node_list[count], path_node_list[count + 1]
path_node_list[count],
path_node_list[count + 1],
)
link_id = edge_dict[0].get("id")
link = links[link_id]
@@ -216,7 +226,9 @@ def apply_iers(
else:
# One of the nodes is not operational
if _VERBOSE:
print("Path not valid - one or more nodes not operational")
print(
"Path not valid - one or more nodes not operational"
)
pass
else:
@@ -231,7 +243,9 @@ def apply_iers(
def apply_node_pol(
nodes: Dict[str, NodeUnion],
node_pol: Dict[any, Union[NodeStateInstructionGreen, NodeStateInstructionRed]],
node_pol: Dict[
any, Union[NodeStateInstructionGreen, NodeStateInstructionRed]
],
step: int,
):
"""
@@ -263,16 +277,22 @@ def apply_node_pol(
elif node_pol_type == NodePOLType.OS:
# Change OS state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if isinstance(node, ActiveNode) or isinstance(
node, ServiceNode
):
node.set_software_state_if_not_compromised(state)
elif node_pol_type == NodePOLType.SERVICE:
# Change a service state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ServiceNode):
node.set_service_state_if_not_compromised(service_name, state)
node.set_service_state_if_not_compromised(
service_name, state
)
else:
# Change the file system status
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if isinstance(node, ActiveNode) or isinstance(
node, ServiceNode
):
node.set_file_system_state_if_not_compromised(state)
else:
# PoL is not valid in this time step

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

@@ -176,7 +176,9 @@ def apply_red_agent_iers(
link_id = edge_dict[0].get("id")
link = links[link_id]
# Check whether the new load exceeds the bandwidth
if (link.get_current_load() + load) > link.get_bandwidth():
if (
link.get_current_load() + load
) > link.get_bandwidth():
link_capacity_exceeded = True
if _VERBOSE:
print("Link capacity exceeded")
@@ -190,7 +192,8 @@ def apply_red_agent_iers(
while count < path_node_list_length - 1:
# Get the link between the next two nodes
edge_dict = network.get_edge_data(
path_node_list[count], path_node_list[count + 1]
path_node_list[count],
path_node_list[count + 1],
)
link_id = edge_dict[0].get("id")
link = links[link_id]
@@ -200,16 +203,23 @@ def apply_red_agent_iers(
# This IER is now valid, so set it to running
ier_value.set_is_running(True)
if _VERBOSE:
print("Red IER was allowed to run in step " + str(step))
print(
"Red IER was allowed to run in step "
+ str(step)
)
else:
# One of the nodes is not operational
if _VERBOSE:
print("Path not valid - one or more nodes not operational")
print(
"Path not valid - one or more nodes not operational"
)
pass
else:
if _VERBOSE:
print("Red IER was NOT allowed to run in step " + str(step))
print(
"Red IER was NOT allowed to run in step " + str(step)
)
print("Source, Dest or ACL were not valid")
pass
# ------------------------------------
@@ -264,7 +274,9 @@ def apply_red_agent_node_pol(
passed_checks = True
elif initiator == NodePOLInitiator.IER:
# Need to check there is a red IER incoming
passed_checks = is_red_ier_incoming(target_node, iers, pol_type)
passed_checks = is_red_ier_incoming(
target_node, iers, pol_type
)
elif initiator == NodePOLInitiator.SERVICE:
# Need to check the condition of a service on another node
source_node = nodes[source_node_id]
@@ -308,7 +320,9 @@ def apply_red_agent_node_pol(
target_node.set_file_system_state(state)
else:
if _VERBOSE:
print("Node Red Agent PoL not allowed - did not pass checks")
print(
"Node Red Agent PoL not allowed - did not pass checks"
)
else:
# PoL is not valid in this time step
pass
@@ -323,7 +337,10 @@ def is_red_ier_incoming(node, iers, node_pol_type):
node_id = node.node_id
for ier_key, ier_value in iers.items():
if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id:
if (
ier_value.get_is_running()
and ier_value.get_dest_node_id() == node_id
):
if (
node_pol_type == NodePOLType.OPERATING
or node_pol_type == NodePOLType.OS