Synced with dev

src/primaite/common/config_values_main.py

@@ -1,94 +0,0 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""The config class."""
-
-
-class ConfigValuesMain(object):
-    """Class to hold main config values."""
-
-    def __init__(self):
-        """Init."""
-        # Generic
-        self.agent_identifier = ""  # the agent in use
-        self.observation_config = None  # observation space config
-        self.num_episodes = 0  # number of episodes to train over
-        self.num_steps = 0  # number of steps in an episode
-        self.time_delay = 0  # delay between steps (ms) - applies to generic agents only
-        self.config_filename_use_case = ""  # the filename for the Use Case config file
-        self.session_type = ""  # the session type to run (TRAINING or EVALUATION)
-
-        # Environment
-        self.observation_space_high_value = (
-            0  # The high value for the observation space
-        )
-
-        # Reward values
-        # Generic
-        self.all_ok = 0
-        # Node Hardware State
-        self.off_should_be_on = 0
-        self.off_should_be_resetting = 0
-        self.on_should_be_off = 0
-        self.on_should_be_resetting = 0
-        self.resetting_should_be_on = 0
-        self.resetting_should_be_off = 0
-        self.resetting = 0
-        # Node Software or Service State
-        self.good_should_be_patching = 0
-        self.good_should_be_compromised = 0
-        self.good_should_be_overwhelmed = 0
-        self.patching_should_be_good = 0
-        self.patching_should_be_compromised = 0
-        self.patching_should_be_overwhelmed = 0
-        self.patching = 0
-        self.compromised_should_be_good = 0
-        self.compromised_should_be_patching = 0
-        self.compromised_should_be_overwhelmed = 0
-        self.compromised = 0
-        self.overwhelmed_should_be_good = 0
-        self.overwhelmed_should_be_patching = 0
-        self.overwhelmed_should_be_compromised = 0
-        self.overwhelmed = 0
-        # Node File System State
-        self.good_should_be_repairing = 0
-        self.good_should_be_restoring = 0
-        self.good_should_be_corrupt = 0
-        self.good_should_be_destroyed = 0
-        self.repairing_should_be_good = 0
-        self.repairing_should_be_restoring = 0
-        self.repairing_should_be_corrupt = 0
-        self.repairing_should_be_destroyed = (
-            0  # Repairing does not fix destroyed state - you need to restore
-        )
-        self.repairing = 0
-        self.restoring_should_be_good = 0
-        self.restoring_should_be_repairing = 0
-        self.restoring_should_be_corrupt = (
-            0  # Not the optimal method (as repair will fix corruption)
-        )
-        self.restoring_should_be_destroyed = 0
-        self.restoring = 0
-        self.corrupt_should_be_good = 0
-        self.corrupt_should_be_repairing = 0
-        self.corrupt_should_be_restoring = 0
-        self.corrupt_should_be_destroyed = 0
-        self.corrupt = 0
-        self.destroyed_should_be_good = 0
-        self.destroyed_should_be_repairing = 0
-        self.destroyed_should_be_restoring = 0
-        self.destroyed_should_be_corrupt = 0
-        self.destroyed = 0
-        self.scanning = 0
-        # IER status
-        self.red_ier_running = 0
-        self.green_ier_blocked = 0
-
-        # 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
-

src/primaite/common/training_config.py

@@ -1,91 +1,95 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
+# # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""The config class."""
+# """The config class."""
-from dataclasses import dataclass
+# from dataclasses import dataclass
-
+#
-from primaite.common.enums import ActionType
+# from primaite.common.enums import ActionType
-
+#
-
+#
-@dataclass()
+# @dataclass()
-class TrainingConfig:
+# class TrainingConfig:
-    """Class to hold main config values."""
+#     """Class to hold main config values."""
-
+#
-    # Generic
+#     # Generic
-    agent_identifier: str  # The Red Agent algo/class to be used
+#     agent_identifier: str  # The Red Agent algo/class to be used
-    action_type: ActionType  # type of action to use (NODE/ACL/ANY)
+#     action_type: ActionType  # type of action to use (NODE/ACL/ANY)
-    num_episodes: int  # number of episodes to train over
+#     num_episodes: int  # number of episodes to train over
-    num_steps: int  # number of steps in an episode
+#     num_steps: int  # number of steps in an episode
-    time_delay: int  # delay between steps (ms) - applies to generic agents only
+#     time_delay: int  # delay between steps (ms) - applies to generic agents only
-    # file
+#     # file
-    session_type: str  # the session type to run (TRAINING or EVALUATION)
+#     session_type: str  # the session type to run (TRAINING or EVALUATION)
-    load_agent: str  # Determine whether to load an agent from file
+#     load_agent: str  # Determine whether to load an agent from file
-    agent_load_file: str  # File path and file name of agent if you're loading one in
+#     agent_load_file: str  # File path and file name of agent if you're loading one in
-
+#
-    # Environment
+#     # Environment
-    observation_space_high_value: int  # The high value for the observation space
+#     observation_space_high_value: int  # The high value for the observation space
-
+#
-    # Reward values
+#     # Reward values
-    # Generic
+#     # Generic
-    all_ok: int
+#     all_ok: int
-    # Node Hardware State
+#     # Node Hardware State
-    off_should_be_on: int
+#     off_should_be_on: int
-    off_should_be_resetting: int
+#     off_should_be_resetting: int
-    on_should_be_off: int
+#     on_should_be_off: int
-    on_should_be_resetting: int
+#     on_should_be_resetting: int
-    resetting_should_be_on: int
+#     resetting_should_be_on: int
-    resetting_should_be_off: int
+#     resetting_should_be_off: int
-    resetting: int
+#     resetting: int
-    # Node Software or Service State
+#     # Node Software or Service State
-    good_should_be_patching: int
+#     good_should_be_patching: int
-    good_should_be_compromised: int
+#     good_should_be_compromised: int
-    good_should_be_overwhelmed: int
+#     good_should_be_overwhelmed: int
-    patching_should_be_good: int
+#     patching_should_be_good: int
-    patching_should_be_compromised: int
+#     patching_should_be_compromised: int
-    patching_should_be_overwhelmed: int
+#     patching_should_be_overwhelmed: int
-    patching: int
+#     patching: int
-    compromised_should_be_good: int
+#     compromised_should_be_good: int
-    compromised_should_be_patching: int
+#     compromised_should_be_patching: int
-    compromised_should_be_overwhelmed: int
+#     compromised_should_be_overwhelmed: int
-    compromised: int
+#     compromised: int
-    overwhelmed_should_be_good: int
+#     overwhelmed_should_be_good: int
-    overwhelmed_should_be_patching: int
+#     overwhelmed_should_be_patching: int
-    overwhelmed_should_be_compromised: int
+#     overwhelmed_should_be_compromised: int
-    overwhelmed: int
+#     overwhelmed: int
-    # Node File System State
+#     # Node File System State
-    good_should_be_repairing: int
+#     good_should_be_repairing: int
-    good_should_be_restoring: int
+#     good_should_be_restoring: int
-    good_should_be_corrupt: int
+#     good_should_be_corrupt: int
-    good_should_be_destroyed: int
+#     good_should_be_destroyed: int
-    repairing_should_be_good: int
+#     repairing_should_be_good: int
-    repairing_should_be_restoring: int
+#     repairing_should_be_restoring: int
-    repairing_should_be_corrupt: int
+#     repairing_should_be_corrupt: int
-    repairing_should_be_destroyed: int  # Repairing does not fix destroyed state - you need to restore
+#     repairing_should_be_destroyed: int  # Repairing does not fix destroyed state - you need to restore
-
+#
-    repairing: int
+#     repairing: int
-    restoring_should_be_good: int
+#     restoring_should_be_good: int
-    restoring_should_be_repairing: int
+#     restoring_should_be_repairing: int
-    restoring_should_be_corrupt: int  # Not the optimal method (as repair will fix corruption)
+#     restoring_should_be_corrupt: int  # Not the optimal method (as repair will fix corruption)
-
+#
-    restoring_should_be_destroyed: int
+#     restoring_should_be_destroyed: int
-    restoring: int
+#     restoring: int
-    corrupt_should_be_good: int
+#     corrupt_should_be_good: int
-    corrupt_should_be_repairing: int
+#     corrupt_should_be_repairing: int
-    corrupt_should_be_restoring: int
+#     corrupt_should_be_restoring: int
-    corrupt_should_be_destroyed: int
+#     corrupt_should_be_destroyed: int
-    corrupt: int
+#     corrupt: int
-    destroyed_should_be_good: int
+#     destroyed_should_be_good: int
-    destroyed_should_be_repairing: int
+#     destroyed_should_be_repairing: int
-    destroyed_should_be_restoring: int
+#     destroyed_should_be_restoring: int
-    destroyed_should_be_corrupt: int
+#     destroyed_should_be_corrupt: int
-    destroyed: int
+#     destroyed: int
-    scanning: int
+#     scanning: int
-    # IER status
+#     # IER status
-    red_ier_running: int
+#     red_ier_running: int
-    green_ier_blocked: int
+#     green_ier_blocked: int
-
+#
-    # Patching / Reset
+#     # Patching / Reset
-    os_patching_duration: int  # The time taken to patch the OS
+#     os_patching_duration: int  # The time taken to patch the OS
-    node_reset_duration: int  # The time taken to reset a node (hardware)
+#     node_reset_duration: int  # The time taken to reset a node (hardware)
-    service_patching_duration: int  # The time taken to patch a service
+#     node_booting_duration = 0  # The Time taken to turn on the node
-    file_system_repairing_limit: int  # The time take to repair a file
+#     node_shutdown_duration = 0  # The time taken to turn off the node
-    file_system_restoring_limit: int  # The time take to restore a file
+#     service_patching_duration: int  # The time taken to patch a service
-    file_system_scanning_limit: int  # The time taken to scan the file system
+#     file_system_repairing_limit: int  # The time take to repair a file
+#     file_system_restoring_limit: int  # The time take to restore a file
+#     file_system_scanning_limit: int  # The time taken to scan the file system
+#     # Patching / Reset
+#

src/primaite/config/training_config.py

@@ -119,6 +119,12 @@ class TrainingConfig:
     node_reset_duration: int = 5
     "The time taken to reset a node (hardware)."
 
+    node_booting_duration: int = 3
+    "The Time taken to turn on the node."
+
+    node_shutdown_duration: int = 2
+    "The time taken to turn off the node."
+
     service_patching_duration: int = 5
     "The time taken to patch a service."
 
@@ -207,7 +213,8 @@ def load(file_path: Union[str, Path],
 
 
 def convert_legacy_training_config_dict(
-        legacy_config_dict: Dict[str, Any], num_steps: int = 256,
+        legacy_config_dict: Dict[str, Any],
+        num_steps: int = 256,
         action_type: str = "ANY"
 ) -> Dict[str, Any]:
     """
@@ -220,7 +227,10 @@ def convert_legacy_training_config_dict(
         don't have action_type values.
     :return: The converted training config dict.
     """
-    config_dict = {"num_steps": num_steps, "action_type": action_type}
+    config_dict = {
+        "num_steps": num_steps,
+        "action_type": action_type
+    }
     for legacy_key, value in legacy_config_dict.items():
         new_key = _get_new_key_from_legacy(legacy_key)
         if new_key:
@@ -296,6 +306,8 @@ def _get_new_key_from_legacy(legacy_key: str) -> str:
         "greenIerBlocked": "green_ier_blocked",
         "osPatchingDuration": "os_patching_duration",
         "nodeResetDuration": "node_reset_duration",
+        "nodeBootingDuration": "node_booting_duration",
+        "nodeShutdownDuration": "node_shutdown_duration",
         "servicePatchingDuration": "service_patching_duration",
         "fileSystemRepairingLimit": "file_system_repairing_limit",
         "fileSystemRestoringLimit": "file_system_restoring_limit",

src/primaite/environment/primaite_env.py

@@ -1,6 +1,5 @@
 # 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
@@ -15,7 +14,8 @@ from gym import Env, spaces
 from matplotlib import pyplot as plt
 
 from primaite.acl.access_control_list import AccessControlList
-from primaite.agents.utils import is_valid_acl_action_extra, is_valid_node_action
+from primaite.agents.utils import is_valid_acl_action_extra, \
+    is_valid_node_action
 from primaite.common.custom_typing import NodeUnion
 from primaite.common.enums import (
     ActionType,
@@ -25,23 +25,25 @@ from primaite.common.enums import (
     NodePOLType,
     NodeType,
     Priority,
-    SoftwareState,
+    SoftwareState, ObservationType,
 )
 from primaite.common.service import Service
-from primaite.environment.observations import ObservationsHandler
 from primaite.config import training_config
 from primaite.config.training_config import TrainingConfig
+from primaite.environment.observations import ObservationsHandler
 from primaite.environment.reward import calculate_reward_function
 from primaite.links.link import Link
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.node import Node
-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.passive_node import PassiveNode
 from primaite.nodes.service_node import ServiceNode
 from primaite.pol.green_pol import apply_iers, apply_node_pol
 from primaite.pol.ier import IER
-from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
+from primaite.pol.red_agent_pol import apply_red_agent_iers, \
+    apply_red_agent_node_pol
 from primaite.transactions.transaction import Transaction
 
 _LOGGER = logging.getLogger(__name__)

tests/test_resetting_node.py

@@ -3,8 +3,8 @@ import pytest
 
 from primaite.common.enums import FileSystemState, HardwareState, SoftwareState, NodeType, Priority
 from primaite.common.service import Service
+from primaite.config.training_config import TrainingConfig
 from primaite.nodes.active_node import ActiveNode
-from primaite.common.config_values_main import ConfigValuesMain
 from primaite.nodes.service_node import ServiceNode
 
 
@@ -27,7 +27,7 @@ def test_node_resets_correctly(starting_operating_state, expected_operating_stat
         ip_address = "192.168.0.1",
         software_state = SoftwareState.COMPROMISED,
         file_system_state = FileSystemState.CORRUPT,
-        config_values = ConfigValuesMain(),
+        config_values=TrainingConfig()
     )
 
     for x in range(5):

tests/test_training_config.py

@@ -18,7 +18,8 @@ def test_legacy_lay_down_config_yaml_conversion():
 
     converted_dict = training_config.convert_legacy_training_config_dict(legacy_dict)
 
-    assert converted_dict == new_dict
+    for key, value in new_dict.items():
+        assert converted_dict[key] == value
 
 
 def test_create_config_values_main_from_file():