From 6245ad92986856b327b8228c9d46fc6a38f90079 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Thu, 25 May 2023 21:03:11 +0100 Subject: [PATCH] #1355 - Carried out full renaming in node.py, active_node.py, passive_node.py, and service_node.py to make params and variable names explicit. - Made the same renaming in the yaml laydown config files. - Added Type hints wherever I've been. - Added a custom NodeType in custom_typing.py to encompass the Union of ActiveNode, PassiveNode, ServiceNode. --- .flake8 | 1 + docs/source/about.rst | 26 +- docs/source/config.rst | 32 +-- src/primaite/acl/access_control_list.py | 3 +- src/primaite/common/config_values_main.py | 6 +- src/primaite/common/custom_typing.py | 8 + src/primaite/common/enums.py | 20 +- src/primaite/common/service.py | 73 +----- src/primaite/config/config_1_DDOS_BASIC.yaml | 84 +++--- src/primaite/config/config_2_DDOS_BASIC.yaml | 140 +++++----- .../config/config_3_DOS_VERY_BASIC.yaml | 56 ++-- .../config/config_5_DATA_MANIPULATION.yaml | 126 ++++----- src/primaite/config/config_UNIT_TEST.yaml | 126 ++++----- src/primaite/config/config_main.yaml | 4 +- src/primaite/environment/primaite_env.py | 214 ++++++++-------- src/primaite/environment/reward.py | 183 +++++++------- src/primaite/links/link.py | 3 +- src/primaite/main.py | 8 +- src/primaite/nodes/active_node.py | 239 ++++++++---------- src/primaite/nodes/node.py | 139 +++------- .../nodes/node_state_instruction_red.py | 7 +- src/primaite/nodes/passive_node.py | 40 ++- src/primaite/nodes/service_node.py | 238 +++++++++-------- src/primaite/pol/green_pol.py | 72 ++++-- src/primaite/pol/red_agent_pol.py | 94 ++++--- ...ne_node_states_on_off_lay_down_config.yaml | 14 +- .../one_node_states_on_off_main_config.yaml | 4 +- tests/conftest.py | 8 +- tests/test_active_node.py | 54 ++-- tests/test_reward.py | 2 +- tests/test_service_node.py | 22 +- 31 files changed, 965 insertions(+), 1081 deletions(-) create mode 100644 src/primaite/common/custom_typing.py diff --git a/.flake8 b/.flake8 index 77c6984b..398d14fb 100644 --- a/.flake8 +++ b/.flake8 @@ -8,5 +8,6 @@ extend-ignore = E203 E712 D401 + F811 exclude = docs/source/* diff --git a/docs/source/about.rst b/docs/source/about.rst index 2235f139..809b0ebe 100644 --- a/docs/source/about.rst +++ b/docs/source/about.rst @@ -11,7 +11,7 @@ PrimAITE provides the following features: * A flexible network / system laydown based on the Python networkx framework​ * Nodes and links (edges) host Python classes in order to present attributes and methods (and hence, a more representative model of a platform / system)​ * A ‘green agent’ Information Exchange Requirement (IER) function allows the representation of traffic (protocols and loading) on any / all links. Application of IERs is based on the status of node operating systems and services​ -* A ‘green agent’ node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. Operating state, Operating System state, Service state, File System state)​ +* A ‘green agent’ node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. Hardware state, Software State, Service state, File System state)​ * An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP, destination IP, protocol and port). Application of IERs adheres to any ACL restrictions​ * Presents an OpenAI Gym interface to the environment, allowing integration with any OpenAI Gym compliant defensive agents ​ * Red agent activity based on ‘red’ IERs and ‘red’ PoL @@ -31,12 +31,12 @@ 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)​ -* Operating State (ON, OFF, RESETTING - enumeration)​ +* Hardware State (ON, OFF, RESETTING - enumeration)​ Active Nodes also have the following attributes (Class: Active Node):​ * IP Address​ -* Operating System State (GOOD, PATCHING, COMPROMISED - enumeration)​ +* Software State (GOOD, PATCHING, COMPROMISED - enumeration)​ * File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration) Service Nodes also have the following attributes (Class: Service Node)​: @@ -101,7 +101,7 @@ The status changes that can be made to a node are as follows: * All Nodes: - * Operating State: + * Hardware State: * ON * OFF @@ -109,7 +109,7 @@ The status changes that can be made to a node are as follows: * Active Nodes and Service Nodes: - * Operating System State: + * Software State: * GOOD * PATCHING - when a status of patching is entered, the node will automatically exit this state after a number of steps (as defined by the osPatchingDuration configuration item) after which it returns to a GOOD state @@ -185,7 +185,7 @@ Observation Spaces The OpenAI Gym observation space provides the status of all nodes and links across the whole system:​ -* Nodes (in terms of operating state, operating system state, file system state and services state) ​ +* Nodes (in terms of hardware state, Software State, file system state and services state) ​ * Links (in terms of current loading for each service/protocol) An example observation space is provided below: @@ -196,8 +196,8 @@ An example observation space is provided below: * - - ID - - Operating State - - O/S State + - Hardware State + - SoftwareState - File System State - Service / Protocol A - Service / Protocol B @@ -249,13 +249,13 @@ The observation space is a 6 x 6 Box type (OpenAI Gym Space) in this example. Th For the nodes, the following values are represented: * ID - * Operating State: + * Hardware State: * 1 = ON * 2 = OFF * 3 = RESETTING - * O/S State: + * SoftwareState: * 1 = GOOD * 2 = PATCHING @@ -281,8 +281,8 @@ For the nodes, the following values are represented: For the links, the following statuses are represented: * ID - * Operating State = N/A - * O/S State = N/A + * Hardware State = N/A + * SoftwareState = N/A * Protocol = loading in bits/s Action Spaces @@ -300,7 +300,7 @@ The choice of action space used during a training session is determined in the c The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym multidiscrete type, as follows: * [0, num nodes] - Node ID (0 = nothing, node ID) - * [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = O/S state, 3 = service state, 4 = file system state) + * [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = SoftwareState, 3 = service state, 4 = file system state) * [0, 3] - Action on property (0 = nothing, 1 = on / scan, 2 = off / repair, 3 = reset / patch / restore) * [0, num services] - Resolves to service ID (0 = nothing, resolves to service) diff --git a/docs/source/config.rst b/docs/source/config.rst index 8a50e9a2..88399973 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -57,31 +57,31 @@ The config_main.yaml file consists of the following attributes: The score to give when the current situation (for a given component) is no different from that expected in the baseline (i.e. as though no blue or red agent actions had been undertaken) -* **Node Operating State [offShouldBeOn]** [int] +* **Node Hardware State [offShouldBeOn]** [int] The score to give when the node should be on, but is off -* **Node Operating State [offShouldBeResetting]** [int] +* **Node Hardware State [offShouldBeResetting]** [int] The score to give when the node should be resetting, but is off -* **Node Operating State [onShouldBeOff]** [int] +* **Node Hardware State [onShouldBeOff]** [int] The score to give when the node should be off, but is on -* **Node Operating State [onShouldBeResetting]** [int] +* **Node Hardware State [onShouldBeResetting]** [int] The score to give when the node should be resetting, but is on -* **Node Operating State [resettingShouldBeOn]** [int] +* **Node Hardware State [resettingShouldBeOn]** [int] The score to give when the node should be on, but is resetting -* **Node Operating State [resettingShouldBeOff]** [int] +* **Node Hardware State [resettingShouldBeOff]** [int] The score to give when the node should be off, but is resetting -* **Node Operating State [resetting]** [int] +* **Node Hardware State [resetting]** [int] The score to give when the node is resetting @@ -261,7 +261,7 @@ The config_main.yaml file consists of the following attributes: * **nodeResetDuration** [int] - The number of steps to take when resetting a node's operating state + The number of steps to take when resetting a node's hardware state * **servicePatchingDuration** [int] @@ -306,13 +306,13 @@ The config_[name].yaml file consists of the following attributes: * **id** [int]: Unique ID for this YAML item * **name** [freetext]: Human-readable name of the component - * **baseType** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services - * **nodeType** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER + * **node_class** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services + * **node_type** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER * **priority** [enum]: Provides a priority for each node. Can be one of P1, P2, P3, P4 or P5 (which P1 being the highest) - * **hardwareState** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING - * **ipAddress** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx - * **softwareState** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED - * **fileSystemState** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING + * **hardware_state** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING + * **ip_address** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx + * **software_state** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED + * **file_system_state** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING * **services**: For each service associated with the node: * **name** [freetext]: Free-text name of the service, but must match one of the services defined for the system in the services list @@ -367,7 +367,7 @@ The config_[name].yaml file consists of the following attributes: * **nodeId** [int]: The ID of the node to apply the PoL to * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list - * **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) + * **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) * **itemType: RED_POL** @@ -380,7 +380,7 @@ The config_[name].yaml file consists of the following attributes: * **initiator** [enum]: What initiates the PoL. Can be DIRECT, IER or SERVICE * **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE * **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list - * **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state) + * **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state) * **sourceNodeId** [int] The ID of the source node containing the service to check (used for SERVICE initiator) * **sourceNodeService** [freetext]: The service on the source node to check (used for SERVICE initiator). Must match a value in the services list for this node * **sourceNodeServiceState** [enum]: The state of the source node service to check (used for SERVICE initiator). Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 7b22187a..284ed764 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -1,5 +1,6 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A class that implements the access control list implementation for the network.""" +from typing import Dict from primaite.acl.acl_rule import ACLRule @@ -9,7 +10,7 @@ class AccessControlList: def __init__(self): """Init.""" - self.acl = {} # A dictionary of ACL Rules + self.acl: Dict[str, AccessControlList] = {} # A dictionary of ACL Rules def check_address_match(self, _rule, _source_ip_address, _dest_ip_address): """ diff --git a/src/primaite/common/config_values_main.py b/src/primaite/common/config_values_main.py index d6fe22a9..3493f9d2 100644 --- a/src/primaite/common/config_values_main.py +++ b/src/primaite/common/config_values_main.py @@ -2,7 +2,7 @@ """The config class.""" -class config_values_main(object): +class ConfigValuesMain(object): """Class to hold main config values.""" def __init__(self): @@ -23,7 +23,7 @@ class config_values_main(object): # Reward values # Generic self.all_ok = 0 - # Node Operating State + # Node Hardware State self.off_should_be_on = 0 self.off_should_be_resetting = 0 self.on_should_be_off = 0 @@ -31,7 +31,7 @@ class config_values_main(object): self.resetting_should_be_on = 0 self.resetting_should_be_off = 0 self.resetting = 0 - # Node O/S or Service State + # Node Software or Service State self.good_should_be_patching = 0 self.good_should_be_compromised = 0 self.good_should_be_overwhelmed = 0 diff --git a/src/primaite/common/custom_typing.py b/src/primaite/common/custom_typing.py new file mode 100644 index 00000000..03ce6a3b --- /dev/null +++ b/src/primaite/common/custom_typing.py @@ -0,0 +1,8 @@ +from typing import Type, Union + +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] +"""A Union of ActiveNode, PassiveNode, and ServiceNode.""" diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index b2854b64..0e00c9e4 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -4,7 +4,7 @@ from enum import Enum -class TYPE(Enum): +class NodeType(Enum): """Node type enumeration.""" CCTV = 1 @@ -19,7 +19,7 @@ class TYPE(Enum): SERVER = 10 -class PRIORITY(Enum): +class Priority(Enum): """Node priority enumeration.""" P1 = 1 @@ -29,7 +29,7 @@ class PRIORITY(Enum): P5 = 5 -class HARDWARE_STATE(Enum): +class HardwareState(Enum): """Node hardware state enumeration.""" ON = 1 @@ -37,8 +37,8 @@ class HARDWARE_STATE(Enum): RESETTING = 3 -class SOFTWARE_STATE(Enum): - """O/S or Service state enumeration.""" +class SoftwareState(Enum): + """Software or Service state enumeration.""" GOOD = 1 PATCHING = 2 @@ -46,7 +46,7 @@ class SOFTWARE_STATE(Enum): OVERWHELMED = 4 -class NODE_POL_TYPE(Enum): +class NodePOLType(Enum): """Node Pattern of Life type enumeration.""" OPERATING = 1 @@ -55,7 +55,7 @@ class NODE_POL_TYPE(Enum): FILE = 4 -class NODE_POL_INITIATOR(Enum): +class NodePOLInitiator(Enum): """Node Pattern of Life initiator enumeration.""" DIRECT = 1 @@ -63,7 +63,7 @@ class NODE_POL_INITIATOR(Enum): SERVICE = 3 -class PROTOCOL(Enum): +class Protocol(Enum): """Service protocol enumeration.""" LDAP = 0 @@ -76,14 +76,14 @@ class PROTOCOL(Enum): NONE = 7 -class ACTION_TYPE(Enum): +class ActionType(Enum): """Action type enumeration.""" NODE = 0 ACL = 1 -class FILE_SYSTEM_STATE(Enum): +class FileSystemState(Enum): """File System State.""" GOOD = 1 diff --git a/src/primaite/common/service.py b/src/primaite/common/service.py index 010a59e5..51403879 100644 --- a/src/primaite/common/service.py +++ b/src/primaite/common/service.py @@ -1,83 +1,28 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """The Service class.""" -from primaite.common.enums import SOFTWARE_STATE +from primaite.common.enums import SoftwareState class Service(object): """Service class.""" - def __init__(self, _name, _port, _state): + def __init__(self, name: str, port: str, software_state: SoftwareState): """ Init. - Args: - _name: The service name - _port: The service port - _state: The service state + :param name: The service name. + :param port: The service port. + :param software_state: The service SoftwareState. """ - self.name = _name - self.port = _port - self.state = _state + self.name = name + self.port = port + self.software_state = software_state self.patching_count = 0 - def set_name(self, _name): - """ - Sets the service name. - - Args: - _name: The service name - """ - self.name = _name - - def get_name(self): - """ - Gets the service name. - - Returns: - The service name - """ - return self.name - - def set_port(self, _port): - """ - Sets the service port. - - Args: - _port: The service port - """ - self.port = _port - - def get_port(self): - """ - Gets the service port. - - Returns: - The service port - """ - return self.port - - def set_state(self, _state): - """ - Sets the service state. - - Args: - _state: The service state - """ - self.state = _state - - def get_state(self): - """ - Gets the service state. - - Returns: - The service state - """ - return self.state - def reduce_patching_count(self): """Reduces the patching count for the service.""" self.patching_count -= 1 if self.patching_count <= 0: self.patching_count = 0 - self.state = SOFTWARE_STATE.GOOD + self.software_state = SoftwareState.GOOD diff --git a/src/primaite/config/config_1_DDOS_BASIC.yaml b/src/primaite/config/config_1_DDOS_BASIC.yaml index 8de32355..ada813f3 100644 --- a/src/primaite/config/config_1_DDOS_BASIC.yaml +++ b/src/primaite/config/config_1_DDOS_BASIC.yaml @@ -9,77 +9,77 @@ serviceList: - name: TCP - itemType: NODE - id: '1' + node_id: '1' name: PC1 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.2 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '2' + node_id: '2' name: SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.3 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '3' + node_id: '3' name: PC2 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.4 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '4' + node_id: '4' name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.5 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.5 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '5' + node_id: '5' name: SWITCH2 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.6 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.6 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '6' + node_id: '6' name: SWITCH3 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.7 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.7 + software_state: GOOD + file_system_state: GOOD - itemType: LINK id: '7' name: link1 diff --git a/src/primaite/config/config_2_DDOS_BASIC.yaml b/src/primaite/config/config_2_DDOS_BASIC.yaml index 08c28a5e..425fe013 100644 --- a/src/primaite/config/config_2_DDOS_BASIC.yaml +++ b/src/primaite/config/config_2_DDOS_BASIC.yaml @@ -9,133 +9,133 @@ serviceList: - name: TCP - itemType: NODE - id: '1' + node_id: '1' name: PC1 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.11 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.11 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '2' + node_id: '2' name: PC2 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '3' + node_id: '3' name: PC3 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.13 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.13 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '4' + node_id: '4' name: PC4 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.20.14 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.20.14 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '5' + node_id: '5' name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.2 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '6' + node_id: '6' name: IDS - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.4 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '7' + node_id: '7' name: SWITCH2 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.3 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '8' + node_id: '8' name: LOP1 - baseType: SERVICE - nodeType: LOP + node_class: SERVICE + node_type: LOP priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '9' + node_id: '9' name: SERVER1 - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.14 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.14 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '10' + node_id: '10' name: SERVER2 - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.20.15 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.20.15 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' diff --git a/src/primaite/config/config_3_DOS_VERY_BASIC.yaml b/src/primaite/config/config_3_DOS_VERY_BASIC.yaml index 9174bb29..8c9b84a6 100644 --- a/src/primaite/config/config_3_DOS_VERY_BASIC.yaml +++ b/src/primaite/config/config_3_DOS_VERY_BASIC.yaml @@ -9,53 +9,53 @@ serviceList: - name: TCP - itemType: NODE - id: '1' + node_id: '1' name: PC1 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.2 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.2 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '2' + node_id: '2' name: PC2 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.3 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.3 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '3' + node_id: '3' name: SWITCH1 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.1.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.1 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '4' + node_id: '4' name: SERVER1 - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.4 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.4 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' diff --git a/src/primaite/config/config_5_DATA_MANIPULATION.yaml b/src/primaite/config/config_5_DATA_MANIPULATION.yaml index 95e7b1e4..3b29ff4a 100644 --- a/src/primaite/config/config_5_DATA_MANIPULATION.yaml +++ b/src/primaite/config/config_5_DATA_MANIPULATION.yaml @@ -13,15 +13,15 @@ - name: TCP_SQL - name: UDP - itemType: NODE - id: '1' + node_id: '1' name: CLIENT_1 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.11 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.11 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -30,39 +30,39 @@ port: '53' state: GOOD - itemType: NODE - id: '2' + node_id: '2' name: CLIENT_2 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '3' + node_id: '3' name: SWITCH_1 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.10.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.1 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '4' + node_id: '4' name: SECURITY_SUITE - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.10 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.10 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -71,15 +71,15 @@ port: '53' state: GOOD - itemType: NODE - id: '5' + node_id: '5' name: MANAGEMENT_CONSOLE - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -88,25 +88,25 @@ port: '53' state: GOOD - itemType: NODE - id: '6' + node_id: '6' name: SWITCH_2 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.2.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.1 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '7' + node_id: '7' name: WEB_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.10 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.10 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -115,15 +115,15 @@ port: '1433' state: GOOD - itemType: NODE - id: '8' + node_id: '8' name: DATABASE_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.14 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.14 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -135,15 +135,15 @@ port: '53' state: GOOD - itemType: NODE - id: '9' + node_id: '9' name: BACKUP_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.16 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.16 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' diff --git a/src/primaite/config/config_UNIT_TEST.yaml b/src/primaite/config/config_UNIT_TEST.yaml index 95e7b1e4..3b29ff4a 100644 --- a/src/primaite/config/config_UNIT_TEST.yaml +++ b/src/primaite/config/config_UNIT_TEST.yaml @@ -13,15 +13,15 @@ - name: TCP_SQL - name: UDP - itemType: NODE - id: '1' + node_id: '1' name: CLIENT_1 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.11 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.11 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -30,39 +30,39 @@ port: '53' state: GOOD - itemType: NODE - id: '2' + node_id: '2' name: CLIENT_2 - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.10.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' state: GOOD - itemType: NODE - id: '3' + node_id: '3' name: SWITCH_1 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.10.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.10.1 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '4' + node_id: '4' name: SECURITY_SUITE - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.10 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.10 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -71,15 +71,15 @@ port: '53' state: GOOD - itemType: NODE - id: '5' + node_id: '5' name: MANAGEMENT_CONSOLE - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.1.12 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.1.12 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -88,25 +88,25 @@ port: '53' state: GOOD - itemType: NODE - id: '6' + node_id: '6' name: SWITCH_2 - baseType: ACTIVE - nodeType: SWITCH + node_class: ACTIVE + node_type: SWITCH priority: P2 - hardwareState: 'ON' - ipAddress: 192.168.2.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.1 + software_state: GOOD + file_system_state: GOOD - itemType: NODE - id: '7' + node_id: '7' name: WEB_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.10 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.10 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -115,15 +115,15 @@ port: '1433' state: GOOD - itemType: NODE - id: '8' + node_id: '8' name: DATABASE_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.14 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.14 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' @@ -135,15 +135,15 @@ port: '53' state: GOOD - itemType: NODE - id: '9' + node_id: '9' name: BACKUP_SERVER - baseType: SERVICE - nodeType: SERVER + node_class: SERVICE + node_type: SERVER priority: P5 - hardwareState: 'ON' - ipAddress: 192.168.2.16 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.2.16 + software_state: GOOD + file_system_state: GOOD services: - name: TCP port: '80' diff --git a/src/primaite/config/config_main.yaml b/src/primaite/config/config_main.yaml index bb957d06..b31a73b7 100644 --- a/src/primaite/config/config_main.yaml +++ b/src/primaite/config/config_main.yaml @@ -26,7 +26,7 @@ observationSpaceHighValue: 1000000000 # Reward values # Generic allOk: 0 -# Node Operating State +# Node Hardware State offShouldBeOn: -10 offShouldBeResetting: -5 onShouldBeOff: -2 @@ -34,7 +34,7 @@ onShouldBeResetting: -5 resettingShouldBeOn: -5 resettingShouldBeOff: -2 resetting: -3 -# Node O/S or Service State +# Node Software or Service State goodShouldBePatching: 2 goodShouldBeCompromised: 5 goodShouldBeOverwhelmed: 5 diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index f6c30eff..5d68009d 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -6,6 +6,7 @@ import csv import logging import os.path from datetime import datetime +from typing import Dict import networkx as nx import numpy as np @@ -14,20 +15,22 @@ 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.enums import ( - ACTION_TYPE, - FILE_SYSTEM_STATE, - HARDWARE_STATE, - NODE_POL_INITIATOR, - NODE_POL_TYPE, - PRIORITY, - SOFTWARE_STATE, - TYPE, + ActionType, + FileSystemState, + HardwareState, + NodePOLInitiator, + NodePOLType, + NodeType, + Priority, + SoftwareState, ) from primaite.common.service import Service 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_red import NodeStateInstructionRed from primaite.nodes.passive_node import PassiveNode @@ -77,19 +80,19 @@ class Primaite(Env): self.agent_identifier = self.config_values.agent_identifier # Create a dictionary to hold all the nodes - self.nodes = {} + self.nodes: Dict[str, NodeType] = {} # Create a dictionary to hold a reference set of nodes - self.nodes_reference = {} + self.nodes_reference: Dict[str, NodeType] = {} # Create a dictionary to hold all the links - self.links = {} + self.links: Dict[str, Link] = {} # Create a dictionary to hold a reference set of links - self.links_reference = {} + self.links_reference: Dict[str, Link] = {} # Create a dictionary to hold all the green IERs (this will come from an external source) - self.green_iers = {} + self.green_iers: Dict[str, IER] = {} # Create a dictionary to hold all the node PoLs (this will come from an external source) self.node_pol = {} @@ -190,8 +193,8 @@ class Primaite(Env): # For each item, we send: # - [For Nodes] | [For Links] # - node ID | link ID - # - operating state | N/A - # - operating system state | N/A + # - hardware state | N/A + # - Software State | N/A # - file system state | N/A # - service A state | service A loading # - service B state | service B loading @@ -205,7 +208,7 @@ class Primaite(Env): # observation space num_items = self.num_links + self.num_nodes # Set the number of observation parameters, being # of services plus id, - # operating state, file system state and O/S state (i.e. 4) + # hardware state, file system state and SoftwareState (i.e. 4) self.num_observation_parameters = ( self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS ) @@ -222,13 +225,13 @@ class Primaite(Env): self.env_obs = np.zeros(self.observation_shape, dtype=np.int64) # Define Action Space - depends on action space type (Node or ACL) - if self.action_type == ACTION_TYPE.NODE: + if self.action_type == ActionType.NODE: _LOGGER.info("Action space type NODE selected") # Terms (for node action space): # [0, num nodes] - node ID (0 = nothing, node ID) - # [0, 4] - what property it's acting on (0 = nothing, state, o/s state, service state, file system state) - # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) - # [0, num services] - resolves to service ID (0 = nothing, resolves to service) + # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa + # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa + # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa self.action_space = spaces.MultiDiscrete( [ self.num_nodes, @@ -402,7 +405,7 @@ class Primaite(Env): self.step_count, self.config_values, ) - print(f" Step {self.step_count} Reward: {str(reward)}") + # print(f" Step {self.step_count} Reward: {str(reward)}") self.total_reward += reward if self.step_count == self.episode_steps: self.average_reward = self.total_reward / self.step_count @@ -441,7 +444,7 @@ class Primaite(Env): """Output the link status of all links to the console.""" for link_key, link_value in self.links.items(): print("Link ID: " + link_value.get_id()) - for protocol in link_value.get_protocol_list(): + for protocol in link_value.protocol_list: print( " Protocol: " + protocol.get_name().name @@ -457,7 +460,7 @@ class Primaite(Env): _action: The action space from the agent """ # At the moment, actions are only affecting nodes - if self.action_type == ACTION_TYPE.NODE: + if self.action_type == ActionType.NODE: self.apply_actions_to_nodes(_action) else: self.apply_actions_to_acl(_action) @@ -484,32 +487,32 @@ class Primaite(Env): # This is the do nothing action return elif node_property == 1: - # This is an action on the node Operating State + # This is an action on the node Hardware State if property_action == 0: # Do nothing return elif property_action == 1: # Turn on (only applicable if it's OFF, not if it's patching) - if node.get_state() == HARDWARE_STATE.OFF: + if node.hardware_state == HardwareState.OFF: node.turn_on() elif property_action == 2: # Turn off node.turn_off() elif property_action == 3: # Reset (only applicable if it's ON) - if node.get_state() == HARDWARE_STATE.ON: + if node.hardware_state == HardwareState.ON: node.reset() else: return elif node_property == 2: if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - # This is an action on the node Operating System State + # This is an action on the node Software State if property_action == 0: # Do nothing return elif property_action == 1: # Patch (valid action if it's good or compromised) - node.set_os_state(SOFTWARE_STATE.PATCHING) + node.software_state = SoftwareState.PATCHING else: # Node is not of Active or Service Type return @@ -523,7 +526,7 @@ class Primaite(Env): elif property_action == 1: # Patch (valid action if it's good or compromised) node.set_service_state( - self.services_list[service_index], SOFTWARE_STATE.PATCHING + self.services_list[service_index], SoftwareState.PATCHING ) else: # Node is not of Service Type @@ -540,14 +543,11 @@ class Primaite(Env): elif property_action == 2: # Repair # You cannot repair a destroyed file system - it needs restoring - if ( - node.get_file_system_state_actual() - != FILE_SYSTEM_STATE.DESTROYED - ): - node.set_file_system_state(FILE_SYSTEM_STATE.REPAIRING) + if node.file_system_state_actual != FileSystemState.DESTROYED: + node.set_file_system_state(FileSystemState.REPAIRING) elif property_action == 3: # Restore - node.set_file_system_state(FILE_SYSTEM_STATE.RESTORING) + node.set_file_system_state(FileSystemState.RESTORING) else: # Node is not of Active Type return @@ -584,7 +584,7 @@ class Primaite(Env): else: node = list(self.nodes.values())[action_source_ip - 1] if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): - acl_rule_source = node.get_ip_address() + acl_rule_source = node.ip_address else: return # Destination IP value @@ -593,7 +593,7 @@ class Primaite(Env): else: node = list(self.nodes.values())[action_destination_ip - 1] if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): - acl_rule_destination = node.get_ip_address() + acl_rule_destination = node.ip_address else: return # Protocol value @@ -636,13 +636,13 @@ class Primaite(Env): e.g. reset / patching status """ for node_key, node in self.nodes.items(): - if node.get_state() == HARDWARE_STATE.RESETTING: + if node.hardware_state == HardwareState.RESETTING: node.update_resetting_status() else: pass if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): node.update_file_system_state() - if node.get_os_state() == SOFTWARE_STATE.PATCHING: + if node.software_state == SoftwareState.PATCHING: node.update_os_patching_status() else: pass @@ -654,13 +654,13 @@ class Primaite(Env): pass for node_key, node in self.nodes_reference.items(): - if node.get_state() == HARDWARE_STATE.RESETTING: + if node.hardware_state == HardwareState.RESETTING: node.update_resetting_status() else: pass if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): node.update_file_system_state() - if node.get_os_state() == SOFTWARE_STATE.PATCHING: + if node.software_state == SoftwareState.PATCHING: node.update_os_patching_status() else: pass @@ -677,13 +677,11 @@ class Primaite(Env): # Do nodes first for node_key, node in self.nodes.items(): - self.env_obs[item_index][0] = int(node.get_id()) - self.env_obs[item_index][1] = node.get_state().value + self.env_obs[item_index][0] = int(node.node_id) + self.env_obs[item_index][1] = node.hardware_state.value if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - self.env_obs[item_index][2] = node.get_os_state().value - self.env_obs[item_index][ - 3 - ] = node.get_file_system_state_observed().value + self.env_obs[item_index][2] = node.software_state.value + self.env_obs[item_index][3] = node.file_system_state_observed.value else: self.env_obs[item_index][2] = 0 self.env_obs[item_index][3] = 0 @@ -768,14 +766,14 @@ class Primaite(Env): item: A config data item """ # All nodes have these parameters - node_id = item["id"] + node_id = item["node_id"] node_name = item["name"] - node_base_type = item["baseType"] - node_type = TYPE[item["nodeType"]] - node_priority = PRIORITY[item["priority"]] - node_hardware_state = HARDWARE_STATE[item["hardwareState"]] + node_class = item["node_class"] + node_type = NodeType[item["node_type"]] + node_priority = Priority[item["priority"]] + node_hardware_state = HardwareState[item["hardware_state"]] - if node_base_type == "PASSIVE": + if node_class == "PASSIVE": node = PassiveNode( node_id, node_name, @@ -784,11 +782,11 @@ class Primaite(Env): node_hardware_state, self.config_values, ) - elif node_base_type == "ACTIVE": - # Active nodes have IP address, operating system state and file system state - node_ip_address = item["ipAddress"] - node_software_state = SOFTWARE_STATE[item["softwareState"]] - node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] + elif node_class == "ACTIVE": + # Active nodes have IP address, Software State and file system state + node_ip_address = item["ip_address"] + node_software_state = SoftwareState[item["software_state"]] + node_file_system_state = FileSystemState[item["file_system_state"]] node = ActiveNode( node_id, node_name, @@ -800,11 +798,11 @@ class Primaite(Env): node_file_system_state, self.config_values, ) - elif node_base_type == "SERVICE": - # Service nodes have IP address, operating system state, file system state and list of services - node_ip_address = item["ipAddress"] - node_software_state = SOFTWARE_STATE[item["softwareState"]] - node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] + elif node_class == "SERVICE": + # Service nodes have IP address, Software State, file system state and list of services + node_ip_address = item["ip_address"] + node_software_state = SoftwareState[item["software_state"]] + node_file_system_state = FileSystemState[item["file_system_state"]] node = ServiceNode( node_id, node_name, @@ -820,7 +818,7 @@ class Primaite(Env): for service in node_services: service_protocol = service["name"] service_port = service["port"] - service_state = SOFTWARE_STATE[service["state"]] + service_state = SoftwareState[service["state"]] node.add_service(Service(service_protocol, service_port, service_state)) else: # Bad formatting @@ -841,7 +839,7 @@ class Primaite(Env): # Add node to network (reference) self.network_reference.add_nodes_from([node_ref]) - def create_link(self, item): + def create_link(self, item: Dict): """ Creates a link from config data. @@ -854,8 +852,8 @@ class Primaite(Env): link_source = item["source"] link_destination = item["destination"] - source_node = self.nodes[link_source] - dest_node = self.nodes[link_destination] + source_node: Node = self.nodes[link_source] + dest_node: Node = self.nodes[link_destination] # Add link to network self.network.add_edge(source_node, dest_node, id=link_name) @@ -864,14 +862,14 @@ class Primaite(Env): self.links[link_name] = Link( link_id, link_bandwidth, - source_node.get_name(), - dest_node.get_name(), + source_node.name, + dest_node.name, self.services_list, ) # Reference - source_node_ref = self.nodes_reference[link_source] - dest_node_ref = self.nodes_reference[link_destination] + source_node_ref: Node = self.nodes_reference[link_source] + dest_node_ref: Node = self.nodes_reference[link_destination] # Add link to network (reference) self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name) @@ -880,8 +878,8 @@ class Primaite(Env): self.links_reference[link_name] = Link( link_id, link_bandwidth, - source_node_ref.get_name(), - dest_node_ref.get_name(), + source_node_ref.name, + dest_node_ref.name, self.services_list, ) @@ -956,18 +954,18 @@ class Primaite(Env): pol_start_step = item["startStep"] pol_end_step = item["endStep"] pol_node = item["nodeId"] - pol_type = NODE_POL_TYPE[item["type"]] + pol_type = NodePOLType[item["type"]] - # State depends on whether this is Operating, O/S, file system or Service PoL type - if pol_type == NODE_POL_TYPE.OPERATING: - pol_state = HARDWARE_STATE[item["state"]] + # State depends on whether this is Operating, Software, file system or Service PoL type + if pol_type == NodePOLType.OPERATING: + pol_state = HardwareState[item["state"]] pol_protocol = "" - elif pol_type == NODE_POL_TYPE.FILE: - pol_state = FILE_SYSTEM_STATE[item["state"]] + elif pol_type == NodePOLType.FILE: + pol_state = FileSystemState[item["state"]] pol_protocol = "" else: pol_protocol = item["protocol"] - pol_state = SOFTWARE_STATE[item["state"]] + pol_state = SoftwareState[item["state"]] self.node_pol[pol_id] = NodeStateInstructionGreen( pol_id, @@ -990,17 +988,17 @@ class Primaite(Env): pol_start_step = item["startStep"] pol_end_step = item["endStep"] pol_target_node_id = item["targetNodeId"] - pol_initiator = NODE_POL_INITIATOR[item["initiator"]] - pol_type = NODE_POL_TYPE[item["type"]] + pol_initiator = NodePOLInitiator[item["initiator"]] + pol_type = NodePOLType[item["type"]] pol_protocol = item["protocol"] - # State depends on whether this is Operating, O/S, file system or Service PoL type - if pol_type == NODE_POL_TYPE.OPERATING: - pol_state = HARDWARE_STATE[item["state"]] - elif pol_type == NODE_POL_TYPE.FILE: - pol_state = FILE_SYSTEM_STATE[item["state"]] + # State depends on whether this is Operating, Software, file system or Service PoL type + if pol_type == NodePOLType.OPERATING: + pol_state = HardwareState[item["state"]] + elif pol_type == NodePOLType.FILE: + pol_state = FileSystemState[item["state"]] else: - pol_state = SOFTWARE_STATE[item["state"]] + pol_state = SoftwareState[item["state"]] pol_source_node_id = item["sourceNodeId"] pol_source_node_service = item["sourceNodeService"] @@ -1080,7 +1078,7 @@ class Primaite(Env): Args: item: A config data item representing action info """ - self.action_type = ACTION_TYPE[action_info["type"]] + self.action_type = ActionType[action_info["type"]] def get_steps_info(self, steps_info): """ @@ -1126,38 +1124,38 @@ class Primaite(Env): item: A config data item """ # All nodes have these parameters - node_id = item["id"] - node_base_type = item["baseType"] - node_hardware_state = HARDWARE_STATE[item["hardwareState"]] + node_id = item["node_id"] + node_class = item["node_class"] + node_hardware_state: HardwareState = HardwareState[item["hardware_state"]] - node = self.nodes[node_id] + node: NodeType = self.nodes[node_id] node_ref = self.nodes_reference[node_id] # Reset the hardware state (common for all node types) - node.set_state(node_hardware_state) - node_ref.set_state(node_hardware_state) + node.hardware_state = node_hardware_state + node_ref.hardware_state = node_hardware_state - if node_base_type == "ACTIVE": - # Active nodes have operating system state - node_software_state = SOFTWARE_STATE[item["softwareState"]] - node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] - node.set_os_state(node_software_state) - node_ref.set_os_state(node_software_state) + if node_class == "ACTIVE": + # Active nodes have Software State + node_software_state = SoftwareState[item["software_state"]] + node_file_system_state = FileSystemState[item["file_system_state"]] + node.software_state = node_software_state + node_ref.software_state = node_software_state node.set_file_system_state(node_file_system_state) node_ref.set_file_system_state(node_file_system_state) - elif node_base_type == "SERVICE": - # Service nodes have operating system state and list of services - node_software_state = SOFTWARE_STATE[item["softwareState"]] - node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] - node.set_os_state(node_software_state) - node_ref.set_os_state(node_software_state) + elif node_class == "SERVICE": + # Service nodes have Software State and list of services + node_software_state = SoftwareState[item["software_state"]] + node_file_system_state = FileSystemState[item["file_system_state"]] + node.software_state = node_software_state + node_ref.software_state = node_software_state node.set_file_system_state(node_file_system_state) node_ref.set_file_system_state(node_file_system_state) # Update service states node_services = item["services"] for service in node_services: service_protocol = service["name"] - service_state = SOFTWARE_STATE[service["state"]] + service_state = SoftwareState[service["state"]] # Update node service state node.set_service_state(service_protocol, service_state) # Update reference node service state diff --git a/src/primaite/environment/reward.py b/src/primaite/environment/reward.py index 548306b0..007dcdc8 100644 --- a/src/primaite/environment/reward.py +++ b/src/primaite/environment/reward.py @@ -1,6 +1,9 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """Implements reward function.""" -from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE +from typing import Dict + +from primaite.common.enums import FileSystemState, HardwareState, SoftwareState +from primaite.common.service import Service from primaite.nodes.active_node import ActiveNode from primaite.nodes.service_node import ServiceNode @@ -28,17 +31,17 @@ def calculate_reward_function( """ reward_value = 0 - # For each node, compare operating state, o/s operating state, service states + # For each node, compare hardware state, SoftwareState, service states for node_key, final_node in final_nodes.items(): initial_node = initial_nodes[node_key] reference_node = reference_nodes[node_key] - # Operating State + # Hardware State reward_value += score_node_operating_state( final_node, initial_node, reference_node, config_values ) - # Operating System State + # Software State if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode): reward_value += score_node_os_state( final_node, initial_node, reference_node, config_values @@ -80,7 +83,7 @@ def calculate_reward_function( def score_node_operating_state(final_node, initial_node, reference_node, config_values): """ - Calculates score relating to the operating state of a node. + Calculates score relating to the hardware state of a node. Args: final_node: The node after red and blue agents take effect @@ -89,9 +92,9 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_ config_values: Config values """ score = 0 - final_node_operating_state = final_node.get_state() - initial_node_operating_state = initial_node.get_state() - reference_node_operating_state = reference_node.get_state() + final_node_operating_state = final_node.hardware_state + initial_node_operating_state = initial_node.hardware_state + reference_node_operating_state = reference_node.hardware_state if final_node_operating_state == reference_node_operating_state: # All is well - we're no different from the reference situation @@ -99,26 +102,26 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_ else: # We're different from the reference situation # Need to compare initial and reference (current) state of node (i.e. at every step) - if initial_node_operating_state == HARDWARE_STATE.ON: - if reference_node_operating_state == HARDWARE_STATE.OFF: + if initial_node_operating_state == HardwareState.ON: + if reference_node_operating_state == HardwareState.OFF: score += config_values.off_should_be_on - elif reference_node_operating_state == HARDWARE_STATE.RESETTING: + elif reference_node_operating_state == HardwareState.RESETTING: score += config_values.resetting_should_be_on else: pass - elif initial_node_operating_state == HARDWARE_STATE.OFF: - if reference_node_operating_state == HARDWARE_STATE.ON: + elif initial_node_operating_state == HardwareState.OFF: + if reference_node_operating_state == HardwareState.ON: score += config_values.on_should_be_off - elif reference_node_operating_state == HARDWARE_STATE.RESETTING: + elif reference_node_operating_state == HardwareState.RESETTING: score += config_values.resetting_should_be_off else: pass - elif initial_node_operating_state == HARDWARE_STATE.RESETTING: - if reference_node_operating_state == HARDWARE_STATE.ON: + elif initial_node_operating_state == HardwareState.RESETTING: + if reference_node_operating_state == HardwareState.ON: score += config_values.on_should_be_resetting - elif reference_node_operating_state == HARDWARE_STATE.OFF: + elif reference_node_operating_state == HardwareState.OFF: score += config_values.off_should_be_resetting - elif reference_node_operating_state == HARDWARE_STATE.RESETTING: + elif reference_node_operating_state == HardwareState.RESETTING: score += config_values.resetting else: pass @@ -130,7 +133,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_ def score_node_os_state(final_node, initial_node, reference_node, config_values): """ - Calculates score relating to the operating system state of a node. + Calculates score relating to the Software State of a node. Args: final_node: The node after red and blue agents take effect @@ -139,9 +142,9 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values) config_values: Config values """ score = 0 - final_node_os_state = final_node.get_os_state() - initial_node_os_state = initial_node.get_os_state() - reference_node_os_state = reference_node.get_os_state() + final_node_os_state = final_node.software_state + initial_node_os_state = initial_node.software_state + reference_node_os_state = reference_node.software_state if final_node_os_state == reference_node_os_state: # All is well - we're no different from the reference situation @@ -149,28 +152,28 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values) else: # We're different from the reference situation # Need to compare initial and reference (current) state of node (i.e. at every step) - if initial_node_os_state == SOFTWARE_STATE.GOOD: - if reference_node_os_state == SOFTWARE_STATE.PATCHING: + if initial_node_os_state == SoftwareState.GOOD: + if reference_node_os_state == SoftwareState.PATCHING: score += config_values.patching_should_be_good - elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED: + elif reference_node_os_state == SoftwareState.COMPROMISED: score += config_values.compromised_should_be_good else: pass - elif initial_node_os_state == SOFTWARE_STATE.PATCHING: - if reference_node_os_state == SOFTWARE_STATE.GOOD: + elif initial_node_os_state == SoftwareState.PATCHING: + if reference_node_os_state == SoftwareState.GOOD: score += config_values.good_should_be_patching - elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED: + elif reference_node_os_state == SoftwareState.COMPROMISED: score += config_values.compromised_should_be_patching - elif reference_node_os_state == SOFTWARE_STATE.PATCHING: + elif reference_node_os_state == SoftwareState.PATCHING: score += config_values.patching else: pass - elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED: - if reference_node_os_state == SOFTWARE_STATE.GOOD: + elif initial_node_os_state == SoftwareState.COMPROMISED: + if reference_node_os_state == SoftwareState.GOOD: score += config_values.good_should_be_compromised - elif reference_node_os_state == SOFTWARE_STATE.PATCHING: + elif reference_node_os_state == SoftwareState.PATCHING: score += config_values.patching_should_be_compromised - elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED: + elif reference_node_os_state == SoftwareState.COMPROMISED: score += config_values.compromised else: pass @@ -191,59 +194,59 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va config_values: Config values """ score = 0 - final_node_services = final_node.get_services() - initial_node_services = initial_node.get_services() - reference_node_services = reference_node.get_services() + final_node_services: Dict[str, Service] = final_node.services + initial_node_services: Dict[str, Service] = initial_node.services + reference_node_services: Dict[str, Service] = reference_node.services for service_key, final_service in final_node_services.items(): reference_service = reference_node_services[service_key] initial_service = initial_node_services[service_key] - if final_service.get_state() == reference_service.get_state(): + if final_service.software_state == reference_service.software_state: # All is well - we're no different from the reference situation score += config_values.all_ok else: # We're different from the reference situation # Need to compare initial and reference state of node (i.e. at every step) - if initial_service.get_state() == SOFTWARE_STATE.GOOD: - if reference_service.get_state() == SOFTWARE_STATE.PATCHING: + if initial_service.software_state == SoftwareState.GOOD: + if reference_service.software_state == SoftwareState.PATCHING: score += config_values.patching_should_be_good - elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED: + elif reference_service.software_state == SoftwareState.COMPROMISED: score += config_values.compromised_should_be_good - elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED: + elif reference_service.software_state == SoftwareState.OVERWHELMED: score += config_values.overwhelmed_should_be_good else: pass - elif initial_service.get_state() == SOFTWARE_STATE.PATCHING: - if reference_service.get_state() == SOFTWARE_STATE.GOOD: + elif initial_service.software_state == SoftwareState.PATCHING: + if reference_service.software_state == SoftwareState.GOOD: score += config_values.good_should_be_patching - elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED: + elif reference_service.software_state == SoftwareState.COMPROMISED: score += config_values.compromised_should_be_patching - elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED: + elif reference_service.software_state == SoftwareState.OVERWHELMED: score += config_values.overwhelmed_should_be_patching - elif reference_service.get_state() == SOFTWARE_STATE.PATCHING: + elif reference_service.software_state == SoftwareState.PATCHING: score += config_values.patching else: pass - elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED: - if reference_service.get_state() == SOFTWARE_STATE.GOOD: + elif initial_service.software_state == SoftwareState.COMPROMISED: + if reference_service.software_state == SoftwareState.GOOD: score += config_values.good_should_be_compromised - elif reference_service.get_state() == SOFTWARE_STATE.PATCHING: + elif reference_service.software_state == SoftwareState.PATCHING: score += config_values.patching_should_be_compromised - elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED: + elif reference_service.software_state == SoftwareState.COMPROMISED: score += config_values.compromised - elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED: + elif reference_service.software_state == SoftwareState.OVERWHELMED: score += config_values.overwhelmed_should_be_compromised else: pass - elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED: - if reference_service.get_state() == SOFTWARE_STATE.GOOD: + elif initial_service.software_state == SoftwareState.OVERWHELMED: + if reference_service.software_state == SoftwareState.GOOD: score += config_values.good_should_be_overwhelmed - elif reference_service.get_state() == SOFTWARE_STATE.PATCHING: + elif reference_service.software_state == SoftwareState.PATCHING: score += config_values.patching_should_be_overwhelmed - elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED: + elif reference_service.software_state == SoftwareState.COMPROMISED: score += config_values.compromised_should_be_overwhelmed - elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED: + elif reference_service.software_state == SoftwareState.OVERWHELMED: score += config_values.overwhelmed else: pass @@ -263,12 +266,12 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu reference_node: The node if there had been no red or blue effect """ score = 0 - final_node_file_system_state = final_node.get_file_system_state_actual() - initial_node_file_system_state = initial_node.get_file_system_state_actual() - reference_node_file_system_state = reference_node.get_file_system_state_actual() + final_node_file_system_state = final_node.file_system_state_actual + initial_node_file_system_state = initial_node.file_system_state_actual + reference_node_file_system_state = reference_node.file_system_state_actual - final_node_scanning_state = final_node.is_scanning_file_system() - reference_node_scanning_state = reference_node.is_scanning_file_system() + final_node_scanning_state = final_node.file_system_scanning + reference_node_scanning_state = reference_node.file_system_scanning # File System State if final_node_file_system_state == reference_node_file_system_state: @@ -277,66 +280,66 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu else: # We're different from the reference situation # Need to compare initial and reference state of node (i.e. at every step) - if initial_node_file_system_state == FILE_SYSTEM_STATE.GOOD: - if reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: + if initial_node_file_system_state == FileSystemState.GOOD: + if reference_node_file_system_state == FileSystemState.REPAIRING: score += config_values.repairing_should_be_good - elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: + elif reference_node_file_system_state == FileSystemState.RESTORING: score += config_values.restoring_should_be_good - elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: + elif reference_node_file_system_state == FileSystemState.CORRUPT: score += config_values.corrupt_should_be_good - elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: + elif reference_node_file_system_state == FileSystemState.DESTROYED: score += config_values.destroyed_should_be_good else: pass - elif initial_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: - if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: + elif initial_node_file_system_state == FileSystemState.REPAIRING: + if reference_node_file_system_state == FileSystemState.GOOD: score += config_values.good_should_be_repairing - elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: + elif reference_node_file_system_state == FileSystemState.RESTORING: score += config_values.restoring_should_be_repairing - elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: + elif reference_node_file_system_state == FileSystemState.CORRUPT: score += config_values.corrupt_should_be_repairing - elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: + elif reference_node_file_system_state == FileSystemState.DESTROYED: score += config_values.destroyed_should_be_repairing - elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: + elif reference_node_file_system_state == FileSystemState.REPAIRING: score += config_values.repairing else: pass - elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: - if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: + elif initial_node_file_system_state == FileSystemState.RESTORING: + if reference_node_file_system_state == FileSystemState.GOOD: score += config_values.good_should_be_restoring - elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: + elif reference_node_file_system_state == FileSystemState.REPAIRING: score += config_values.repairing_should_be_restoring - elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: + elif reference_node_file_system_state == FileSystemState.CORRUPT: score += config_values.corrupt_should_be_restoring - elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: + elif reference_node_file_system_state == FileSystemState.DESTROYED: score += config_values.destroyed_should_be_restoring - elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: + elif reference_node_file_system_state == FileSystemState.RESTORING: score += config_values.restoring else: pass - elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: - if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: + elif initial_node_file_system_state == FileSystemState.CORRUPT: + if reference_node_file_system_state == FileSystemState.GOOD: score += config_values.good_should_be_corrupt - elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: + elif reference_node_file_system_state == FileSystemState.REPAIRING: score += config_values.repairing_should_be_corrupt - elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: + elif reference_node_file_system_state == FileSystemState.RESTORING: score += config_values.restoring_should_be_corrupt - elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: + elif reference_node_file_system_state == FileSystemState.DESTROYED: score += config_values.destroyed_should_be_corrupt - elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: + elif reference_node_file_system_state == FileSystemState.CORRUPT: score += config_values.corrupt else: pass - elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: - if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: + elif initial_node_file_system_state == FileSystemState.DESTROYED: + if reference_node_file_system_state == FileSystemState.GOOD: score += config_values.good_should_be_destroyed - elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: + elif reference_node_file_system_state == FileSystemState.REPAIRING: score += config_values.repairing_should_be_destroyed - elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: + elif reference_node_file_system_state == FileSystemState.RESTORING: score += config_values.restoring_should_be_destroyed - elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: + elif reference_node_file_system_state == FileSystemState.CORRUPT: score += config_values.corrupt_should_be_destroyed - elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: + elif reference_node_file_system_state == FileSystemState.DESTROYED: score += config_values.destroyed else: pass diff --git a/src/primaite/links/link.py b/src/primaite/links/link.py index b2b4e106..90235e9f 100644 --- a/src/primaite/links/link.py +++ b/src/primaite/links/link.py @@ -1,5 +1,6 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """The link class.""" +from typing import List from primaite.common.protocol import Protocol @@ -22,7 +23,7 @@ class Link(object): self.bandwidth = _bandwidth self.source_node_name = _source_node_name self.dest_node_name = _dest_node_name - self.protocol_list = [] + self.protocol_list: List[Protocol] = [] # Add the default protocols for protocol_name in _services: diff --git a/src/primaite/main.py b/src/primaite/main.py index b130b16e..0963fa7e 100644 --- a/src/primaite/main.py +++ b/src/primaite/main.py @@ -15,7 +15,7 @@ from stable_baselines3 import A2C, PPO from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.ppo import MlpPolicy as PPOMlp -from primaite.common.config_values_main import config_values_main +from primaite.common.config_values_main import ConfigValuesMain from primaite.environment.primaite_env import Primaite from primaite.transactions.transactions_to_file import write_transaction_to_file @@ -178,7 +178,7 @@ def load_config_values(): # Reward values # Generic config_values.all_ok = int(config_data["allOk"]) - # Node Operating State + # Node Hardware State config_values.off_should_be_on = int(config_data["offShouldBeOn"]) config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"]) config_values.on_should_be_off = int(config_data["onShouldBeOff"]) @@ -186,7 +186,7 @@ def load_config_values(): config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"]) config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"]) config_values.resetting = int(config_data["resetting"]) - # Node O/S or Service State + # Node Software or Service State config_values.good_should_be_patching = int(config_data["goodShouldBePatching"]) config_values.good_should_be_compromised = int( config_data["goodShouldBeCompromised"] @@ -331,7 +331,7 @@ try: config_file_main = open("config/config_main.yaml", "r") config_data = yaml.safe_load(config_file_main) # Create a config class - config_values = config_values_main() + config_values = ConfigValuesMain() # Load in config data load_config_values() except Exception: diff --git a/src/primaite/nodes/active_node.py b/src/primaite/nodes/active_node.py index 0538f41d..83b7ab9f 100644 --- a/src/primaite/nodes/active_node.py +++ b/src/primaite/nodes/active_node.py @@ -3,7 +3,14 @@ import logging from typing import Final -from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE +from primaite.common.config_values_main import ConfigValuesMain +from primaite.common.enums import ( + FileSystemState, + HardwareState, + NodeType, + Priority, + SoftwareState, +) from primaite.nodes.node import Node _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) @@ -14,216 +21,174 @@ class ActiveNode(Node): def __init__( self, - _id, - _name, - _type, - _priority, - _state, - _ip_address, - _os_state, - _file_system_state, - _config_values, + node_id: str, + name: str, + node_type: NodeType, + priority: Priority, + hardware_state: HardwareState, + ip_address: str, + software_state: SoftwareState, + file_system_state: FileSystemState, + config_values: ConfigValuesMain, ): """ Init. - Args: - _id: The node ID - _name: The node name - _type: The node type (enum) - _priority: The node priority (enum) - _state: The node state (enum) - _ip_address: The node IP address - _os_state: The node Operating System state - _file_system_state: The node file system state - _config_values: The config values + :param node_id: The node ID + :param name: The node name + :param node_type: The node type (enum) + :param priority: The node priority (enum) + :param hardware_state: The node Hardware State + :param ip_address: The node IP address + :param software_state: The node Software State + :param file_system_state: The node file system state + :param config_values: The config values """ - super().__init__(_id, _name, _type, _priority, _state, _config_values) - self.ip_address = _ip_address - # Related to O/S - self.os_state = _os_state - self.patching_count = 0 + super().__init__( + node_id, name, node_type, priority, hardware_state, config_values + ) + self.ip_address: str = ip_address + # Related to Software + self._software_state: SoftwareState = software_state + self.patching_count: int = 0 # Related to File System - self.file_system_state_actual = _file_system_state - self.file_system_state_observed = _file_system_state - self.file_system_scanning = False - self.file_system_scanning_count = 0 - self.file_system_action_count = 0 + self.file_system_state_actual: FileSystemState = file_system_state + self.file_system_state_observed: FileSystemState = file_system_state + self.file_system_scanning: bool = False + self.file_system_scanning_count: int = 0 + self.file_system_action_count: int = 0 - def set_ip_address(self, _ip_address): + @property + def software_state(self) -> SoftwareState: """ - Sets IP address. + Get the software_state. - Args: - _ip_address: IP address + :return: The software_state. """ - self.ip_address = _ip_address + return self._software_state - def get_ip_address(self): + @software_state.setter + def software_state(self, software_state: SoftwareState): """ - Gets IP address. + Get the software_state. - Returns: - IP address + :param software_state: Software State. """ - return self.ip_address - - def set_os_state(self, _os_state): - """ - Sets operating system state. - - Args: - _os_state: Operating system state - """ - if self.operating_state != HARDWARE_STATE.OFF: - self.os_state = _os_state - if _os_state == SOFTWARE_STATE.PATCHING: + if self.hardware_state != HardwareState.OFF: + self._software_state = software_state + if software_state == SoftwareState.PATCHING: self.patching_count = self.config_values.os_patching_duration else: _LOGGER.info( - f"The Nodes operating state is OFF so OS State cannot be " + f"The Nodes hardware state is OFF so OS State cannot be " f"changed. " - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node Operating System State:{self.os_state}" + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.software_state:{self._software_state}" ) - def set_os_state_if_not_compromised(self, _os_state): + def set_software_state_if_not_compromised(self, software_state: SoftwareState): """ - Sets operating system state if the node is not compromised. + Sets Software State if the node is not compromised. Args: - _os_state: Operating system state + software_state: Software State """ - if self.operating_state != HARDWARE_STATE.OFF: - if self.os_state != SOFTWARE_STATE.COMPROMISED: - self.os_state = _os_state - if _os_state == SOFTWARE_STATE.PATCHING: + if self.hardware_state != HardwareState.OFF: + if self._software_state != SoftwareState.COMPROMISED: + self._software_state = software_state + if software_state == SoftwareState.PATCHING: self.patching_count = self.config_values.os_patching_duration else: _LOGGER.info( - f"The Nodes operating state is OFF so OS State cannot be changed." - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node Operating System State:{self.os_state}", + f"The Nodes hardware state is OFF so OS State cannot be changed." + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.software_state:{self._software_state}" ) - def get_os_state(self): - """ - Gets operating system state. - - Returns: - Operating system state - """ - return self.os_state - def update_os_patching_status(self): """Updates operating system status based on patching cycle.""" self.patching_count -= 1 if self.patching_count <= 0: self.patching_count = 0 - self.os_state = SOFTWARE_STATE.GOOD + self._software_state = SoftwareState.GOOD - def set_file_system_state(self, _file_system_state): + def set_file_system_state(self, file_system_state: FileSystemState): """ Sets the file system state (actual and observed). Args: - _file_system_state: File system state + file_system_state: File system state """ - if self.operating_state != HARDWARE_STATE.OFF: - self.file_system_state_actual = _file_system_state + if self.hardware_state != HardwareState.OFF: + self.file_system_state_actual = file_system_state - if _file_system_state == FILE_SYSTEM_STATE.REPAIRING: + if file_system_state == FileSystemState.REPAIRING: self.file_system_action_count = ( self.config_values.file_system_repairing_limit ) - self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING - elif _file_system_state == FILE_SYSTEM_STATE.RESTORING: + self.file_system_state_observed = FileSystemState.REPAIRING + elif file_system_state == FileSystemState.RESTORING: self.file_system_action_count = ( self.config_values.file_system_restoring_limit ) - self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING - elif _file_system_state == FILE_SYSTEM_STATE.GOOD: - self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD + self.file_system_state_observed = FileSystemState.RESTORING + elif file_system_state == FileSystemState.GOOD: + self.file_system_state_observed = FileSystemState.GOOD else: _LOGGER.info( - f"The Nodes operating state is OFF so File System State " + f"The Nodes hardware state is OFF so File System State " f"cannot be changed. " - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node File System State:{self.file_system_state_actual}", + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.file_system_state.actual:{self.file_system_state_actual}" ) - def set_file_system_state_if_not_compromised(self, _file_system_state): + def set_file_system_state_if_not_compromised( + self, file_system_state: FileSystemState + ): """ Sets the file system state (actual and observed) if not in a compromised state. Use for green PoL to prevent it overturning a compromised state Args: - _file_system_state: File system state + file_system_state: File system state """ - if self.operating_state != HARDWARE_STATE.OFF: + if self.hardware_state != HardwareState.OFF: if ( - self.file_system_state_actual != FILE_SYSTEM_STATE.CORRUPT - and self.file_system_state_actual != FILE_SYSTEM_STATE.DESTROYED + self.file_system_state_actual != FileSystemState.CORRUPT + and self.file_system_state_actual != FileSystemState.DESTROYED ): - self.file_system_state_actual = _file_system_state + self.file_system_state_actual = file_system_state - if _file_system_state == FILE_SYSTEM_STATE.REPAIRING: + if file_system_state == FileSystemState.REPAIRING: self.file_system_action_count = ( self.config_values.file_system_repairing_limit ) - self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING - elif _file_system_state == FILE_SYSTEM_STATE.RESTORING: + self.file_system_state_observed = FileSystemState.REPAIRING + elif file_system_state == FileSystemState.RESTORING: self.file_system_action_count = ( self.config_values.file_system_restoring_limit ) - self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING - elif _file_system_state == FILE_SYSTEM_STATE.GOOD: - self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD + self.file_system_state_observed = FileSystemState.RESTORING + elif file_system_state == FileSystemState.GOOD: + self.file_system_state_observed = FileSystemState.GOOD else: _LOGGER.info( - f"The Nodes operating state is OFF so File System State (if not " + f"The Nodes hardware state is OFF so File System State (if not " f"compromised) cannot be changed. " - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node File System State:{self.file_system_state_actual}", + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.file_system_state.actual:{self.file_system_state_actual}" ) - def get_file_system_state_actual(self): - """ - Gets file system state (actual). - - Returns: - File system state (actual) - """ - return self.file_system_state_actual - - def get_file_system_state_observed(self): - """ - Gets file system state (observed). - - Returns: - File system state (observed) - """ - return self.file_system_state_observed - def start_file_system_scan(self): """Starts a file system scan.""" self.file_system_scanning = True self.file_system_scanning_count = self.config_values.file_system_scanning_limit - def is_scanning_file_system(self): - """ - Gets true/false on whether file system is being scanned. - - Returns: - True if file system is being scanned - """ - return self.file_system_scanning - def update_file_system_state(self): """Updates file system status based on scanning/restore/repair cycle.""" # Deprecate both the action count (for restoring or reparing) and the scanning count @@ -234,11 +199,11 @@ class ActiveNode(Node): if self.file_system_action_count <= 0: self.file_system_action_count = 0 if ( - self.file_system_state_actual == FILE_SYSTEM_STATE.REPAIRING - or self.file_system_state_actual == FILE_SYSTEM_STATE.RESTORING + self.file_system_state_actual == FileSystemState.REPAIRING + or self.file_system_state_actual == FileSystemState.RESTORING ): - self.file_system_state_actual = FILE_SYSTEM_STATE.GOOD - self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD + self.file_system_state_actual = FileSystemState.GOOD + self.file_system_state_observed = FileSystemState.GOOD # Scanning updates if self.file_system_scanning == True and self.file_system_scanning_count < 0: diff --git a/src/primaite/nodes/node.py b/src/primaite/nodes/node.py index 1d9fd554..449ceb50 100644 --- a/src/primaite/nodes/node.py +++ b/src/primaite/nodes/node.py @@ -1,135 +1,56 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """The base Node class.""" -from primaite.common.enums import HARDWARE_STATE +from typing import Final + +from primaite.common.config_values_main import ConfigValuesMain +from primaite.common.enums import HardwareState, NodeType, Priority class Node: """Node class.""" - def __init__(self, _id, _name, _type, _priority, _state, _config_values): + def __init__( + self, + node_id: str, + name: str, + node_type: NodeType, + priority: Priority, + hardware_state: HardwareState, + config_values: ConfigValuesMain, + ): """ Init. - Args: - _id: The node id - _name: The name of the node - _type: The type of the node - _priority: The priority of the node - _state: The state of the node + :param node_id: The node id. + :param name: The name of the node. + :param node_type: The type of the node. + :param priority: The priority of the node. + :param hardware_state: The state of the node. + :param config_values: Config values. """ - self.id = _id - self.name = _name - self.type = _type - self.priority = _priority - self.operating_state = _state - self.resetting_count = 0 - self.config_values = _config_values + self.node_id: Final[str] = node_id + self.name: Final[str] = name + self.node_type: Final[NodeType] = node_type + self.priority = priority + self.hardware_state: HardwareState = hardware_state + self.resetting_count: int = 0 + self.config_values: ConfigValuesMain = config_values def __repr__(self): """Returns the name of the node.""" return self.name - def set_id(self, _id): - """ - Sets the node ID. - - Args: - _id: The node ID - """ - self.id = _id - - def get_id(self): - """ - Gets the node ID. - - Returns: - The node ID - """ - return self.id - - def set_name(self, _name): - """ - Sets the node name. - - Args: - _name: The node name - """ - self.name = _name - - def get_name(self): - """ - Gets the node name. - - Returns: - The node name - """ - return self.name - - def set_type(self, _type): - """ - Sets the node type. - - Args: - _type: The node type - """ - self.type = _type - - def get_type(self): - """ - Gets the node type. - - Returns: - The node type - """ - return self.type - - def set_priority(self, _priority): - """ - Sets the node priority. - - Args: - _priority: The node priority - """ - self.priority = _priority - - def get_priority(self): - """ - Gets the node priority. - - Returns: - The node priority - """ - return self.priority - - def set_state(self, _state): - """ - Sets the node state. - - Args: - _state: The node state - """ - self.operating_state = _state - - def get_state(self): - """ - Gets the node operating state. - - Returns: - The node operating state - """ - return self.operating_state - def turn_on(self): """Sets the node state to ON.""" - self.operating_state = HARDWARE_STATE.ON + self.hardware_state = HardwareState.ON def turn_off(self): """Sets the node state to OFF.""" - self.operating_state = HARDWARE_STATE.OFF + self.hardware_state = HardwareState.OFF def reset(self): """Sets the node state to Resetting and starts the reset count.""" - self.operating_state = HARDWARE_STATE.RESETTING + self.hardware_state = HardwareState.RESETTING self.resetting_count = self.config_values.node_reset_duration def update_resetting_status(self): @@ -137,4 +58,4 @@ class Node: self.resetting_count -= 1 if self.resetting_count <= 0: self.resetting_count = 0 - self.operating_state = HARDWARE_STATE.ON + self.hardware_state = HardwareState.ON diff --git a/src/primaite/nodes/node_state_instruction_red.py b/src/primaite/nodes/node_state_instruction_red.py index 56ee615b..7f62fe24 100644 --- a/src/primaite/nodes/node_state_instruction_red.py +++ b/src/primaite/nodes/node_state_instruction_red.py @@ -1,5 +1,6 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """Defines node behaviour for Green PoL.""" +from primaite.common.enums import NodePOLType class NodeStateInstructionRed(object): @@ -12,7 +13,7 @@ class NodeStateInstructionRed(object): _end_step, _target_node_id, _pol_initiator, - _pol_type, + _pol_type: NodePOLType, pol_protocol, _pol_state, _pol_source_node_id, @@ -40,7 +41,7 @@ class NodeStateInstructionRed(object): self.end_step = _end_step self.target_node_id = _target_node_id self.initiator = _pol_initiator - self.pol_type = _pol_type + self.pol_type: NodePOLType = _pol_type self.service_name = pol_protocol # Not used when not a service instruction self.state = _pol_state self.source_node_id = _pol_source_node_id @@ -83,7 +84,7 @@ class NodeStateInstructionRed(object): """ return self.initiator - def get_pol_type(self): + def get_pol_type(self) -> NodePOLType: """ Gets the node pattern of life type (enum). diff --git a/src/primaite/nodes/passive_node.py b/src/primaite/nodes/passive_node.py index ef221181..c6980e4c 100644 --- a/src/primaite/nodes/passive_node.py +++ b/src/primaite/nodes/passive_node.py @@ -1,32 +1,44 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """The Passive Node class (i.e. an actuator).""" - +from primaite.common.config_values_main import ConfigValuesMain +from primaite.common.enums import HardwareState, NodeType, Priority from primaite.nodes.node import Node class PassiveNode(Node): """The Passive Node class.""" - def __init__(self, _id, _name, _type, _priority, _state, _config_values): + def __init__( + self, + node_id: str, + name: str, + node_type: NodeType, + priority: Priority, + hardware_state: HardwareState, + config_values: ConfigValuesMain, + ): """ Init. - Args: - _id: The node id - _name: The name of the node - _type: The type of the node - _priority: The priority of the node - _state: The state of the node + :param node_id: The node id. + :param name: The name of the node. + :param node_type: The type of the node. + :param priority: The priority of the node. + :param hardware_state: The state of the node. + :param config_values: Config values. """ # Pass through to Super for now - super().__init__(_id, _name, _type, _priority, _state, _config_values) + super().__init__( + node_id, name, node_type, priority, hardware_state, config_values + ) - def get_ip_address(self): + @property + def ip_address(self) -> str: """ - Gets the node IP address. + Gets the node IP address as an empty string. - Returns: - The node IP address + No concept of IP address for passive nodes for now. + + :return: The node IP address. """ - # No concept of IP address for passive nodes for now return "" diff --git a/src/primaite/nodes/service_node.py b/src/primaite/nodes/service_node.py index 8bc8cc41..161f9249 100644 --- a/src/primaite/nodes/service_node.py +++ b/src/primaite/nodes/service_node.py @@ -1,9 +1,17 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A Service Node (i.e. not an actuator).""" import logging -from typing import Final +from typing import Dict, Final -from primaite.common.enums import HARDWARE_STATE, SOFTWARE_STATE +from primaite.common.config_values_main import ConfigValuesMain +from primaite.common.enums import ( + FileSystemState, + HardwareState, + NodeType, + Priority, + SoftwareState, +) +from primaite.common.service import Service from primaite.nodes.active_node import ActiveNode _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) @@ -14,185 +22,169 @@ class ServiceNode(ActiveNode): def __init__( self, - _id, - _name, - _type, - _priority, - _state, - _ip_address, - _os_state, - _file_system_state, - _config_values, + node_id: str, + name: str, + node_type: NodeType, + priority: Priority, + hardware_state: HardwareState, + ip_address: str, + software_state: SoftwareState, + file_system_state: FileSystemState, + config_values: ConfigValuesMain, ): """ Init. - Args: - _id: The node id - _name: The name of the node - _type: The type of the node - _priority: The priority of the node - _state: The state of the node - _ipAddress: The IP address of the node - _osState: The operating system state of the node - _file_system_state: The file system state of the node + :param node_id: The node ID + :param name: The node name + :param node_type: The node type (enum) + :param priority: The node priority (enum) + :param hardware_state: The node Hardware State + :param ip_address: The node IP address + :param software_state: The node Software State + :param file_system_state: The node file system state + :param config_values: The config values """ super().__init__( - _id, - _name, - _type, - _priority, - _state, - _ip_address, - _os_state, - _file_system_state, - _config_values, + node_id, + name, + node_type, + priority, + hardware_state, + ip_address, + software_state, + file_system_state, + config_values, ) - self.services = {} + self.services: Dict[str, Service] = {} - def add_service(self, _service): + def add_service(self, service: Service): """ Adds a service to the node. - Args: - _service: The service to add + :param service: The service to add """ - self.services[_service.get_name()] = _service + self.services[service.name] = service - def get_services(self): - """ - Gets the dictionary of services on this node. - - Returns: - Dictionary of services on this node - """ - return self.services - - def has_service(self, _protocol): + def has_service(self, protocol_name: str) -> bool: """ Indicates whether a service is on a node. - Returns: - True if service (protocol) is on the node + :param protocol_name: The service (protocol)e. + :return: True if service (protocol) is on the node, otherwise False. """ for service_key, service_value in self.services.items(): - if service_key == _protocol: + if service_key == protocol_name: return True - else: - pass return False - def service_running(self, _protocol): + def service_running(self, protocol_name: str) -> bool: """ Indicates whether a service is in a running state on the node. - Returns: - True if service (protocol) is in a running state on the node + :param protocol_name: The service (protocol) + :return: True if service (protocol) is in a running state on the + node, otherwise False. """ for service_key, service_value in self.services.items(): - if service_key == _protocol: - if service_value.get_state() != SOFTWARE_STATE.PATCHING: + if service_key == protocol_name: + if service_value.software_state != SoftwareState.PATCHING: return True else: return False - else: - pass return False - def service_is_overwhelmed(self, _protocol): + def service_is_overwhelmed(self, protocol_name: str) -> bool: """ Indicates whether a service is in an overwhelmed state on the node. - Returns: - True if service (protocol) is in an overwhelmed state on the node + :param protocol_name: The service (protocol) + :return: True if service (protocol) is in an overwhelmed state on the + node, otherwise False. """ for service_key, service_value in self.services.items(): - if service_key == _protocol: - if service_value.get_state() == SOFTWARE_STATE.OVERWHELMED: + if service_key == protocol_name: + if service_value.software_state == SoftwareState.OVERWHELMED: return True else: return False - else: - pass return False - def set_service_state(self, _protocol, _state): + def set_service_state(self, protocol_name: str, software_state: SoftwareState): """ - Sets the state of a service (protocol) on the node. + Sets the software_state of a service (protocol) on the node. - Args: - _protocol: The service (protocol) - _state: The state value + :param protocol_name: The service (protocol). + :param software_state: The software_state. """ - if self.operating_state != HARDWARE_STATE.OFF: - for service_key, service_value in self.services.items(): - if service_key == _protocol: - # Can't set to compromised if you're in a patching state - if ( - _state == SOFTWARE_STATE.COMPROMISED - and service_value.get_state() != SOFTWARE_STATE.PATCHING - ) or _state != SOFTWARE_STATE.COMPROMISED: - service_value.set_state(_state) - else: - # Do nothing - pass - if _state == SOFTWARE_STATE.PATCHING: + if self.hardware_state != HardwareState.OFF: + service_key = protocol_name + service_value = self.services.get(service_key) + if service_value: + # Can't set to compromised if you're in a patching state + if ( + software_state == SoftwareState.COMPROMISED + and service_value.software_state != SoftwareState.PATCHING + ) or software_state != SoftwareState.COMPROMISED: + service_value.software_state = software_state + if software_state == SoftwareState.PATCHING: + service_value.patching_count = ( + self.config_values.service_patching_duration + ) + else: + _LOGGER.info( + f"The Nodes hardware state is OFF so the state of a service " + f"cannot be changed. " + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.services[]:{protocol_name}, " + f"Node.services[].software_state:{software_state}" + ) + + def set_service_state_if_not_compromised( + self, protocol_name: str, software_state: SoftwareState + ): + """ + Sets the software_state of a service (protocol) on the node. + + Done if the software_state is not "compromised". + + :param protocol_name: The service (protocol). + :param software_state: The software_state. + """ + if self.hardware_state != HardwareState.OFF: + service_key = protocol_name + service_value = self.services.get(service_key) + if service_value: + if service_value.software_state != SoftwareState.COMPROMISED: + service_value.software_state = software_state + if software_state == SoftwareState.PATCHING: service_value.patching_count = ( self.config_values.service_patching_duration ) - else: - # Do nothing - pass else: _LOGGER.info( - f"The Nodes operating state is OFF so the state of a service " + f"The Nodes hardware state is OFF so the state of a service " f"cannot be changed. " - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node Service Protocol:{_protocol}, " - f"Node Service State: {_state}" + f"Node.node_id:{self.node_id}, " + f"Node.hardware_state:{self.hardware_state}, " + f"Node.services[]:{protocol_name}, " + f"Node.services[].software_state:{software_state}" ) - def set_service_state_if_not_compromised(self, _protocol, _state): - """ - Sets the state of a service (protocol) on the node if the operating state is not "compromised". - - Args: - _protocol: The service (protocol) - _state: The state value - """ - if self.operating_state != HARDWARE_STATE.OFF: - for service_key, service_value in self.services.items(): - if service_key == _protocol: - if service_value.get_state() != SOFTWARE_STATE.COMPROMISED: - service_value.set_state(_state) - if _state == SOFTWARE_STATE.PATCHING: - service_value.patching_count = ( - self.config_values.service_patching_duration - ) - else: - _LOGGER.info( - f"The Nodes operating state is OFF so the state of a service " - f"cannot be changed. " - f"Node:{self.id}, " - f"Node Operating State:{self.operating_state}, " - f"Node Service Protocol:{_protocol}, " - f"Node Service State:{_state}" - ) - - def get_service_state(self, _protocol): + def get_service_state(self, protocol_name): """ Gets the state of a service. - Returns: - The state of the service + :return: The software_state of the service. """ - for service_key, service_value in self.services.items(): - if service_key == _protocol: - return service_value.get_state() + service_key = protocol_name + service_value = self.services.get(service_key) + if service_value: + return service_value.software_state def update_services_patching_status(self): """Updates the patching counter for any service that are patching.""" for service_key, service_value in self.services.items(): - if service_value.get_state() == SOFTWARE_STATE.PATCHING: + if service_value.software_state == SoftwareState.PATCHING: service_value.reduce_patching_count() diff --git a/src/primaite/pol/green_pol.py b/src/primaite/pol/green_pol.py index ecb7dd2e..e8dbd2c2 100644 --- a/src/primaite/pol/green_pol.py +++ b/src/primaite/pol/green_pol.py @@ -1,16 +1,30 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """Implements Pattern of Life on the network (nodes and links).""" +from typing import Dict, Union -from networkx import shortest_path +from networkx import MultiGraph, shortest_path -from primaite.common.enums import HARDWARE_STATE, NODE_POL_TYPE, SOFTWARE_STATE, TYPE +from primaite.acl.access_control_list import AccessControlList +from primaite.common.custom_typing import NodeType +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_red import NodeStateInstructionRed from primaite.nodes.service_node import ServiceNode +from primaite.pol.ier import IER _VERBOSE = False -def apply_iers(network, nodes, links, iers, acl, step): +def apply_iers( + network: MultiGraph, + nodes: Dict[str, NodeType], + links: Dict[str, Link], + iers: Dict[str, IER], + acl: AccessControlList, + step: int, +): """ Applies IERs to the links (link pattern of life). @@ -51,25 +65,25 @@ def apply_iers(network, nodes, links, iers, acl, step): dest_node = nodes[dest_node_id] # 1. Check the source node situation - if source_node.get_type() == TYPE.SWITCH: + if source_node.node_type == NodeType.SWITCH: # It's a switch if ( - source_node.get_state() == HARDWARE_STATE.ON - and source_node.get_os_state() != SOFTWARE_STATE.PATCHING + source_node.hardware_state == HardwareState.ON + and source_node.software_state != SoftwareState.PATCHING ): source_valid = True else: # IER no longer valid source_valid = False - elif source_node.get_type() == TYPE.ACTUATOR: + elif source_node.node_type == NodeType.ACTUATOR: # It's an actuator # TO DO pass else: # It's not a switch or an actuator (so active node) if ( - source_node.get_state() == HARDWARE_STATE.ON - and source_node.get_os_state() != SOFTWARE_STATE.PATCHING + source_node.hardware_state == HardwareState.ON + and source_node.software_state != SoftwareState.PATCHING ): if source_node.has_service(protocol): if source_node.service_running( @@ -87,24 +101,24 @@ def apply_iers(network, nodes, links, iers, acl, step): source_valid = False # 2. Check the dest node situation - if dest_node.get_type() == TYPE.SWITCH: + if dest_node.node_type == NodeType.SWITCH: # It's a switch if ( - dest_node.get_state() == HARDWARE_STATE.ON - and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING + dest_node.hardware_state == HardwareState.ON + and dest_node.software_state != SoftwareState.PATCHING ): dest_valid = True else: # IER no longer valid dest_valid = False - elif dest_node.get_type() == TYPE.ACTUATOR: + elif dest_node.node_type == NodeType.ACTUATOR: # It's an actuator pass else: # It's not a switch or an actuator (so active node) if ( - dest_node.get_state() == HARDWARE_STATE.ON - and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING + dest_node.hardware_state == HardwareState.ON + and dest_node.software_state != SoftwareState.PATCHING ): if dest_node.has_service(protocol): if dest_node.service_running( @@ -123,15 +137,15 @@ def apply_iers(network, nodes, links, iers, acl, step): # 3. Check that the ACL doesn't block it acl_block = acl.is_blocked( - source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port + source_node.ip_address, dest_node.ip_address, protocol, port ) if acl_block: if _VERBOSE: print( "ACL block on source: " - + source_node.get_ip_address() + + source_node.ip_address + ", dest: " - + dest_node.get_ip_address() + + dest_node.ip_address + ", protocol: " + protocol + ", port: " @@ -156,8 +170,8 @@ def apply_iers(network, nodes, links, iers, acl, step): # We might have a switch in the path, so check all nodes are operational for node in path_node_list: if ( - node.get_state() != HARDWARE_STATE.ON - or node.get_os_state() == SOFTWARE_STATE.PATCHING + node.hardware_state != HardwareState.ON + or node.software_state == SoftwareState.PATCHING ): path_valid = False @@ -215,7 +229,11 @@ def apply_iers(network, nodes, links, iers, acl, step): pass -def apply_node_pol(nodes, node_pol, step): +def apply_node_pol( + nodes: Dict[str, NodeType], + node_pol: Dict[any, Union[NodeStateInstructionGreen, NodeStateInstructionRed]], + step: int, +): """ Applies node pattern of life. @@ -239,15 +257,15 @@ def apply_node_pol(nodes, node_pol, step): # continue -------------------------- node = nodes[node_id] - if node_pol_type == NODE_POL_TYPE.OPERATING: - # Change operating state - node.set_state(state) - elif node_pol_type == NODE_POL_TYPE.OS: + if node_pol_type == NodePOLType.OPERATING: + # Change hardware state + node.hardware_state = state + 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): - node.set_os_state_if_not_compromised(state) - elif node_pol_type == NODE_POL_TYPE.SERVICE: + 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): diff --git a/src/primaite/pol/red_agent_pol.py b/src/primaite/pol/red_agent_pol.py index 14d4fbd8..6a060373 100644 --- a/src/primaite/pol/red_agent_pol.py +++ b/src/primaite/pol/red_agent_pol.py @@ -1,22 +1,35 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """Implements POL on the network (nodes and links) resulting from the red agent attack.""" +from typing import Dict -from networkx import shortest_path +from networkx import MultiGraph, shortest_path +from primaite.acl.access_control_list import AccessControlList +from primaite.common.custom_typing import NodeType from primaite.common.enums import ( - HARDWARE_STATE, - NODE_POL_INITIATOR, - NODE_POL_TYPE, - SOFTWARE_STATE, - TYPE, + HardwareState, + NodePOLInitiator, + NodePOLType, + NodeType, + SoftwareState, ) +from primaite.links.link import Link from primaite.nodes.active_node import ActiveNode +from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed from primaite.nodes.service_node import ServiceNode +from primaite.pol.ier import IER _VERBOSE = False -def apply_red_agent_iers(network, nodes, links, iers, acl, step): +def apply_red_agent_iers( + network: MultiGraph, + nodes: Dict[str, NodeType], + links: Dict[str, Link], + iers: Dict[str, IER], + acl: AccessControlList, + step: int, +): """ Applies IERs to the links (link POL) resulting from red agent attack. @@ -54,25 +67,25 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step): dest_node = nodes[dest_node_id] # 1. Check the source node situation - if source_node.get_type() == TYPE.SWITCH: + if source_node.node_type == NodeType.SWITCH: # It's a switch - if source_node.get_state() == HARDWARE_STATE.ON: + if source_node.hardware_state == HardwareState.ON: source_valid = True else: # IER no longer valid source_valid = False - elif source_node.get_type() == TYPE.ACTUATOR: + elif source_node.node_type == NodeType.ACTUATOR: # It's an actuator # TO DO pass else: # It's not a switch or an actuator (so active node) - if source_node.get_state() == HARDWARE_STATE.ON: + if source_node.hardware_state == HardwareState.ON: if source_node.has_service(protocol): # Red agents IERs can only be valid if the source service is in a compromised state if ( source_node.get_service_state(protocol) - == SOFTWARE_STATE.COMPROMISED + == SoftwareState.COMPROMISED ): source_valid = True else: @@ -86,19 +99,19 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step): source_valid = False # 2. Check the dest node situation - if dest_node.get_type() == TYPE.SWITCH: + if dest_node.node_type == NodeType.SWITCH: # It's a switch - if dest_node.get_state() == HARDWARE_STATE.ON: + if dest_node.hardware_state == HardwareState.ON: dest_valid = True else: # IER no longer valid dest_valid = False - elif dest_node.get_type() == TYPE.ACTUATOR: + elif dest_node.node_type == NodeType.ACTUATOR: # It's an actuator pass else: # It's not a switch or an actuator (so active node) - if dest_node.get_state() == HARDWARE_STATE.ON: + if dest_node.hardware_state == HardwareState.ON: if dest_node.has_service(protocol): # We don't care what state the destination service is in for an IER dest_valid = True @@ -112,15 +125,15 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step): # 3. Check that the ACL doesn't block it acl_block = acl.is_blocked( - source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port + source_node.ip_address, dest_node.ip_address, protocol, port ) if acl_block: if _VERBOSE: print( "ACL block on source: " - + source_node.get_ip_address() + + source_node.ip_address + ", dest: " - + dest_node.get_ip_address() + + dest_node.ip_address + ", protocol: " + protocol + ", port: " @@ -145,7 +158,7 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step): # We might have a switch in the path, so check all nodes are operational # We're assuming here that red agents can get past switches that are patching for node in path_node_list: - if node.get_state() != HARDWARE_STATE.ON: + if node.hardware_state != HardwareState.ON: path_valid = False if path_valid: @@ -207,7 +220,12 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step): pass -def apply_red_agent_node_pol(nodes, iers, node_pol, step): +def apply_red_agent_node_pol( + nodes: NodeType, + iers: Dict[str, IER], + node_pol: Dict[str, NodeStateInstructionRed], + step: int, +): """ Applies node pattern of life. @@ -238,22 +256,22 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step): if step >= start_step and step <= stop_step: # continue -------------------------- - target_node = nodes[target_node_id] + target_node: NodeType = nodes[target_node_id] # Based the action taken on the initiator type - if initiator == NODE_POL_INITIATOR.DIRECT: + if initiator == NodePOLInitiator.DIRECT: # No conditions required, just apply the change passed_checks = True - elif initiator == NODE_POL_INITIATOR.IER: + elif initiator == NodePOLInitiator.IER: # Need to check there is a red IER incoming passed_checks = is_red_ier_incoming(target_node, iers, pol_type) - elif initiator == NODE_POL_INITIATOR.SERVICE: + elif initiator == NodePOLInitiator.SERVICE: # Need to check the condition of a service on another node source_node = nodes[source_node_id] if source_node.has_service(source_node_service_name): if ( source_node.get_service_state(source_node_service_name) - == SOFTWARE_STATE[source_node_service_state_value] + == SoftwareState[source_node_service_state_value] ): passed_checks = True else: @@ -269,16 +287,16 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step): # Only apply the PoL if the checks have passed (based on the initiator type) if passed_checks: # Apply the change - if pol_type == NODE_POL_TYPE.OPERATING: - # Change operating state - target_node.set_state(state) - elif pol_type == NODE_POL_TYPE.OS: + if pol_type == NodePOLType.OPERATING: + # Change hardware state + target_node.hardware_state = state + elif pol_type == NodePOLType.OS: # Change OS state if isinstance(target_node, ActiveNode) or isinstance( target_node, ServiceNode ): - target_node.set_os_state(state) - elif pol_type == NODE_POL_TYPE.SERVICE: + target_node.software_state = state + elif pol_type == NodePOLType.SERVICE: # Change a service state if isinstance(target_node, ServiceNode): target_node.set_service_state(service_name, state) @@ -302,18 +320,18 @@ def is_red_ier_incoming(node, iers, node_pol_type): TODO: Write more descriptive docstring with params and returns. """ - node_id = node.get_id() + 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 ( - node_pol_type == NODE_POL_TYPE.OPERATING - or node_pol_type == NODE_POL_TYPE.OS - or node_pol_type == NODE_POL_TYPE.FILE + node_pol_type == NodePOLType.OPERATING + or node_pol_type == NodePOLType.OS + or node_pol_type == NodePOLType.FILE ): - # It's looking to change operating state, file system or O/S state, so valid + # It's looking to change hardware state, file system or SoftwareState, so valid return True - elif node_pol_type == NODE_POL_TYPE.SERVICE: + elif node_pol_type == NodePOLType.SERVICE: # Check if the service is present on the node and running ier_protocol = ier_value.get_protocol() if isinstance(node, ServiceNode): diff --git a/tests/config/one_node_states_on_off_lay_down_config.yaml b/tests/config/one_node_states_on_off_lay_down_config.yaml index a543f019..355760bf 100644 --- a/tests/config/one_node_states_on_off_lay_down_config.yaml +++ b/tests/config/one_node_states_on_off_lay_down_config.yaml @@ -9,15 +9,15 @@ serviceList: - name: ftp - itemType: NODE - id: '1' + node_id: '1' name: node - baseType: SERVICE - nodeType: COMPUTER + node_class: SERVICE + node_type: COMPUTER priority: P1 - hardwareState: 'ON' - ipAddress: 192.168.0.1 - softwareState: GOOD - fileSystemState: GOOD + hardware_state: 'ON' + ip_address: 192.168.0.1 + software_state: GOOD + file_system_state: GOOD services: - name: ftp port: '21' diff --git a/tests/config/one_node_states_on_off_main_config.yaml b/tests/config/one_node_states_on_off_main_config.yaml index 6f6bb4e6..603d03dc 100644 --- a/tests/config/one_node_states_on_off_main_config.yaml +++ b/tests/config/one_node_states_on_off_main_config.yaml @@ -26,7 +26,7 @@ observationSpaceHighValue: 1000000000 # Reward values # Generic allOk: 0 -# Node Operating State +# Node Hardware State offShouldBeOn: -10 offShouldBeResetting: -5 onShouldBeOff: -2 @@ -34,7 +34,7 @@ onShouldBeResetting: -5 resettingShouldBeOn: -5 resettingShouldBeOff: -2 resetting: -3 -# Node O/S or Service State +# Node Software or Service State goodShouldBePatching: 2 goodShouldBeCompromised: 5 goodShouldBeOverwhelmed: 5 diff --git a/tests/conftest.py b/tests/conftest.py index 00f226a9..1e987223 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,7 +5,7 @@ from typing import Union import yaml -from primaite.common.config_values_main import config_values_main +from primaite.common.config_values_main import ConfigValuesMain from primaite.environment.primaite_env import Primaite ACTION_SPACE_NODE_VALUES = 1 @@ -32,7 +32,7 @@ def _get_primaite_env_from_config( # Reward values # Generic config_values.all_ok = int(config_data["allOk"]) - # Node Operating State + # Node Hardware State config_values.off_should_be_on = int(config_data["offShouldBeOn"]) config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"]) config_values.on_should_be_off = int(config_data["onShouldBeOff"]) @@ -40,7 +40,7 @@ def _get_primaite_env_from_config( config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"]) config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"]) config_values.resetting = int(config_data["resetting"]) - # Node O/S or Service State + # Node Software or Service State config_values.good_should_be_patching = int(config_data["goodShouldBePatching"]) config_values.good_should_be_compromised = int( config_data["goodShouldBeCompromised"] @@ -160,7 +160,7 @@ def _get_primaite_env_from_config( config_file_main = open(main_config_path, "r") config_data = yaml.safe_load(config_file_main) # Create a config class - config_values = config_values_main() + config_values = ConfigValuesMain() # Load in config data load_config_values() env = Primaite(config_values, []) diff --git a/tests/test_active_node.py b/tests/test_active_node.py index 760b2709..addc595c 100644 --- a/tests/test_active_node.py +++ b/tests/test_active_node.py @@ -1,22 +1,22 @@ """Used to test Active Node functions.""" import pytest -from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE +from primaite.common.enums import FileSystemState, HardwareState, SoftwareState from primaite.nodes.active_node import ActiveNode @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), - (HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), + (HardwareState.OFF, SoftwareState.GOOD), + (HardwareState.ON, SoftwareState.OVERWHELMED), ], ) def test_os_state_change(operating_state, expected_state): """ - Test that a node cannot change its operating system state. + Test that a node cannot change its Software State. - When its operating state is OFF. + When its hardware state is OFF. """ active_node = ActiveNode( 0, @@ -25,28 +25,28 @@ def test_os_state_change(operating_state, expected_state): "1", operating_state, "192.168.0.1", - SOFTWARE_STATE.GOOD, + SoftwareState.GOOD, "GOOD", 1, ) - active_node.set_os_state(SOFTWARE_STATE.OVERWHELMED) + active_node.software_state = SoftwareState.OVERWHELMED - assert active_node.get_os_state() == expected_state + assert active_node.software_state == expected_state @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), - (HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), + (HardwareState.OFF, SoftwareState.GOOD), + (HardwareState.ON, SoftwareState.OVERWHELMED), ], ) def test_os_state_change_if_not_compromised(operating_state, expected_state): """ - Test that a node cannot change its operating system state. + Test that a node cannot change its Software State. - If not compromised) when its operating state is OFF. + If not compromised) when its hardware state is OFF. """ active_node = ActiveNode( 0, @@ -55,25 +55,25 @@ def test_os_state_change_if_not_compromised(operating_state, expected_state): "1", operating_state, "192.168.0.1", - SOFTWARE_STATE.GOOD, + SoftwareState.GOOD, "GOOD", 1, ) - active_node.set_os_state_if_not_compromised(SOFTWARE_STATE.OVERWHELMED) + active_node.set_software_state_if_not_compromised(SoftwareState.OVERWHELMED) - assert active_node.get_os_state() == expected_state + assert active_node.software_state == expected_state @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD), - (HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT), + (HardwareState.OFF, FileSystemState.GOOD), + (HardwareState.ON, FileSystemState.CORRUPT), ], ) def test_file_system_change(operating_state, expected_state): - """Test that a node cannot change its file system state when its operating state is ON.""" + """Test that a node cannot change its file system state when its hardware state is ON.""" active_node = ActiveNode( 0, "node", @@ -82,27 +82,27 @@ def test_file_system_change(operating_state, expected_state): operating_state, "192.168.0.1", "COMPROMISED", - FILE_SYSTEM_STATE.GOOD, + FileSystemState.GOOD, 1, ) - active_node.set_file_system_state(FILE_SYSTEM_STATE.CORRUPT) + active_node.set_file_system_state(FileSystemState.CORRUPT) - assert active_node.get_file_system_state_actual() == expected_state + assert active_node.file_system_state_actual == expected_state @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD), - (HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT), + (HardwareState.OFF, FileSystemState.GOOD), + (HardwareState.ON, FileSystemState.CORRUPT), ], ) def test_file_system_change_if_not_compromised(operating_state, expected_state): """ Test that a node cannot change its file system state. - If not compromised) when its operating state is OFF. + If not compromised) when its hardware state is OFF. """ active_node = ActiveNode( 0, @@ -112,10 +112,10 @@ def test_file_system_change_if_not_compromised(operating_state, expected_state): operating_state, "192.168.0.1", "GOOD", - FILE_SYSTEM_STATE.GOOD, + FileSystemState.GOOD, 1, ) - active_node.set_file_system_state_if_not_compromised(FILE_SYSTEM_STATE.CORRUPT) + active_node.set_file_system_state_if_not_compromised(FileSystemState.CORRUPT) - assert active_node.get_file_system_state_actual() == expected_state + assert active_node.file_system_state_actual == expected_state diff --git a/tests/test_reward.py b/tests/test_reward.py index 9387d0a4..10dfb79c 100644 --- a/tests/test_reward.py +++ b/tests/test_reward.py @@ -19,7 +19,7 @@ def test_rewards_are_being_penalised_at_each_step_function(): File System State: goodShouldBeCorrupt = 5 (between Steps 1 & 3) Hardware State: onShouldBeOff = -2 (between Steps 4 & 6) Service State: goodShouldBeCompromised = 5 (between Steps 7 & 9) - Operating System State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12) + Software State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12) Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26 Step Count: 13 diff --git a/tests/test_service_node.py b/tests/test_service_node.py index 3c517f2b..4383fc1b 100644 --- a/tests/test_service_node.py +++ b/tests/test_service_node.py @@ -1,7 +1,7 @@ """Used to test Service Node functions.""" import pytest -from primaite.common.enums import HARDWARE_STATE, SOFTWARE_STATE +from primaite.common.enums import HardwareState, SoftwareState from primaite.common.service import Service from primaite.nodes.service_node import ServiceNode @@ -9,15 +9,15 @@ from primaite.nodes.service_node import ServiceNode @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), - (HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), + (HardwareState.OFF, SoftwareState.GOOD), + (HardwareState.ON, SoftwareState.OVERWHELMED), ], ) def test_service_state_change(operating_state, expected_state): """ Test that a node cannot change the state of a running service. - When its operating state is OFF. + When its hardware state is OFF. """ service_node = ServiceNode( 0, @@ -30,10 +30,10 @@ def test_service_state_change(operating_state, expected_state): "RESTORING", 1, ) - service = Service("TCP", 80, SOFTWARE_STATE.GOOD) + service = Service("TCP", 80, SoftwareState.GOOD) service_node.add_service(service) - service_node.set_service_state("TCP", SOFTWARE_STATE.OVERWHELMED) + service_node.set_service_state("TCP", SoftwareState.OVERWHELMED) assert service_node.get_service_state("TCP") == expected_state @@ -41,15 +41,15 @@ def test_service_state_change(operating_state, expected_state): @pytest.mark.parametrize( "operating_state, expected_state", [ - (HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), - (HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), + (HardwareState.OFF, SoftwareState.GOOD), + (HardwareState.ON, SoftwareState.OVERWHELMED), ], ) def test_service_state_change_if_not_comprised(operating_state, expected_state): """ Test that a node cannot change the state of a running service. - If not compromised when its operating state is ON. + If not compromised when its hardware state is ON. """ service_node = ServiceNode( 0, @@ -62,9 +62,9 @@ def test_service_state_change_if_not_comprised(operating_state, expected_state): "RESTORING", 1, ) - service = Service("TCP", 80, SOFTWARE_STATE.GOOD) + service = Service("TCP", 80, SoftwareState.GOOD) service_node.add_service(service) - service_node.set_service_state_if_not_compromised("TCP", SOFTWARE_STATE.OVERWHELMED) + service_node.set_service_state_if_not_compromised("TCP", SoftwareState.OVERWHELMED) assert service_node.get_service_state("TCP") == expected_state