From a1d7f9e57c99bf3319f8410e3ce69c7124c2a4e2 Mon Sep 17 00:00:00 2001 From: Brian Kanyora Date: Fri, 2 Jun 2023 14:54:23 +0100 Subject: [PATCH] feature\898: --- src/primaite/common/config_values_main.py | 3 + src/primaite/common/enums.py | 2 + src/primaite/config/config_main.yaml | 2 + src/primaite/environment/primaite_env.py | 2 +- src/primaite/main.py | 2 + src/primaite/nodes/active_node.py | 7 ++ src/primaite/nodes/node.py | 21 +++++- src/primaite/nodes/service_node.py | 6 ++ tests/test_resetting_node.py | 91 +++++++++++++++++++++++ 9 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 tests/test_resetting_node.py diff --git a/src/primaite/common/config_values_main.py b/src/primaite/common/config_values_main.py index 3493f9d2..372ad2b2 100644 --- a/src/primaite/common/config_values_main.py +++ b/src/primaite/common/config_values_main.py @@ -84,7 +84,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 0e00c9e4..44da60e1 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/config/config_main.yaml b/src/primaite/config/config_main.yaml index b31a73b7..5c2025a2 100644 --- a/src/primaite/config/config_main.yaml +++ b/src/primaite/config/config_main.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/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 99c7c09f..808aca20 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/main.py b/src/primaite/main.py index 0963fa7e..98750cad 100644 --- a/src/primaite/main.py +++ b/src/primaite/main.py @@ -290,6 +290,8 @@ def load_config_values(): # Patching / Reset durations config_values.os_patching_duration = int(config_data["osPatchingDuration"]) config_values.node_reset_duration = int(config_data["nodeResetDuration"]) + config_values.node_booting_duration = int(config_data["nodeBootingDuration"]) + config_values.node_shutting_down_duration = int(config_data["nodeShutdownDuration"]) config_values.service_patching_duration = int( config_data["servicePatchingDuration"] ) diff --git a/src/primaite/nodes/active_node.py b/src/primaite/nodes/active_node.py index 83b7ab9f..d1c69089 100644 --- a/src/primaite/nodes/active_node.py +++ b/src/primaite/nodes/active_node.py @@ -210,3 +210,10 @@ 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): + super().update_resetting_status() + if self.resetting_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 449ceb50..384ba57f 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: ConfigValuesMain = 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 161f9249..f4a9b452 100644 --- a/src/primaite/nodes/service_node.py +++ b/src/primaite/nodes/service_node.py @@ -188,3 +188,9 @@ 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 diff --git a/tests/test_resetting_node.py b/tests/test_resetting_node.py new file mode 100644 index 00000000..089007db --- /dev/null +++ b/tests/test_resetting_node.py @@ -0,0 +1,91 @@ +"""Used to test Active Node functions.""" +import pytest + +from primaite.common.enums import FileSystemState, HardwareState, SoftwareState +from primaite.nodes.active_node import ActiveNode +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 = "COMPUTER", + priority = "1", + hardware_state = starting_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_resetting_status() + + 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. + """ + 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_booting_status() + + assert active_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 + +