diff --git a/docs/source/about.rst b/docs/source/about.rst index 47511c1b..1f4669fe 100644 --- a/docs/source/about.rst +++ b/docs/source/about.rst @@ -31,7 +31,7 @@ An inheritance model has been adopted in order to model nodes. All nodes have th * Name * Type (e.g. computer, switch, RTU - enumeration) * Priority (P1, P2, P3, P4 or P5 - enumeration) -* Hardware State (ON, OFF, RESETTING - enumeration) +* Hardware State (ON, OFF, RESETTING, SHUTTING_DOWN, BOOTING - enumeration) Active Nodes also have the following attributes (Class: Active Node): @@ -106,6 +106,8 @@ The status changes that can be made to a node are as follows: * ON * OFF * RESETTING - when a status of resetting is entered, the node will automatically exit this state after a number of steps (as defined by the nodeResetDuration configuration item) after which it returns to an ON state + * BOOTING + * SHUTTING_DOWN * Active Nodes and Service Nodes: @@ -254,6 +256,8 @@ For the nodes, the following values are represented: * 1 = ON * 2 = OFF * 3 = RESETTING + * 4 = SHUTTING_DOWN + * 5 = BOOTING * SoftwareState: @@ -303,7 +307,7 @@ Each ``node_info`` contains the following: .. code-block:: [ - hardware_state (0=none, 1=ON, 2=OFF, 3=RESETTING) + hardware_state (0=none, 1=ON, 2=OFF, 3=RESETTING, 4=SHUTTING_DOWN, 5=BOOTING) software_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) file_system_state (0=none, 1=GOOD, 2=CORRUPT, 3=DESTROYED, 4=REPAIRING, 5=RESTORING) service1_state (0=none, 1=GOOD, 2=PATCHING, 3=COMPROMISED) diff --git a/src/primaite/common/config_values_main.py b/src/primaite/common/config_values_main.py index f822b77f..a852e5e2 100644 --- a/src/primaite/common/config_values_main.py +++ b/src/primaite/common/config_values_main.py @@ -85,7 +85,10 @@ class ConfigValuesMain(object): # Patching / Reset self.os_patching_duration = 0 # The time taken to patch the OS self.node_reset_duration = 0 # The time taken to reset a node (hardware) + self.node_booting_duration = 0 # The Time taken to turn on the node + self.node_shutdown_duration = 0 # The time taken to turn off the node self.service_patching_duration = 0 # The time taken to patch a service self.file_system_repairing_limit = 0 # The time take to repair a file self.file_system_restoring_limit = 0 # The time take to restore a file self.file_system_scanning_limit = 0 # The time taken to scan the file system + diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index c268a766..68ad80f2 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -35,6 +35,8 @@ class HardwareState(Enum): ON = 1 OFF = 2 RESETTING = 3 + SHUTTING_DOWN = 4 + BOOTING = 5 class SoftwareState(Enum): diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index d3c45c79..1aa12f84 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -1,6 +1,6 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """Main environment module containing the PRIMmary AI Training Evironment (Primaite) class.""" - +import time import copy import csv import logging diff --git a/src/primaite/nodes/active_node.py b/src/primaite/nodes/active_node.py index 54b0c642..57fa4c68 100644 --- a/src/primaite/nodes/active_node.py +++ b/src/primaite/nodes/active_node.py @@ -210,3 +210,17 @@ class ActiveNode(Node): self.file_system_state_observed = self.file_system_state_actual self.file_system_scanning = False self.file_system_scanning_count = 0 + + def update_resetting_status(self): + """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): + """Updates the booting software and file state to GOOD.""" + super().update_booting_status() + if self.booting_count <= 0: + self.file_system_state_actual = FileSystemState.GOOD + self.software_state = SoftwareState.GOOD diff --git a/src/primaite/nodes/node.py b/src/primaite/nodes/node.py index 8fd69c78..00cd01c2 100644 --- a/src/primaite/nodes/node.py +++ b/src/primaite/nodes/node.py @@ -35,6 +35,8 @@ class Node: self.hardware_state: HardwareState = hardware_state self.resetting_count: int = 0 self.config_values: TrainingConfig = config_values + self.booting_count: int = 0 + self.shutting_down_count: int = 0 def __repr__(self): """Returns the name of the node.""" @@ -42,11 +44,12 @@ class Node: def turn_on(self): """Sets the node state to ON.""" - self.hardware_state = HardwareState.ON - + self.hardware_state = HardwareState.BOOTING + self.booting_count = self.config_values.node_booting_duration def turn_off(self): """Sets the node state to OFF.""" self.hardware_state = HardwareState.OFF + self.shutting_down_count = self.config_values.node_shutdown_duration def reset(self): """Sets the node state to Resetting and starts the reset count.""" @@ -59,3 +62,17 @@ class Node: if self.resetting_count <= 0: self.resetting_count = 0 self.hardware_state = HardwareState.ON + + def update_booting_status(self): + """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): + """Updates the shutdown count""" + self.shutting_down_count -= 1 + if self.shutting_down_count <= 0: + self.shutting_down_count = 0 + self.hardware_state = HardwareState.OFF diff --git a/src/primaite/nodes/service_node.py b/src/primaite/nodes/service_node.py index 5e20f783..84a7c587 100644 --- a/src/primaite/nodes/service_node.py +++ b/src/primaite/nodes/service_node.py @@ -188,3 +188,15 @@ class ServiceNode(ActiveNode): 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): + 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): + super().update_booting_status() + if self.booting_count <= 0: + for service in self.services.values(): + service.software_state =SoftwareState.GOOD diff --git a/tests/config/legacy/legacy_training_config.yaml b/tests/config/legacy/legacy_training_config.yaml index b31a73b7..5c2025a2 100644 --- a/tests/config/legacy/legacy_training_config.yaml +++ b/tests/config/legacy/legacy_training_config.yaml @@ -83,6 +83,8 @@ greenIerBlocked: -10 # Patching / Reset durations osPatchingDuration: 5 # The time taken to patch the OS nodeResetDuration: 5 # The time taken to reset a node (hardware) +nodeBootingDuration: 3 # The Time taken to turn on the node +nodeShutdownDuration: 2 # The time taken to turn off the node servicePatchingDuration: 5 # The time taken to patch a service fileSystemRepairingLimit: 5 # The time take to repair the file system fileSystemRestoringLimit: 5 # The time take to restore the file system diff --git a/tests/test_resetting_node.py b/tests/test_resetting_node.py new file mode 100644 index 00000000..0337c35a --- /dev/null +++ b/tests/test_resetting_node.py @@ -0,0 +1,103 @@ +"""Used to test Active Node functions.""" +import pytest + +from primaite.common.enums import FileSystemState, HardwareState, SoftwareState, NodeType, Priority +from primaite.common.service import Service +from primaite.nodes.active_node import ActiveNode +from primaite.common.config_values_main import ConfigValuesMain +from primaite.nodes.service_node import ServiceNode + + +@pytest.mark.parametrize( + "starting_operating_state, expected_operating_state", + [ + (HardwareState.RESETTING, HardwareState.ON) + ], +) +def test_node_resets_correctly(starting_operating_state, expected_operating_state): + """ + Tests that a node resets correctly. + """ + active_node = ActiveNode( + node_id = "0", + name = "node", + node_type = NodeType.COMPUTER, + priority = Priority.P1, + hardware_state = starting_operating_state, + ip_address = "192.168.0.1", + software_state = SoftwareState.COMPROMISED, + file_system_state = FileSystemState.CORRUPT, + config_values = ConfigValuesMain(), + ) + + for x in range(5): + active_node.update_resetting_status() + + assert active_node.software_state == SoftwareState.GOOD + assert active_node.file_system_state_actual == FileSystemState.GOOD + assert active_node.hardware_state == expected_operating_state + +@pytest.mark.parametrize( + "operating_state, expected_operating_state", + [ + (HardwareState.BOOTING, HardwareState.ON) + ], +) +def test_node_boots_correctly(operating_state, expected_operating_state): + """ + Tests that a node boots correctly. + """ + service_node = ServiceNode( + node_id = 0, + name = "node", + node_type = "COMPUTER", + priority = "1", + hardware_state = operating_state, + ip_address = "192.168.0.1", + software_state = SoftwareState.GOOD, + file_system_state = "GOOD", + config_values = 1, + ) + service_attributes = Service( + name = "node", + port = "80", + software_state = SoftwareState.COMPROMISED + ) + service_node.add_service( + service_attributes + ) + + for x in range(5): + service_node.update_booting_status() + + assert service_attributes.software_state == SoftwareState.GOOD + assert service_node.hardware_state == expected_operating_state + +@pytest.mark.parametrize( + "operating_state, expected_operating_state", + [ + (HardwareState.SHUTTING_DOWN, HardwareState.OFF) + ], +) +def test_node_shutdown_correctly(operating_state, expected_operating_state): + """ + Tests that a node shutdown correctly. + """ + active_node = ActiveNode( + node_id = 0, + name = "node", + node_type = "COMPUTER", + priority = "1", + hardware_state = operating_state, + ip_address = "192.168.0.1", + software_state = SoftwareState.GOOD, + file_system_state = "GOOD", + config_values = 1, + ) + + for x in range(5): + active_node.update_shutdown_status() + + assert active_node.hardware_state == expected_operating_state + +