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

#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.

Browse Source 
- 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.
This commit is contained in:
Chris McCarthy
2023-05-25 21:03:11 +01:00
parent 3ac2399115
commit 6245ad9298
31 changed files with 965 additions and 1081 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.flake8
View File

@@ -8,5 +8,6 @@ extend-ignore =
E203 E203
E712 E712
D401 D401
F811
exclude = exclude =
docs/source/* docs/source/*

26
docs/source/about.rst
View File

@@ -11,7 +11,7 @@ PrimAITE provides the following features:
* A flexible network / system laydown based on the Python networkx framework​ * 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)​ * 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’ 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​ * 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 ​ * 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 * 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 * Name
* Type (e.g. computer, switch, RTU - enumeration)​ * Type (e.g. computer, switch, RTU - enumeration)​
* Priority (P1, P2, P3, P4 or P5 - 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):​ Active Nodes also have the following attributes (Class: Active Node):​
* IP Address​ * IP Address​
* Operating System State (GOOD, PATCHING, COMPROMISED - enumeration)​ * Software State (GOOD, PATCHING, COMPROMISED - enumeration)​
* File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration) * File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration)
Service Nodes also have the following attributes (Class: Service Node)​: 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: * All Nodes:
* Operating State: * Hardware State:
* ON * ON
* OFF * OFF
@@ -109,7 +109,7 @@ The status changes that can be made to a node are as follows:
* Active Nodes and Service Nodes: * Active Nodes and Service Nodes:
* Operating System State: * Software State:
* GOOD * 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 * 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:​ 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) * Links (in terms of current loading for each service/protocol)
An example observation space is provided below: An example observation space is provided below:
@@ -196,8 +196,8 @@ An example observation space is provided below:
* - * -
- ID - ID
- Operating State - Hardware State
- O/S State - SoftwareState
- File System State - File System State
- Service / Protocol A - Service / Protocol A
- Service / Protocol B - 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: For the nodes, the following values are represented:
* ID * ID
* Operating State: * Hardware State:
* 1 = ON * 1 = ON
* 2 = OFF * 2 = OFF
* 3 = RESETTING * 3 = RESETTING
* O/S State: * SoftwareState:
* 1 = GOOD * 1 = GOOD
* 2 = PATCHING * 2 = PATCHING
@@ -281,8 +281,8 @@ For the nodes, the following values are represented:
For the links, the following statuses are represented: For the links, the following statuses are represented:
* ID * ID
* Operating State = N/A * Hardware State = N/A
* O/S State = N/A * SoftwareState = N/A
* Protocol = loading in bits/s * Protocol = loading in bits/s
Action Spaces 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: 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, 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, 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) * [0, num services] - Resolves to service ID (0 = nothing, resolves to service)

32
docs/source/config.rst
View File

@@ -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) 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 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 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 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 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 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 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 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] * **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] * **servicePatchingDuration** [int]
@@ -306,13 +306,13 @@ The config_[name].yaml file consists of the following attributes:
* **id** [int]: Unique ID for this YAML item * **id** [int]: Unique ID for this YAML item
* **name** [freetext]: Human-readable name of the component * **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 * **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
* **nodeType** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER * **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) * **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 * **hardware_state** [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 * **ip_address** [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 * **software_state** [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 * **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: * **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 * **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 * **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 * **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 * **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** * **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 * **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 * **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 * **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) * **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 * **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 * **sourceNodeServiceState** [enum]: The state of the source node service to check (used for SERVICE initiator). Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED

3
src/primaite/acl/access_control_list.py
View File

@@ -1,5 +1,6 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""A class that implements the access control list implementation for the network.""" """A class that implements the access control list implementation for the network."""
from typing import Dict
from primaite.acl.acl_rule import ACLRule from primaite.acl.acl_rule import ACLRule
@@ -9,7 +10,7 @@ class AccessControlList:
def __init__(self): def __init__(self):
"""Init.""" """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): def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
""" """

6
src/primaite/common/config_values_main.py
View File

@@ -2,7 +2,7 @@
"""The config class.""" """The config class."""
class config_values_main(object): class ConfigValuesMain(object):
"""Class to hold main config values.""" """Class to hold main config values."""
def __init__(self): def __init__(self):
@@ -23,7 +23,7 @@ class config_values_main(object):
# Reward values # Reward values
# Generic # Generic
self.all_ok = 0 self.all_ok = 0
# Node Operating State # Node Hardware State
self.off_should_be_on = 0 self.off_should_be_on = 0
self.off_should_be_resetting = 0 self.off_should_be_resetting = 0
self.on_should_be_off = 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_on = 0
self.resetting_should_be_off = 0 self.resetting_should_be_off = 0
self.resetting = 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_patching = 0
self.good_should_be_compromised = 0 self.good_should_be_compromised = 0
self.good_should_be_overwhelmed = 0 self.good_should_be_overwhelmed = 0

8
src/primaite/common/custom_typing.py Normal file
View File

@@ -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."""

20
src/primaite/common/enums.py
View File

@@ -4,7 +4,7 @@
from enum import Enum from enum import Enum
class TYPE(Enum): class NodeType(Enum):
"""Node type enumeration.""" """Node type enumeration."""
CCTV = 1 CCTV = 1
@@ -19,7 +19,7 @@ class TYPE(Enum):
SERVER = 10 SERVER = 10
class PRIORITY(Enum): class Priority(Enum):
"""Node priority enumeration.""" """Node priority enumeration."""
P1 = 1 P1 = 1
@@ -29,7 +29,7 @@ class PRIORITY(Enum):
P5 = 5 P5 = 5
class HARDWARE_STATE(Enum): class HardwareState(Enum):
"""Node hardware state enumeration.""" """Node hardware state enumeration."""
ON = 1 ON = 1
@@ -37,8 +37,8 @@ class HARDWARE_STATE(Enum):
RESETTING = 3 RESETTING = 3
class SOFTWARE_STATE(Enum): class SoftwareState(Enum):
"""O/S or Service state enumeration.""" """Software or Service state enumeration."""
GOOD = 1 GOOD = 1
PATCHING = 2 PATCHING = 2
@@ -46,7 +46,7 @@ class SOFTWARE_STATE(Enum):
OVERWHELMED = 4 OVERWHELMED = 4
class NODE_POL_TYPE(Enum): class NodePOLType(Enum):
"""Node Pattern of Life type enumeration.""" """Node Pattern of Life type enumeration."""
OPERATING = 1 OPERATING = 1
@@ -55,7 +55,7 @@ class NODE_POL_TYPE(Enum):
FILE = 4 FILE = 4
class NODE_POL_INITIATOR(Enum): class NodePOLInitiator(Enum):
"""Node Pattern of Life initiator enumeration.""" """Node Pattern of Life initiator enumeration."""
DIRECT = 1 DIRECT = 1
@@ -63,7 +63,7 @@ class NODE_POL_INITIATOR(Enum):
SERVICE = 3 SERVICE = 3
class PROTOCOL(Enum): class Protocol(Enum):
"""Service protocol enumeration.""" """Service protocol enumeration."""
LDAP = 0 LDAP = 0
@@ -76,14 +76,14 @@ class PROTOCOL(Enum):
NONE = 7 NONE = 7
class ACTION_TYPE(Enum): class ActionType(Enum):
"""Action type enumeration.""" """Action type enumeration."""
NODE = 0 NODE = 0
ACL = 1 ACL = 1
class FILE_SYSTEM_STATE(Enum): class FileSystemState(Enum):
"""File System State.""" """File System State."""
GOOD = 1 GOOD = 1

73
src/primaite/common/service.py
View File

@@ -1,83 +1,28 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The Service class.""" """The Service class."""
from primaite.common.enums import SOFTWARE_STATE from primaite.common.enums import SoftwareState
class Service(object): class Service(object):
"""Service class.""" """Service class."""
def __init__(self, _name, _port, _state): def __init__(self, name: str, port: str, software_state: SoftwareState):
""" """
Init. Init.
Args: :param name: The service name.
_name: The service name :param port: The service port.
_port: The service port :param software_state: The service SoftwareState.
_state: The service state
""" """
self.name = _name self.name = name
self.port = _port self.port = port
self.state = _state self.software_state = software_state
self.patching_count = 0 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): def reduce_patching_count(self):
"""Reduces the patching count for the service.""" """Reduces the patching count for the service."""
self.patching_count -= 1 self.patching_count -= 1
if self.patching_count <= 0: if self.patching_count <= 0:
self.patching_count = 0 self.patching_count = 0
self.state = SOFTWARE_STATE.GOOD self.software_state = SoftwareState.GOOD

84
src/primaite/config/config_1_DDOS_BASIC.yaml
View File

@@ -9,77 +9,77 @@
serviceList: serviceList:
- name: TCP - name: TCP
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: PC1 name: PC1
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.2 ip_address: 192.168.1.2
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '2' node_id: '2'
name: SERVER name: SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.3 ip_address: 192.168.1.3
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '3' node_id: '3'
name: PC2 name: PC2
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.4 ip_address: 192.168.1.4
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '4' node_id: '4'
name: SWITCH1 name: SWITCH1
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.5 ip_address: 192.168.1.5
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '5' node_id: '5'
name: SWITCH2 name: SWITCH2
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.6 ip_address: 192.168.1.6
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '6' node_id: '6'
name: SWITCH3 name: SWITCH3
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.7 ip_address: 192.168.1.7
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: LINK - itemType: LINK
id: '7' id: '7'
name: link1 name: link1

140
src/primaite/config/config_2_DDOS_BASIC.yaml
View File

@@ -9,133 +9,133 @@
serviceList: serviceList:
- name: TCP - name: TCP
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: PC1 name: PC1
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.11 ip_address: 192.168.10.11
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '2' node_id: '2'
name: PC2 name: PC2
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.12 ip_address: 192.168.10.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '3' node_id: '3'
name: PC3 name: PC3
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.13 ip_address: 192.168.10.13
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '4' node_id: '4'
name: PC4 name: PC4
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.20.14 ip_address: 192.168.20.14
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '5' node_id: '5'
name: SWITCH1 name: SWITCH1
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.2 ip_address: 192.168.1.2
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '6' node_id: '6'
name: IDS name: IDS
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.4 ip_address: 192.168.1.4
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '7' node_id: '7'
name: SWITCH2 name: SWITCH2
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.3 ip_address: 192.168.1.3
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '8' node_id: '8'
name: LOP1 name: LOP1
baseType: SERVICE node_class: SERVICE
nodeType: LOP node_type: LOP
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.12 ip_address: 192.168.1.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '9' node_id: '9'
name: SERVER1 name: SERVER1
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.14 ip_address: 192.168.10.14
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '10' node_id: '10'
name: SERVER2 name: SERVER2
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.20.15 ip_address: 192.168.20.15
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'

56
src/primaite/config/config_3_DOS_VERY_BASIC.yaml
View File

@@ -9,53 +9,53 @@
serviceList: serviceList:
- name: TCP - name: TCP
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: PC1 name: PC1
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.2 ip_address: 192.168.1.2
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '2' node_id: '2'
name: PC2 name: PC2
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.3 ip_address: 192.168.1.3
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '3' node_id: '3'
name: SWITCH1 name: SWITCH1
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.1 ip_address: 192.168.1.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '4' node_id: '4'
name: SERVER1 name: SERVER1
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.4 ip_address: 192.168.1.4
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'

126
src/primaite/config/config_5_DATA_MANIPULATION.yaml
View File

@@ -13,15 +13,15 @@
- name: TCP_SQL - name: TCP_SQL
- name: UDP - name: UDP
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: CLIENT_1 name: CLIENT_1
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.11 ip_address: 192.168.10.11
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -30,39 +30,39 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '2' node_id: '2'
name: CLIENT_2 name: CLIENT_2
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.12 ip_address: 192.168.10.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '3' node_id: '3'
name: SWITCH_1 name: SWITCH_1
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.1 ip_address: 192.168.10.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '4' node_id: '4'
name: SECURITY_SUITE name: SECURITY_SUITE
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.10 ip_address: 192.168.1.10
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -71,15 +71,15 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '5' node_id: '5'
name: MANAGEMENT_CONSOLE name: MANAGEMENT_CONSOLE
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.12 ip_address: 192.168.1.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -88,25 +88,25 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '6' node_id: '6'
name: SWITCH_2 name: SWITCH_2
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.1 ip_address: 192.168.2.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '7' node_id: '7'
name: WEB_SERVER name: WEB_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.10 ip_address: 192.168.2.10
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -115,15 +115,15 @@
port: '1433' port: '1433'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '8' node_id: '8'
name: DATABASE_SERVER name: DATABASE_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.14 ip_address: 192.168.2.14
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -135,15 +135,15 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '9' node_id: '9'
name: BACKUP_SERVER name: BACKUP_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.16 ip_address: 192.168.2.16
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'

126
src/primaite/config/config_UNIT_TEST.yaml
View File

@@ -13,15 +13,15 @@
- name: TCP_SQL - name: TCP_SQL
- name: UDP - name: UDP
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: CLIENT_1 name: CLIENT_1
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.11 ip_address: 192.168.10.11
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -30,39 +30,39 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '2' node_id: '2'
name: CLIENT_2 name: CLIENT_2
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.12 ip_address: 192.168.10.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '3' node_id: '3'
name: SWITCH_1 name: SWITCH_1
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.10.1 ip_address: 192.168.10.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '4' node_id: '4'
name: SECURITY_SUITE name: SECURITY_SUITE
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.10 ip_address: 192.168.1.10
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -71,15 +71,15 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '5' node_id: '5'
name: MANAGEMENT_CONSOLE name: MANAGEMENT_CONSOLE
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.1.12 ip_address: 192.168.1.12
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -88,25 +88,25 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '6' node_id: '6'
name: SWITCH_2 name: SWITCH_2
baseType: ACTIVE node_class: ACTIVE
nodeType: SWITCH node_type: SWITCH
priority: P2 priority: P2
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.1 ip_address: 192.168.2.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
- itemType: NODE - itemType: NODE
id: '7' node_id: '7'
name: WEB_SERVER name: WEB_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.10 ip_address: 192.168.2.10
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -115,15 +115,15 @@
port: '1433' port: '1433'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '8' node_id: '8'
name: DATABASE_SERVER name: DATABASE_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.14 ip_address: 192.168.2.14
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'
@@ -135,15 +135,15 @@
port: '53' port: '53'
state: GOOD state: GOOD
- itemType: NODE - itemType: NODE
id: '9' node_id: '9'
name: BACKUP_SERVER name: BACKUP_SERVER
baseType: SERVICE node_class: SERVICE
nodeType: SERVER node_type: SERVER
priority: P5 priority: P5
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.2.16 ip_address: 192.168.2.16
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: TCP - name: TCP
port: '80' port: '80'

4
src/primaite/config/config_main.yaml
View File

@@ -26,7 +26,7 @@ observationSpaceHighValue: 1000000000
# Reward values # Reward values
# Generic # Generic
allOk: 0 allOk: 0
# Node Operating State # Node Hardware State
offShouldBeOn: -10 offShouldBeOn: -10
offShouldBeResetting: -5 offShouldBeResetting: -5
onShouldBeOff: -2 onShouldBeOff: -2
@@ -34,7 +34,7 @@ onShouldBeResetting: -5
resettingShouldBeOn: -5 resettingShouldBeOn: -5
resettingShouldBeOff: -2 resettingShouldBeOff: -2
resetting: -3 resetting: -3
# Node O/S or Service State # Node Software or Service State
goodShouldBePatching: 2 goodShouldBePatching: 2
goodShouldBeCompromised: 5 goodShouldBeCompromised: 5
goodShouldBeOverwhelmed: 5 goodShouldBeOverwhelmed: 5

214
src/primaite/environment/primaite_env.py
View File

@@ -6,6 +6,7 @@ import csv
import logging import logging
import os.path import os.path
from datetime import datetime from datetime import datetime
from typing import Dict
import networkx as nx import networkx as nx
import numpy as np import numpy as np
@@ -14,20 +15,22 @@ from gym import Env, spaces
from matplotlib import pyplot as plt from matplotlib import pyplot as plt
from primaite.acl.access_control_list import AccessControlList from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeType
from primaite.common.enums import ( from primaite.common.enums import (
ACTION_TYPE, ActionType,
FILE_SYSTEM_STATE, FileSystemState,
HARDWARE_STATE, HardwareState,
NODE_POL_INITIATOR, NodePOLInitiator,
NODE_POL_TYPE, NodePOLType,
PRIORITY, NodeType,
SOFTWARE_STATE, Priority,
TYPE, SoftwareState,
) )
from primaite.common.service import Service from primaite.common.service import Service
from primaite.environment.reward import calculate_reward_function from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode from primaite.nodes.active_node import ActiveNode
from primaite.nodes.node import Node
from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from primaite.nodes.passive_node import PassiveNode from primaite.nodes.passive_node import PassiveNode
@@ -77,19 +80,19 @@ class Primaite(Env):
self.agent_identifier = self.config_values.agent_identifier self.agent_identifier = self.config_values.agent_identifier
# Create a dictionary to hold all the nodes # 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 # 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 # 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 # 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) # 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) # Create a dictionary to hold all the node PoLs (this will come from an external source)
self.node_pol = {} self.node_pol = {}
@@ -190,8 +193,8 @@ class Primaite(Env):
# For each item, we send: # For each item, we send:
# - [For Nodes] | [For Links] # - [For Nodes] | [For Links]
# - node ID | link ID # - node ID | link ID
# - operating state | N/A # - hardware state | N/A
# - operating system state | N/A # - Software State | N/A
# - file system state | N/A # - file system state | N/A
# - service A state | service A loading # - service A state | service A loading
# - service B state | service B loading # - service B state | service B loading
@@ -205,7 +208,7 @@ class Primaite(Env):
# observation space # observation space
num_items = self.num_links + self.num_nodes num_items = self.num_links + self.num_nodes
# Set the number of observation parameters, being # of services plus id, # 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_observation_parameters = (
self.num_services + self.OBSERVATION_SPACE_FIXED_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) self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
# Define Action Space - depends on action space type (Node or ACL) # 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") _LOGGER.info("Action space type NODE selected")
# Terms (for node action space): # Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID) # [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, 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) # [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) # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa
self.action_space = spaces.MultiDiscrete( self.action_space = spaces.MultiDiscrete(
[ [
self.num_nodes, self.num_nodes,
@@ -402,7 +405,7 @@ class Primaite(Env):
self.step_count, self.step_count,
self.config_values, 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 self.total_reward += reward
if self.step_count == self.episode_steps: if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count 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.""" """Output the link status of all links to the console."""
for link_key, link_value in self.links.items(): for link_key, link_value in self.links.items():
print("Link ID: " + link_value.get_id()) print("Link ID: " + link_value.get_id())
for protocol in link_value.get_protocol_list(): for protocol in link_value.protocol_list:
print( print(
" Protocol: " " Protocol: "
+ protocol.get_name().name + protocol.get_name().name
@@ -457,7 +460,7 @@ class Primaite(Env):
_action: The action space from the agent _action: The action space from the agent
""" """
# At the moment, actions are only affecting nodes # 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) self.apply_actions_to_nodes(_action)
else: else:
self.apply_actions_to_acl(_action) self.apply_actions_to_acl(_action)
@@ -484,32 +487,32 @@ class Primaite(Env):
# This is the do nothing action # This is the do nothing action
return return
elif node_property == 1: 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: if property_action == 0:
# Do nothing # Do nothing
return return
elif property_action == 1: elif property_action == 1:
# Turn on (only applicable if it's OFF, not if it's patching) # 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() node.turn_on()
elif property_action == 2: elif property_action == 2:
# Turn off # Turn off
node.turn_off() node.turn_off()
elif property_action == 3: elif property_action == 3:
# Reset (only applicable if it's ON) # Reset (only applicable if it's ON)
if node.get_state() == HARDWARE_STATE.ON: if node.hardware_state == HardwareState.ON:
node.reset() node.reset()
else: else:
return return
elif node_property == 2: elif node_property == 2:
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): 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: if property_action == 0:
# Do nothing # Do nothing
return return
elif property_action == 1: elif property_action == 1:
# Patch (valid action if it's good or compromised) # Patch (valid action if it's good or compromised)
node.set_os_state(SOFTWARE_STATE.PATCHING) node.software_state = SoftwareState.PATCHING
else: else:
# Node is not of Active or Service Type # Node is not of Active or Service Type
return return
@@ -523,7 +526,7 @@ class Primaite(Env):
elif property_action == 1: elif property_action == 1:
# Patch (valid action if it's good or compromised) # Patch (valid action if it's good or compromised)
node.set_service_state( node.set_service_state(
self.services_list[service_index], SOFTWARE_STATE.PATCHING self.services_list[service_index], SoftwareState.PATCHING
) )
else: else:
# Node is not of Service Type # Node is not of Service Type
@@ -540,14 +543,11 @@ class Primaite(Env):
elif property_action == 2: elif property_action == 2:
# Repair # Repair
# You cannot repair a destroyed file system - it needs restoring # You cannot repair a destroyed file system - it needs restoring
if ( if node.file_system_state_actual != FileSystemState.DESTROYED:
node.get_file_system_state_actual() node.set_file_system_state(FileSystemState.REPAIRING)
!= FILE_SYSTEM_STATE.DESTROYED
):
node.set_file_system_state(FILE_SYSTEM_STATE.REPAIRING)
elif property_action == 3: elif property_action == 3:
# Restore # Restore
node.set_file_system_state(FILE_SYSTEM_STATE.RESTORING) node.set_file_system_state(FileSystemState.RESTORING)
else: else:
# Node is not of Active Type # Node is not of Active Type
return return
@@ -584,7 +584,7 @@ class Primaite(Env):
else: else:
node = list(self.nodes.values())[action_source_ip - 1] node = list(self.nodes.values())[action_source_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_source = node.get_ip_address() acl_rule_source = node.ip_address
else: else:
return return
# Destination IP value # Destination IP value
@@ -593,7 +593,7 @@ class Primaite(Env):
else: else:
node = list(self.nodes.values())[action_destination_ip - 1] node = list(self.nodes.values())[action_destination_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode): if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_destination = node.get_ip_address() acl_rule_destination = node.ip_address
else: else:
return return
# Protocol value # Protocol value
@@ -636,13 +636,13 @@ class Primaite(Env):
e.g. reset / patching status e.g. reset / patching status
""" """
for node_key, node in self.nodes.items(): 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() node.update_resetting_status()
else: else:
pass pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.update_file_system_state() node.update_file_system_state()
if node.get_os_state() == SOFTWARE_STATE.PATCHING: if node.software_state == SoftwareState.PATCHING:
node.update_os_patching_status() node.update_os_patching_status()
else: else:
pass pass
@@ -654,13 +654,13 @@ class Primaite(Env):
pass pass
for node_key, node in self.nodes_reference.items(): 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() node.update_resetting_status()
else: else:
pass pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.update_file_system_state() node.update_file_system_state()
if node.get_os_state() == SOFTWARE_STATE.PATCHING: if node.software_state == SoftwareState.PATCHING:
node.update_os_patching_status() node.update_os_patching_status()
else: else:
pass pass
@@ -677,13 +677,11 @@ class Primaite(Env):
# Do nodes first # Do nodes first
for node_key, node in self.nodes.items(): for node_key, node in self.nodes.items():
self.env_obs[item_index][0] = int(node.get_id()) self.env_obs[item_index][0] = int(node.node_id)
self.env_obs[item_index][1] = node.get_state().value self.env_obs[item_index][1] = node.hardware_state.value
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
self.env_obs[item_index][2] = node.get_os_state().value self.env_obs[item_index][2] = node.software_state.value
self.env_obs[item_index][ self.env_obs[item_index][3] = node.file_system_state_observed.value
3
] = node.get_file_system_state_observed().value
else: else:
self.env_obs[item_index][2] = 0 self.env_obs[item_index][2] = 0
self.env_obs[item_index][3] = 0 self.env_obs[item_index][3] = 0
@@ -768,14 +766,14 @@ class Primaite(Env):
item: A config data item item: A config data item
""" """
# All nodes have these parameters # All nodes have these parameters
node_id = item["id"] node_id = item["node_id"]
node_name = item["name"] node_name = item["name"]
node_base_type = item["baseType"] node_class = item["node_class"]
node_type = TYPE[item["nodeType"]] node_type = NodeType[item["node_type"]]
node_priority = PRIORITY[item["priority"]] node_priority = Priority[item["priority"]]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]] node_hardware_state = HardwareState[item["hardware_state"]]
if node_base_type == "PASSIVE": if node_class == "PASSIVE":
node = PassiveNode( node = PassiveNode(
node_id, node_id,
node_name, node_name,
@@ -784,11 +782,11 @@ class Primaite(Env):
node_hardware_state, node_hardware_state,
self.config_values, self.config_values,
) )
elif node_base_type == "ACTIVE": elif node_class == "ACTIVE":
# Active nodes have IP address, operating system state and file system state # Active nodes have IP address, Software State and file system state
node_ip_address = item["ipAddress"] node_ip_address = item["ip_address"]
node_software_state = SOFTWARE_STATE[item["softwareState"]] node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] node_file_system_state = FileSystemState[item["file_system_state"]]
node = ActiveNode( node = ActiveNode(
node_id, node_id,
node_name, node_name,
@@ -800,11 +798,11 @@ class Primaite(Env):
node_file_system_state, node_file_system_state,
self.config_values, self.config_values,
) )
elif node_base_type == "SERVICE": elif node_class == "SERVICE":
# Service nodes have IP address, operating system state, file system state and list of services # Service nodes have IP address, Software State, file system state and list of services
node_ip_address = item["ipAddress"] node_ip_address = item["ip_address"]
node_software_state = SOFTWARE_STATE[item["softwareState"]] node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] node_file_system_state = FileSystemState[item["file_system_state"]]
node = ServiceNode( node = ServiceNode(
node_id, node_id,
node_name, node_name,
@@ -820,7 +818,7 @@ class Primaite(Env):
for service in node_services: for service in node_services:
service_protocol = service["name"] service_protocol = service["name"]
service_port = service["port"] 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)) node.add_service(Service(service_protocol, service_port, service_state))
else: else:
# Bad formatting # Bad formatting
@@ -841,7 +839,7 @@ class Primaite(Env):
# Add node to network (reference) # Add node to network (reference)
self.network_reference.add_nodes_from([node_ref]) 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. Creates a link from config data.
@@ -854,8 +852,8 @@ class Primaite(Env):
link_source = item["source"] link_source = item["source"]
link_destination = item["destination"] link_destination = item["destination"]
source_node = self.nodes[link_source] source_node: Node = self.nodes[link_source]
dest_node = self.nodes[link_destination] dest_node: Node = self.nodes[link_destination]
# Add link to network # Add link to network
self.network.add_edge(source_node, dest_node, id=link_name) self.network.add_edge(source_node, dest_node, id=link_name)
@@ -864,14 +862,14 @@ class Primaite(Env):
self.links[link_name] = Link( self.links[link_name] = Link(
link_id, link_id,
link_bandwidth, link_bandwidth,
source_node.get_name(), source_node.name,
dest_node.get_name(), dest_node.name,
self.services_list, self.services_list,
) )
# Reference # Reference
source_node_ref = self.nodes_reference[link_source] source_node_ref: Node = self.nodes_reference[link_source]
dest_node_ref = self.nodes_reference[link_destination] dest_node_ref: Node = self.nodes_reference[link_destination]
# Add link to network (reference) # Add link to network (reference)
self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name) 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( self.links_reference[link_name] = Link(
link_id, link_id,
link_bandwidth, link_bandwidth,
source_node_ref.get_name(), source_node_ref.name,
dest_node_ref.get_name(), dest_node_ref.name,
self.services_list, self.services_list,
) )
@@ -956,18 +954,18 @@ class Primaite(Env):
pol_start_step = item["startStep"] pol_start_step = item["startStep"]
pol_end_step = item["endStep"] pol_end_step = item["endStep"]
pol_node = item["nodeId"] 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 # State depends on whether this is Operating, Software, file system or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING: if pol_type == NodePOLType.OPERATING:
pol_state = HARDWARE_STATE[item["state"]] pol_state = HardwareState[item["state"]]
pol_protocol = "" pol_protocol = ""
elif pol_type == NODE_POL_TYPE.FILE: elif pol_type == NodePOLType.FILE:
pol_state = FILE_SYSTEM_STATE[item["state"]] pol_state = FileSystemState[item["state"]]
pol_protocol = "" pol_protocol = ""
else: else:
pol_protocol = item["protocol"] pol_protocol = item["protocol"]
pol_state = SOFTWARE_STATE[item["state"]] pol_state = SoftwareState[item["state"]]
self.node_pol[pol_id] = NodeStateInstructionGreen( self.node_pol[pol_id] = NodeStateInstructionGreen(
pol_id, pol_id,
@@ -990,17 +988,17 @@ class Primaite(Env):
pol_start_step = item["startStep"] pol_start_step = item["startStep"]
pol_end_step = item["endStep"] pol_end_step = item["endStep"]
pol_target_node_id = item["targetNodeId"] pol_target_node_id = item["targetNodeId"]
pol_initiator = NODE_POL_INITIATOR[item["initiator"]] pol_initiator = NodePOLInitiator[item["initiator"]]
pol_type = NODE_POL_TYPE[item["type"]] pol_type = NodePOLType[item["type"]]
pol_protocol = item["protocol"] pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S, file system or Service PoL type # State depends on whether this is Operating, Software, file system or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING: if pol_type == NodePOLType.OPERATING:
pol_state = HARDWARE_STATE[item["state"]] pol_state = HardwareState[item["state"]]
elif pol_type == NODE_POL_TYPE.FILE: elif pol_type == NodePOLType.FILE:
pol_state = FILE_SYSTEM_STATE[item["state"]] pol_state = FileSystemState[item["state"]]
else: else:
pol_state = SOFTWARE_STATE[item["state"]] pol_state = SoftwareState[item["state"]]
pol_source_node_id = item["sourceNodeId"] pol_source_node_id = item["sourceNodeId"]
pol_source_node_service = item["sourceNodeService"] pol_source_node_service = item["sourceNodeService"]
@@ -1080,7 +1078,7 @@ class Primaite(Env):
Args: Args:
item: A config data item representing action info 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): def get_steps_info(self, steps_info):
""" """
@@ -1126,38 +1124,38 @@ class Primaite(Env):
item: A config data item item: A config data item
""" """
# All nodes have these parameters # All nodes have these parameters
node_id = item["id"] node_id = item["node_id"]
node_base_type = item["baseType"] node_class = item["node_class"]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]] 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] node_ref = self.nodes_reference[node_id]
# Reset the hardware state (common for all node types) # Reset the hardware state (common for all node types)
node.set_state(node_hardware_state) node.hardware_state = node_hardware_state
node_ref.set_state(node_hardware_state) node_ref.hardware_state = node_hardware_state
if node_base_type == "ACTIVE": if node_class == "ACTIVE":
# Active nodes have operating system state # Active nodes have Software State
node_software_state = SOFTWARE_STATE[item["softwareState"]] node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] node_file_system_state = FileSystemState[item["file_system_state"]]
node.set_os_state(node_software_state) node.software_state = node_software_state
node_ref.set_os_state(node_software_state) node_ref.software_state = node_software_state
node.set_file_system_state(node_file_system_state) node.set_file_system_state(node_file_system_state)
node_ref.set_file_system_state(node_file_system_state) node_ref.set_file_system_state(node_file_system_state)
elif node_base_type == "SERVICE": elif node_class == "SERVICE":
# Service nodes have operating system state and list of services # Service nodes have Software State and list of services
node_software_state = SOFTWARE_STATE[item["softwareState"]] node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]] node_file_system_state = FileSystemState[item["file_system_state"]]
node.set_os_state(node_software_state) node.software_state = node_software_state
node_ref.set_os_state(node_software_state) node_ref.software_state = node_software_state
node.set_file_system_state(node_file_system_state) node.set_file_system_state(node_file_system_state)
node_ref.set_file_system_state(node_file_system_state) node_ref.set_file_system_state(node_file_system_state)
# Update service states # Update service states
node_services = item["services"] node_services = item["services"]
for service in node_services: for service in node_services:
service_protocol = service["name"] service_protocol = service["name"]
service_state = SOFTWARE_STATE[service["state"]] service_state = SoftwareState[service["state"]]
# Update node service state # Update node service state
node.set_service_state(service_protocol, service_state) node.set_service_state(service_protocol, service_state)
# Update reference node service state # Update reference node service state

183
src/primaite/environment/reward.py
View File

@@ -1,6 +1,9 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements reward function.""" """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.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode from primaite.nodes.service_node import ServiceNode
@@ -28,17 +31,17 @@ def calculate_reward_function(
""" """
reward_value = 0 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(): for node_key, final_node in final_nodes.items():
initial_node = initial_nodes[node_key] initial_node = initial_nodes[node_key]
reference_node = reference_nodes[node_key] reference_node = reference_nodes[node_key]
# Operating State # Hardware State
reward_value += score_node_operating_state( reward_value += score_node_operating_state(
final_node, initial_node, reference_node, config_values final_node, initial_node, reference_node, config_values
) )
# Operating System State # Software State
if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode): if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode):
reward_value += score_node_os_state( reward_value += score_node_os_state(
final_node, initial_node, reference_node, config_values 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): 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: Args:
final_node: The node after red and blue agents take effect 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 config_values: Config values
""" """
score = 0 score = 0
final_node_operating_state = final_node.get_state() final_node_operating_state = final_node.hardware_state
initial_node_operating_state = initial_node.get_state() initial_node_operating_state = initial_node.hardware_state
reference_node_operating_state = reference_node.get_state() reference_node_operating_state = reference_node.hardware_state
if final_node_operating_state == reference_node_operating_state: if final_node_operating_state == reference_node_operating_state:
# All is well - we're no different from the reference situation # 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: else:
# We're different from the reference situation # We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step) # Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_operating_state == HARDWARE_STATE.ON: if initial_node_operating_state == HardwareState.ON:
if reference_node_operating_state == HARDWARE_STATE.OFF: if reference_node_operating_state == HardwareState.OFF:
score += config_values.off_should_be_on 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 score += config_values.resetting_should_be_on
else: else:
pass pass
elif initial_node_operating_state == HARDWARE_STATE.OFF: elif initial_node_operating_state == HardwareState.OFF:
if reference_node_operating_state == HARDWARE_STATE.ON: if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_off 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 score += config_values.resetting_should_be_off
else: else:
pass pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING: elif initial_node_operating_state == HardwareState.RESETTING:
if reference_node_operating_state == HARDWARE_STATE.ON: if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_resetting 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 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 score += config_values.resetting
else: else:
pass 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): 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: Args:
final_node: The node after red and blue agents take effect 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 config_values: Config values
""" """
score = 0 score = 0
final_node_os_state = final_node.get_os_state() final_node_os_state = final_node.software_state
initial_node_os_state = initial_node.get_os_state() initial_node_os_state = initial_node.software_state
reference_node_os_state = reference_node.get_os_state() reference_node_os_state = reference_node.software_state
if final_node_os_state == reference_node_os_state: if final_node_os_state == reference_node_os_state:
# All is well - we're no different from the reference situation # 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: else:
# We're different from the reference situation # We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step) # Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_os_state == SOFTWARE_STATE.GOOD: if initial_node_os_state == SoftwareState.GOOD:
if reference_node_os_state == SOFTWARE_STATE.PATCHING: if reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good 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 score += config_values.compromised_should_be_good
else: else:
pass pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING: elif initial_node_os_state == SoftwareState.PATCHING:
if reference_node_os_state == SOFTWARE_STATE.GOOD: if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching 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 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 score += config_values.patching
else: else:
pass pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED: elif initial_node_os_state == SoftwareState.COMPROMISED:
if reference_node_os_state == SOFTWARE_STATE.GOOD: if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised 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 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 score += config_values.compromised
else: else:
pass pass
@@ -191,59 +194,59 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
config_values: Config values config_values: Config values
""" """
score = 0 score = 0
final_node_services = final_node.get_services() final_node_services: Dict[str, Service] = final_node.services
initial_node_services = initial_node.get_services() initial_node_services: Dict[str, Service] = initial_node.services
reference_node_services = reference_node.get_services() reference_node_services: Dict[str, Service] = reference_node.services
for service_key, final_service in final_node_services.items(): for service_key, final_service in final_node_services.items():
reference_service = reference_node_services[service_key] reference_service = reference_node_services[service_key]
initial_service = initial_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 # All is well - we're no different from the reference situation
score += config_values.all_ok score += config_values.all_ok
else: else:
# We're different from the reference situation # We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step) # Need to compare initial and reference state of node (i.e. at every step)
if initial_service.get_state() == SOFTWARE_STATE.GOOD: if initial_service.software_state == SoftwareState.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.PATCHING: if reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good 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 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 score += config_values.overwhelmed_should_be_good
else: else:
pass pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING: elif initial_service.software_state == SoftwareState.PATCHING:
if reference_service.get_state() == SOFTWARE_STATE.GOOD: if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching 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 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 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 score += config_values.patching
else: else:
pass pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED: elif initial_service.software_state == SoftwareState.COMPROMISED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD: if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised 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 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 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 score += config_values.overwhelmed_should_be_compromised
else: else:
pass pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED: elif initial_service.software_state == SoftwareState.OVERWHELMED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD: if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_overwhelmed 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 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 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 score += config_values.overwhelmed
else: else:
pass 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 reference_node: The node if there had been no red or blue effect
""" """
score = 0 score = 0
final_node_file_system_state = final_node.get_file_system_state_actual() final_node_file_system_state = final_node.file_system_state_actual
initial_node_file_system_state = initial_node.get_file_system_state_actual() initial_node_file_system_state = initial_node.file_system_state_actual
reference_node_file_system_state = reference_node.get_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() final_node_scanning_state = final_node.file_system_scanning
reference_node_scanning_state = reference_node.is_scanning_file_system() reference_node_scanning_state = reference_node.file_system_scanning
# File System State # File System State
if final_node_file_system_state == reference_node_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: else:
# We're different from the reference situation # We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step) # 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 initial_node_file_system_state == FileSystemState.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: if reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_good 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 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 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 score += config_values.destroyed_should_be_good
else: else:
pass pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING: elif initial_node_file_system_state == FileSystemState.REPAIRING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_repairing 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 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 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 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 score += config_values.repairing
else: else:
pass pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING: elif initial_node_file_system_state == FileSystemState.RESTORING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_restoring 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 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 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 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 score += config_values.restoring
else: else:
pass pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT: elif initial_node_file_system_state == FileSystemState.CORRUPT:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_corrupt 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 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 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 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 score += config_values.corrupt
else: else:
pass pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED: elif initial_node_file_system_state == FileSystemState.DESTROYED:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD: if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_destroyed 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 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 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 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 score += config_values.destroyed
else: else:
pass pass

3
src/primaite/links/link.py
View File

@@ -1,5 +1,6 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The link class.""" """The link class."""
from typing import List
from primaite.common.protocol import Protocol from primaite.common.protocol import Protocol
@@ -22,7 +23,7 @@ class Link(object):
self.bandwidth = _bandwidth self.bandwidth = _bandwidth
self.source_node_name = _source_node_name self.source_node_name = _source_node_name
self.dest_node_name = _dest_node_name self.dest_node_name = _dest_node_name
self.protocol_list = [] self.protocol_list: List[Protocol] = []
# Add the default protocols # Add the default protocols
for protocol_name in _services: for protocol_name in _services:

8
src/primaite/main.py
View File

@@ -15,7 +15,7 @@ from stable_baselines3 import A2C, PPO
from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.ppo import MlpPolicy as PPOMlp 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.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import write_transaction_to_file from primaite.transactions.transactions_to_file import write_transaction_to_file
@@ -178,7 +178,7 @@ def load_config_values():
# Reward values # Reward values
# Generic # Generic
config_values.all_ok = int(config_data["allOk"]) 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_on = int(config_data["offShouldBeOn"])
config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"]) config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
config_values.on_should_be_off = int(config_data["onShouldBeOff"]) 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_on = int(config_data["resettingShouldBeOn"])
config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"]) config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
config_values.resetting = int(config_data["resetting"]) 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_patching = int(config_data["goodShouldBePatching"])
config_values.good_should_be_compromised = int( config_values.good_should_be_compromised = int(
config_data["goodShouldBeCompromised"] config_data["goodShouldBeCompromised"]
@@ -331,7 +331,7 @@ try:
config_file_main = open("config/config_main.yaml", "r") config_file_main = open("config/config_main.yaml", "r")
config_data = yaml.safe_load(config_file_main) config_data = yaml.safe_load(config_file_main)
# Create a config class # Create a config class
config_values = config_values_main() config_values = ConfigValuesMain()
# Load in config data # Load in config data
load_config_values() load_config_values()
except Exception: except Exception:

239
src/primaite/nodes/active_node.py
View File

@@ -3,7 +3,14 @@
import logging import logging
from typing import Final 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 from primaite.nodes.node import Node
_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) _LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
@@ -14,216 +21,174 @@ class ActiveNode(Node):
def __init__( def __init__(
self, self,
_id, node_id: str,
_name, name: str,
_type, node_type: NodeType,
_priority, priority: Priority,
_state, hardware_state: HardwareState,
_ip_address, ip_address: str,
_os_state, software_state: SoftwareState,
_file_system_state, file_system_state: FileSystemState,
_config_values, config_values: ConfigValuesMain,
): ):
""" """
Init. Init.
Args: :param node_id: The node ID
_id: The node ID :param name: The node name
_name: The node name :param node_type: The node type (enum)
_type: The node type (enum) :param priority: The node priority (enum)
_priority: The node priority (enum) :param hardware_state: The node Hardware State
_state: The node state (enum) :param ip_address: The node IP address
_ip_address: The node IP address :param software_state: The node Software State
_os_state: The node Operating System state :param file_system_state: The node file system state
_file_system_state: The node file system state :param config_values: The config values
_config_values: The config values
""" """
super().__init__(_id, _name, _type, _priority, _state, _config_values) super().__init__(
self.ip_address = _ip_address node_id, name, node_type, priority, hardware_state, config_values
# Related to O/S )
self.os_state = _os_state self.ip_address: str = ip_address
self.patching_count = 0 # Related to Software
self._software_state: SoftwareState = software_state
self.patching_count: int = 0
# Related to File System # Related to File System
self.file_system_state_actual = _file_system_state self.file_system_state_actual: FileSystemState = file_system_state
self.file_system_state_observed = _file_system_state self.file_system_state_observed: FileSystemState = file_system_state
self.file_system_scanning = False self.file_system_scanning: bool = False
self.file_system_scanning_count = 0 self.file_system_scanning_count: int = 0
self.file_system_action_count = 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: :return: The software_state.
_ip_address: IP address
""" """
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: :param software_state: Software State.
IP address
""" """
return self.ip_address if self.hardware_state != HardwareState.OFF:
self._software_state = software_state
def set_os_state(self, _os_state): if software_state == SoftwareState.PATCHING:
"""
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:
self.patching_count = self.config_values.os_patching_duration self.patching_count = self.config_values.os_patching_duration
else: else:
_LOGGER.info( _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"changed. "
f"Node:{self.id}, " f"Node.node_id:{self.node_id}, "
f"Node Operating State:{self.operating_state}, " f"Node.hardware_state:{self.hardware_state}, "
f"Node Operating System State:{self.os_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: Args:
_os_state: Operating system state software_state: Software State
""" """
if self.operating_state != HARDWARE_STATE.OFF: if self.hardware_state != HardwareState.OFF:
if self.os_state != SOFTWARE_STATE.COMPROMISED: if self._software_state != SoftwareState.COMPROMISED:
self.os_state = _os_state self._software_state = software_state
if _os_state == SOFTWARE_STATE.PATCHING: if software_state == SoftwareState.PATCHING:
self.patching_count = self.config_values.os_patching_duration self.patching_count = self.config_values.os_patching_duration
else: else:
_LOGGER.info( _LOGGER.info(
f"The Nodes operating state is OFF so OS State cannot be changed." f"The Nodes hardware state is OFF so OS State cannot be changed."
f"Node:{self.id}, " f"Node.node_id:{self.node_id}, "
f"Node Operating State:{self.operating_state}, " f"Node.hardware_state:{self.hardware_state}, "
f"Node Operating System State:{self.os_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): def update_os_patching_status(self):
"""Updates operating system status based on patching cycle.""" """Updates operating system status based on patching cycle."""
self.patching_count -= 1 self.patching_count -= 1
if self.patching_count <= 0: if self.patching_count <= 0:
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). Sets the file system state (actual and observed).
Args: 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:
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.file_system_action_count = (
self.config_values.file_system_repairing_limit self.config_values.file_system_repairing_limit
) )
self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING self.file_system_state_observed = FileSystemState.REPAIRING
elif _file_system_state == FILE_SYSTEM_STATE.RESTORING: elif file_system_state == FileSystemState.RESTORING:
self.file_system_action_count = ( self.file_system_action_count = (
self.config_values.file_system_restoring_limit self.config_values.file_system_restoring_limit
) )
self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING self.file_system_state_observed = FileSystemState.RESTORING
elif _file_system_state == FILE_SYSTEM_STATE.GOOD: elif file_system_state == FileSystemState.GOOD:
self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD self.file_system_state_observed = FileSystemState.GOOD
else: else:
_LOGGER.info( _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"cannot be changed. "
f"Node:{self.id}, " f"Node.node_id:{self.node_id}, "
f"Node Operating State:{self.operating_state}, " f"Node.hardware_state:{self.hardware_state}, "
f"Node File System State:{self.file_system_state_actual}", 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. 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 Use for green PoL to prevent it overturning a compromised state
Args: 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 ( if (
self.file_system_state_actual != FILE_SYSTEM_STATE.CORRUPT self.file_system_state_actual != FileSystemState.CORRUPT
and self.file_system_state_actual != FILE_SYSTEM_STATE.DESTROYED 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.file_system_action_count = (
self.config_values.file_system_repairing_limit self.config_values.file_system_repairing_limit
) )
self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING self.file_system_state_observed = FileSystemState.REPAIRING
elif _file_system_state == FILE_SYSTEM_STATE.RESTORING: elif file_system_state == FileSystemState.RESTORING:
self.file_system_action_count = ( self.file_system_action_count = (
self.config_values.file_system_restoring_limit self.config_values.file_system_restoring_limit
) )
self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING self.file_system_state_observed = FileSystemState.RESTORING
elif _file_system_state == FILE_SYSTEM_STATE.GOOD: elif file_system_state == FileSystemState.GOOD:
self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD self.file_system_state_observed = FileSystemState.GOOD
else: else:
_LOGGER.info( _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"compromised) cannot be changed. "
f"Node:{self.id}, " f"Node.node_id:{self.node_id}, "
f"Node Operating State:{self.operating_state}, " f"Node.hardware_state:{self.hardware_state}, "
f"Node File System State:{self.file_system_state_actual}", 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): def start_file_system_scan(self):
"""Starts a file system scan.""" """Starts a file system scan."""
self.file_system_scanning = True self.file_system_scanning = True
self.file_system_scanning_count = self.config_values.file_system_scanning_limit 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): def update_file_system_state(self):
"""Updates file system status based on scanning/restore/repair cycle.""" """Updates file system status based on scanning/restore/repair cycle."""
# Deprecate both the action count (for restoring or reparing) and the scanning count # 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: if self.file_system_action_count <= 0:
self.file_system_action_count = 0 self.file_system_action_count = 0
if ( if (
self.file_system_state_actual == FILE_SYSTEM_STATE.REPAIRING self.file_system_state_actual == FileSystemState.REPAIRING
or self.file_system_state_actual == FILE_SYSTEM_STATE.RESTORING or self.file_system_state_actual == FileSystemState.RESTORING
): ):
self.file_system_state_actual = FILE_SYSTEM_STATE.GOOD self.file_system_state_actual = FileSystemState.GOOD
self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD self.file_system_state_observed = FileSystemState.GOOD
# Scanning updates # Scanning updates
if self.file_system_scanning == True and self.file_system_scanning_count < 0: if self.file_system_scanning == True and self.file_system_scanning_count < 0:

139
src/primaite/nodes/node.py
View File

@@ -1,135 +1,56 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The base Node class.""" """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: class Node:
"""Node class.""" """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. Init.
Args: :param node_id: The node id.
_id: The node id :param name: The name of the node.
_name: The name of the node :param node_type: The type of the node.
_type: The type of the node :param priority: The priority of the node.
_priority: The priority of the node :param hardware_state: The state of the node.
_state: The state of the node :param config_values: Config values.
""" """
self.id = _id self.node_id: Final[str] = node_id
self.name = _name self.name: Final[str] = name
self.type = _type self.node_type: Final[NodeType] = node_type
self.priority = _priority self.priority = priority
self.operating_state = _state self.hardware_state: HardwareState = hardware_state
self.resetting_count = 0 self.resetting_count: int = 0
self.config_values = _config_values self.config_values: ConfigValuesMain = config_values
def __repr__(self): def __repr__(self):
"""Returns the name of the node.""" """Returns the name of the node."""
return self.name 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): def turn_on(self):
"""Sets the node state to ON.""" """Sets the node state to ON."""
self.operating_state = HARDWARE_STATE.ON self.hardware_state = HardwareState.ON
def turn_off(self): def turn_off(self):
"""Sets the node state to OFF.""" """Sets the node state to OFF."""
self.operating_state = HARDWARE_STATE.OFF self.hardware_state = HardwareState.OFF
def reset(self): def reset(self):
"""Sets the node state to Resetting and starts the reset count.""" """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 self.resetting_count = self.config_values.node_reset_duration
def update_resetting_status(self): def update_resetting_status(self):
@@ -137,4 +58,4 @@ class Node:
self.resetting_count -= 1 self.resetting_count -= 1
if self.resetting_count <= 0: if self.resetting_count <= 0:
self.resetting_count = 0 self.resetting_count = 0
self.operating_state = HARDWARE_STATE.ON self.hardware_state = HardwareState.ON

7
src/primaite/nodes/node_state_instruction_red.py
View File

@@ -1,5 +1,6 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Defines node behaviour for Green PoL.""" """Defines node behaviour for Green PoL."""
from primaite.common.enums import NodePOLType
class NodeStateInstructionRed(object): class NodeStateInstructionRed(object):
@@ -12,7 +13,7 @@ class NodeStateInstructionRed(object):
_end_step, _end_step,
_target_node_id, _target_node_id,
_pol_initiator, _pol_initiator,
_pol_type, _pol_type: NodePOLType,
pol_protocol, pol_protocol,
_pol_state, _pol_state,
_pol_source_node_id, _pol_source_node_id,
@@ -40,7 +41,7 @@ class NodeStateInstructionRed(object):
self.end_step = _end_step self.end_step = _end_step
self.target_node_id = _target_node_id self.target_node_id = _target_node_id
self.initiator = _pol_initiator 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.service_name = pol_protocol # Not used when not a service instruction
self.state = _pol_state self.state = _pol_state
self.source_node_id = _pol_source_node_id self.source_node_id = _pol_source_node_id
@@ -83,7 +84,7 @@ class NodeStateInstructionRed(object):
""" """
return self.initiator return self.initiator
def get_pol_type(self): def get_pol_type(self) -> NodePOLType:
""" """
Gets the node pattern of life type (enum). Gets the node pattern of life type (enum).

40
src/primaite/nodes/passive_node.py
View File

@@ -1,32 +1,44 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The Passive Node class (i.e. an actuator).""" """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 from primaite.nodes.node import Node
class PassiveNode(Node): class PassiveNode(Node):
"""The Passive Node class.""" """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. Init.
Args: :param node_id: The node id.
_id: The node id :param name: The name of the node.
_name: The name of the node :param node_type: The type of the node.
_type: The type of the node :param priority: The priority of the node.
_priority: The priority of the node :param hardware_state: The state of the node.
_state: The state of the node :param config_values: Config values.
""" """
# Pass through to Super for now # 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: No concept of IP address for passive nodes for now.
The node IP address
:return: The node IP address.
""" """
# No concept of IP address for passive nodes for now
return "" return ""

238
src/primaite/nodes/service_node.py
View File

@@ -1,9 +1,17 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""A Service Node (i.e. not an actuator).""" """A Service Node (i.e. not an actuator)."""
import logging 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 from primaite.nodes.active_node import ActiveNode
_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) _LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
@@ -14,185 +22,169 @@ class ServiceNode(ActiveNode):
def __init__( def __init__(
self, self,
_id, node_id: str,
_name, name: str,
_type, node_type: NodeType,
_priority, priority: Priority,
_state, hardware_state: HardwareState,
_ip_address, ip_address: str,
_os_state, software_state: SoftwareState,
_file_system_state, file_system_state: FileSystemState,
_config_values, config_values: ConfigValuesMain,
): ):
""" """
Init. Init.
Args: :param node_id: The node ID
_id: The node id :param name: The node name
_name: The name of the node :param node_type: The node type (enum)
_type: The type of the node :param priority: The node priority (enum)
_priority: The priority of the node :param hardware_state: The node Hardware State
_state: The state of the node :param ip_address: The node IP address
_ipAddress: The IP address of the node :param software_state: The node Software State
_osState: The operating system state of the node :param file_system_state: The node file system state
_file_system_state: The file system state of the node :param config_values: The config values
""" """
super().__init__( super().__init__(
_id, node_id,
_name, name,
_type, node_type,
_priority, priority,
_state, hardware_state,
_ip_address, ip_address,
_os_state, software_state,
_file_system_state, file_system_state,
_config_values, 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. Adds a service to the node.
Args: :param service: The service to add
_service: The service to add
""" """
self.services[_service.get_name()] = _service self.services[service.name] = service
def get_services(self): def has_service(self, protocol_name: str) -> bool:
"""
Gets the dictionary of services on this node.
Returns:
Dictionary of services on this node
"""
return self.services
def has_service(self, _protocol):
""" """
Indicates whether a service is on a node. Indicates whether a service is on a node.
Returns: :param protocol_name: The service (protocol)e.
True if service (protocol) is on the node :return: True if service (protocol) is on the node, otherwise False.
""" """
for service_key, service_value in self.services.items(): for service_key, service_value in self.services.items():
if service_key == _protocol: if service_key == protocol_name:
return True return True
else:
pass
return False 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. Indicates whether a service is in a running state on the node.
Returns: :param protocol_name: The service (protocol)
True if service (protocol) is in a running state on the node :return: True if service (protocol) is in a running state on the
node, otherwise False.
""" """
for service_key, service_value in self.services.items(): for service_key, service_value in self.services.items():
if service_key == _protocol: if service_key == protocol_name:
if service_value.get_state() != SOFTWARE_STATE.PATCHING: if service_value.software_state != SoftwareState.PATCHING:
return True return True
else: else:
return False return False
else:
pass
return False 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. Indicates whether a service is in an overwhelmed state on the node.
Returns: :param protocol_name: The service (protocol)
True if service (protocol) is in an overwhelmed state on the node :return: True if service (protocol) is in an overwhelmed state on the
node, otherwise False.
""" """
for service_key, service_value in self.services.items(): for service_key, service_value in self.services.items():
if service_key == _protocol: if service_key == protocol_name:
if service_value.get_state() == SOFTWARE_STATE.OVERWHELMED: if service_value.software_state == SoftwareState.OVERWHELMED:
return True return True
else: else:
return False return False
else:
pass
return False 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: :param protocol_name: The service (protocol).
_protocol: The service (protocol) :param software_state: The software_state.
_state: The state value
""" """
if self.operating_state != HARDWARE_STATE.OFF: if self.hardware_state != HardwareState.OFF:
for service_key, service_value in self.services.items(): service_key = protocol_name
if service_key == _protocol: service_value = self.services.get(service_key)
# Can't set to compromised if you're in a patching state if service_value:
if ( # Can't set to compromised if you're in a patching state
_state == SOFTWARE_STATE.COMPROMISED if (
and service_value.get_state() != SOFTWARE_STATE.PATCHING software_state == SoftwareState.COMPROMISED
) or _state != SOFTWARE_STATE.COMPROMISED: and service_value.software_state != SoftwareState.PATCHING
service_value.set_state(_state) ) or software_state != SoftwareState.COMPROMISED:
else: service_value.software_state = software_state
# Do nothing if software_state == SoftwareState.PATCHING:
pass service_value.patching_count = (
if _state == SOFTWARE_STATE.PATCHING: 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[<key>]:{protocol_name}, "
f"Node.services[<key>].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 = ( service_value.patching_count = (
self.config_values.service_patching_duration self.config_values.service_patching_duration
) )
else:
# Do nothing
pass
else: else:
_LOGGER.info( _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"cannot be changed. "
f"Node:{self.id}, " f"Node.node_id:{self.node_id}, "
f"Node Operating State:{self.operating_state}, " f"Node.hardware_state:{self.hardware_state}, "
f"Node Service Protocol:{_protocol}, " f"Node.services[<key>]:{protocol_name}, "
f"Node Service State: {_state}" f"Node.services[<key>].software_state:{software_state}"
) )
def set_service_state_if_not_compromised(self, _protocol, _state): def get_service_state(self, protocol_name):
"""
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):
""" """
Gets the state of a service. Gets the state of a service.
Returns: :return: The software_state of the service.
The state of the service
""" """
for service_key, service_value in self.services.items(): service_key = protocol_name
if service_key == _protocol: service_value = self.services.get(service_key)
return service_value.get_state() if service_value:
return service_value.software_state
def update_services_patching_status(self): def update_services_patching_status(self):
"""Updates the patching counter for any service that are patching.""" """Updates the patching counter for any service that are patching."""
for service_key, service_value in self.services.items(): 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() service_value.reduce_patching_count()

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

@@ -1,16 +1,30 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements Pattern of Life on the network (nodes and links).""" """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.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.nodes.service_node import ServiceNode
from primaite.pol.ier import IER
_VERBOSE = False _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). 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] dest_node = nodes[dest_node_id]
# 1. Check the source node situation # 1. Check the source node situation
if source_node.get_type() == TYPE.SWITCH: if source_node.node_type == NodeType.SWITCH:
# It's a switch # It's a switch
if ( if (
source_node.get_state() == HARDWARE_STATE.ON source_node.hardware_state == HardwareState.ON
and source_node.get_os_state() != SOFTWARE_STATE.PATCHING and source_node.software_state != SoftwareState.PATCHING
): ):
source_valid = True source_valid = True
else: else:
# IER no longer valid # IER no longer valid
source_valid = False source_valid = False
elif source_node.get_type() == TYPE.ACTUATOR: elif source_node.node_type == NodeType.ACTUATOR:
# It's an actuator # It's an actuator
# TO DO # TO DO
pass pass
else: else:
# It's not a switch or an actuator (so active node) # It's not a switch or an actuator (so active node)
if ( if (
source_node.get_state() == HARDWARE_STATE.ON source_node.hardware_state == HardwareState.ON
and source_node.get_os_state() != SOFTWARE_STATE.PATCHING and source_node.software_state != SoftwareState.PATCHING
): ):
if source_node.has_service(protocol): if source_node.has_service(protocol):
if source_node.service_running( if source_node.service_running(
@@ -87,24 +101,24 @@ def apply_iers(network, nodes, links, iers, acl, step):
source_valid = False source_valid = False
# 2. Check the dest node situation # 2. Check the dest node situation
if dest_node.get_type() == TYPE.SWITCH: if dest_node.node_type == NodeType.SWITCH:
# It's a switch # It's a switch
if ( if (
dest_node.get_state() == HARDWARE_STATE.ON dest_node.hardware_state == HardwareState.ON
and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING and dest_node.software_state != SoftwareState.PATCHING
): ):
dest_valid = True dest_valid = True
else: else:
# IER no longer valid # IER no longer valid
dest_valid = False dest_valid = False
elif dest_node.get_type() == TYPE.ACTUATOR: elif dest_node.node_type == NodeType.ACTUATOR:
# It's an actuator # It's an actuator
pass pass
else: else:
# It's not a switch or an actuator (so active node) # It's not a switch or an actuator (so active node)
if ( if (
dest_node.get_state() == HARDWARE_STATE.ON dest_node.hardware_state == HardwareState.ON
and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING and dest_node.software_state != SoftwareState.PATCHING
): ):
if dest_node.has_service(protocol): if dest_node.has_service(protocol):
if dest_node.service_running( 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 # 3. Check that the ACL doesn't block it
acl_block = acl.is_blocked( 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 acl_block:
if _VERBOSE: if _VERBOSE:
print( print(
"ACL block on source: " "ACL block on source: "
+ source_node.get_ip_address() + source_node.ip_address
+ ", dest: " + ", dest: "
+ dest_node.get_ip_address() + dest_node.ip_address
+ ", protocol: " + ", protocol: "
+ protocol + protocol
+ ", port: " + ", 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 # We might have a switch in the path, so check all nodes are operational
for node in path_node_list: for node in path_node_list:
if ( if (
node.get_state() != HARDWARE_STATE.ON node.hardware_state != HardwareState.ON
or node.get_os_state() == SOFTWARE_STATE.PATCHING or node.software_state == SoftwareState.PATCHING
): ):
path_valid = False path_valid = False
@@ -215,7 +229,11 @@ def apply_iers(network, nodes, links, iers, acl, step):
pass 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. Applies node pattern of life.
@@ -239,15 +257,15 @@ def apply_node_pol(nodes, node_pol, step):
# continue -------------------------- # continue --------------------------
node = nodes[node_id] node = nodes[node_id]
if node_pol_type == NODE_POL_TYPE.OPERATING: if node_pol_type == NodePOLType.OPERATING:
# Change operating state # Change hardware state
node.set_state(state) node.hardware_state = state
elif node_pol_type == NODE_POL_TYPE.OS: elif node_pol_type == NodePOLType.OS:
# Change OS state # Change OS state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this # Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.set_os_state_if_not_compromised(state) node.set_software_state_if_not_compromised(state)
elif node_pol_type == NODE_POL_TYPE.SERVICE: elif node_pol_type == NodePOLType.SERVICE:
# Change a service state # Change a service state
# Don't allow PoL to fix something that is compromised. Only the Blue agent can do this # Don't allow PoL to fix something that is compromised. Only the Blue agent can do this
if isinstance(node, ServiceNode): if isinstance(node, ServiceNode):

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

@@ -1,22 +1,35 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements POL on the network (nodes and links) resulting from the red agent attack.""" """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 ( from primaite.common.enums import (
HARDWARE_STATE, HardwareState,
NODE_POL_INITIATOR, NodePOLInitiator,
NODE_POL_TYPE, NodePOLType,
SOFTWARE_STATE, NodeType,
TYPE, SoftwareState,
) )
from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode 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.nodes.service_node import ServiceNode
from primaite.pol.ier import IER
_VERBOSE = False _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. 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] dest_node = nodes[dest_node_id]
# 1. Check the source node situation # 1. Check the source node situation
if source_node.get_type() == TYPE.SWITCH: if source_node.node_type == NodeType.SWITCH:
# It's a switch # It's a switch
if source_node.get_state() == HARDWARE_STATE.ON: if source_node.hardware_state == HardwareState.ON:
source_valid = True source_valid = True
else: else:
# IER no longer valid # IER no longer valid
source_valid = False source_valid = False
elif source_node.get_type() == TYPE.ACTUATOR: elif source_node.node_type == NodeType.ACTUATOR:
# It's an actuator # It's an actuator
# TO DO # TO DO
pass pass
else: else:
# It's not a switch or an actuator (so active node) # 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): if source_node.has_service(protocol):
# Red agents IERs can only be valid if the source service is in a compromised state # Red agents IERs can only be valid if the source service is in a compromised state
if ( if (
source_node.get_service_state(protocol) source_node.get_service_state(protocol)
== SOFTWARE_STATE.COMPROMISED == SoftwareState.COMPROMISED
): ):
source_valid = True source_valid = True
else: else:
@@ -86,19 +99,19 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
source_valid = False source_valid = False
# 2. Check the dest node situation # 2. Check the dest node situation
if dest_node.get_type() == TYPE.SWITCH: if dest_node.node_type == NodeType.SWITCH:
# It's a switch # It's a switch
if dest_node.get_state() == HARDWARE_STATE.ON: if dest_node.hardware_state == HardwareState.ON:
dest_valid = True dest_valid = True
else: else:
# IER no longer valid # IER no longer valid
dest_valid = False dest_valid = False
elif dest_node.get_type() == TYPE.ACTUATOR: elif dest_node.node_type == NodeType.ACTUATOR:
# It's an actuator # It's an actuator
pass pass
else: else:
# It's not a switch or an actuator (so active node) # 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): if dest_node.has_service(protocol):
# We don't care what state the destination service is in for an IER # We don't care what state the destination service is in for an IER
dest_valid = True 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 # 3. Check that the ACL doesn't block it
acl_block = acl.is_blocked( 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 acl_block:
if _VERBOSE: if _VERBOSE:
print( print(
"ACL block on source: " "ACL block on source: "
+ source_node.get_ip_address() + source_node.ip_address
+ ", dest: " + ", dest: "
+ dest_node.get_ip_address() + dest_node.ip_address
+ ", protocol: " + ", protocol: "
+ protocol + protocol
+ ", port: " + ", 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 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 # We're assuming here that red agents can get past switches that are patching
for node in path_node_list: for node in path_node_list:
if node.get_state() != HARDWARE_STATE.ON: if node.hardware_state != HardwareState.ON:
path_valid = False path_valid = False
if path_valid: if path_valid:
@@ -207,7 +220,12 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
pass 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. 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: if step >= start_step and step <= stop_step:
# continue -------------------------- # continue --------------------------
target_node = nodes[target_node_id] target_node: NodeType = nodes[target_node_id]
# Based the action taken on the initiator type # 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 # No conditions required, just apply the change
passed_checks = True passed_checks = True
elif initiator == NODE_POL_INITIATOR.IER: elif initiator == NodePOLInitiator.IER:
# Need to check there is a red IER incoming # Need to check there is a red IER incoming
passed_checks = is_red_ier_incoming(target_node, iers, pol_type) passed_checks = is_red_ier_incoming(target_node, iers, pol_type)
elif initiator == NODE_POL_INITIATOR.SERVICE: elif initiator == NodePOLInitiator.SERVICE:
# Need to check the condition of a service on another node # Need to check the condition of a service on another node
source_node = nodes[source_node_id] source_node = nodes[source_node_id]
if source_node.has_service(source_node_service_name): if source_node.has_service(source_node_service_name):
if ( if (
source_node.get_service_state(source_node_service_name) 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 passed_checks = True
else: 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) # Only apply the PoL if the checks have passed (based on the initiator type)
if passed_checks: if passed_checks:
# Apply the change # Apply the change
if pol_type == NODE_POL_TYPE.OPERATING: if pol_type == NodePOLType.OPERATING:
# Change operating state # Change hardware state
target_node.set_state(state) target_node.hardware_state = state
elif pol_type == NODE_POL_TYPE.OS: elif pol_type == NodePOLType.OS:
# Change OS state # Change OS state
if isinstance(target_node, ActiveNode) or isinstance( if isinstance(target_node, ActiveNode) or isinstance(
target_node, ServiceNode target_node, ServiceNode
): ):
target_node.set_os_state(state) target_node.software_state = state
elif pol_type == NODE_POL_TYPE.SERVICE: elif pol_type == NodePOLType.SERVICE:
# Change a service state # Change a service state
if isinstance(target_node, ServiceNode): if isinstance(target_node, ServiceNode):
target_node.set_service_state(service_name, state) 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. 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(): for ier_key, ier_value in iers.items():
if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id: if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id:
if ( if (
node_pol_type == NODE_POL_TYPE.OPERATING node_pol_type == NodePOLType.OPERATING
or node_pol_type == NODE_POL_TYPE.OS or node_pol_type == NodePOLType.OS
or node_pol_type == NODE_POL_TYPE.FILE 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 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 # Check if the service is present on the node and running
ier_protocol = ier_value.get_protocol() ier_protocol = ier_value.get_protocol()
if isinstance(node, ServiceNode): if isinstance(node, ServiceNode):

14
tests/config/one_node_states_on_off_lay_down_config.yaml
View File

@@ -9,15 +9,15 @@
serviceList: serviceList:
- name: ftp - name: ftp
- itemType: NODE - itemType: NODE
id: '1' node_id: '1'
name: node name: node
baseType: SERVICE node_class: SERVICE
nodeType: COMPUTER node_type: COMPUTER
priority: P1 priority: P1
hardwareState: 'ON' hardware_state: 'ON'
ipAddress: 192.168.0.1 ip_address: 192.168.0.1
softwareState: GOOD software_state: GOOD
fileSystemState: GOOD file_system_state: GOOD
services: services:
- name: ftp - name: ftp
port: '21' port: '21'

4
tests/config/one_node_states_on_off_main_config.yaml
View File

@@ -26,7 +26,7 @@ observationSpaceHighValue: 1000000000
# Reward values # Reward values
# Generic # Generic
allOk: 0 allOk: 0
# Node Operating State # Node Hardware State
offShouldBeOn: -10 offShouldBeOn: -10
offShouldBeResetting: -5 offShouldBeResetting: -5
onShouldBeOff: -2 onShouldBeOff: -2
@@ -34,7 +34,7 @@ onShouldBeResetting: -5
resettingShouldBeOn: -5 resettingShouldBeOn: -5
resettingShouldBeOff: -2 resettingShouldBeOff: -2
resetting: -3 resetting: -3
# Node O/S or Service State # Node Software or Service State
goodShouldBePatching: 2 goodShouldBePatching: 2
goodShouldBeCompromised: 5 goodShouldBeCompromised: 5
goodShouldBeOverwhelmed: 5 goodShouldBeOverwhelmed: 5

8
tests/conftest.py
View File

@@ -5,7 +5,7 @@ from typing import Union
import yaml 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 from primaite.environment.primaite_env import Primaite
ACTION_SPACE_NODE_VALUES = 1 ACTION_SPACE_NODE_VALUES = 1
@@ -32,7 +32,7 @@ def _get_primaite_env_from_config(
# Reward values # Reward values
# Generic # Generic
config_values.all_ok = int(config_data["allOk"]) 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_on = int(config_data["offShouldBeOn"])
config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"]) config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
config_values.on_should_be_off = int(config_data["onShouldBeOff"]) 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_on = int(config_data["resettingShouldBeOn"])
config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"]) config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
config_values.resetting = int(config_data["resetting"]) 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_patching = int(config_data["goodShouldBePatching"])
config_values.good_should_be_compromised = int( config_values.good_should_be_compromised = int(
config_data["goodShouldBeCompromised"] config_data["goodShouldBeCompromised"]
@@ -160,7 +160,7 @@ def _get_primaite_env_from_config(
config_file_main = open(main_config_path, "r") config_file_main = open(main_config_path, "r")
config_data = yaml.safe_load(config_file_main) config_data = yaml.safe_load(config_file_main)
# Create a config class # Create a config class
config_values = config_values_main() config_values = ConfigValuesMain()
# Load in config data # Load in config data
load_config_values() load_config_values()
env = Primaite(config_values, []) env = Primaite(config_values, [])

54
tests/test_active_node.py
View File

@@ -1,22 +1,22 @@
"""Used to test Active Node functions.""" """Used to test Active Node functions."""
import pytest 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 from primaite.nodes.active_node import ActiveNode
@pytest.mark.parametrize( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), (HardwareState.OFF, SoftwareState.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), (HardwareState.ON, SoftwareState.OVERWHELMED),
], ],
) )
def test_os_state_change(operating_state, expected_state): 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( active_node = ActiveNode(
0, 0,
@@ -25,28 +25,28 @@ def test_os_state_change(operating_state, expected_state):
"1", "1",
operating_state, operating_state,
"192.168.0.1", "192.168.0.1",
SOFTWARE_STATE.GOOD, SoftwareState.GOOD,
"GOOD", "GOOD",
1, 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( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), (HardwareState.OFF, SoftwareState.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), (HardwareState.ON, SoftwareState.OVERWHELMED),
], ],
) )
def test_os_state_change_if_not_compromised(operating_state, expected_state): 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( active_node = ActiveNode(
0, 0,
@@ -55,25 +55,25 @@ def test_os_state_change_if_not_compromised(operating_state, expected_state):
"1", "1",
operating_state, operating_state,
"192.168.0.1", "192.168.0.1",
SOFTWARE_STATE.GOOD, SoftwareState.GOOD,
"GOOD", "GOOD",
1, 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( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD), (HardwareState.OFF, FileSystemState.GOOD),
(HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT), (HardwareState.ON, FileSystemState.CORRUPT),
], ],
) )
def test_file_system_change(operating_state, expected_state): 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( active_node = ActiveNode(
0, 0,
"node", "node",
@@ -82,27 +82,27 @@ def test_file_system_change(operating_state, expected_state):
operating_state, operating_state,
"192.168.0.1", "192.168.0.1",
"COMPROMISED", "COMPROMISED",
FILE_SYSTEM_STATE.GOOD, FileSystemState.GOOD,
1, 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( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD), (HardwareState.OFF, FileSystemState.GOOD),
(HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT), (HardwareState.ON, FileSystemState.CORRUPT),
], ],
) )
def test_file_system_change_if_not_compromised(operating_state, expected_state): def test_file_system_change_if_not_compromised(operating_state, expected_state):
""" """
Test that a node cannot change its file system 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( active_node = ActiveNode(
0, 0,
@@ -112,10 +112,10 @@ def test_file_system_change_if_not_compromised(operating_state, expected_state):
operating_state, operating_state,
"192.168.0.1", "192.168.0.1",
"GOOD", "GOOD",
FILE_SYSTEM_STATE.GOOD, FileSystemState.GOOD,
1, 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

2
tests/test_reward.py
View File

@@ -19,7 +19,7 @@ def test_rewards_are_being_penalised_at_each_step_function():
File System State: goodShouldBeCorrupt = 5 (between Steps 1 & 3) File System State: goodShouldBeCorrupt = 5 (between Steps 1 & 3)
Hardware State: onShouldBeOff = -2 (between Steps 4 & 6) Hardware State: onShouldBeOff = -2 (between Steps 4 & 6)
Service State: goodShouldBeCompromised = 5 (between Steps 7 & 9) 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 Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26
Step Count: 13 Step Count: 13

22
tests/test_service_node.py
View File

@@ -1,7 +1,7 @@
"""Used to test Service Node functions.""" """Used to test Service Node functions."""
import pytest 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.common.service import Service
from primaite.nodes.service_node import ServiceNode from primaite.nodes.service_node import ServiceNode
@@ -9,15 +9,15 @@ from primaite.nodes.service_node import ServiceNode
@pytest.mark.parametrize( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), (HardwareState.OFF, SoftwareState.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), (HardwareState.ON, SoftwareState.OVERWHELMED),
], ],
) )
def test_service_state_change(operating_state, expected_state): def test_service_state_change(operating_state, expected_state):
""" """
Test that a node cannot change the state of a running service. 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( service_node = ServiceNode(
0, 0,
@@ -30,10 +30,10 @@ def test_service_state_change(operating_state, expected_state):
"RESTORING", "RESTORING",
1, 1,
) )
service = Service("TCP", 80, SOFTWARE_STATE.GOOD) service = Service("TCP", 80, SoftwareState.GOOD)
service_node.add_service(service) 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 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( @pytest.mark.parametrize(
"operating_state, expected_state", "operating_state, expected_state",
[ [
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD), (HardwareState.OFF, SoftwareState.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED), (HardwareState.ON, SoftwareState.OVERWHELMED),
], ],
) )
def test_service_state_change_if_not_comprised(operating_state, expected_state): def test_service_state_change_if_not_comprised(operating_state, expected_state):
""" """
Test that a node cannot change the state of a running service. 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( service_node = ServiceNode(
0, 0,
@@ -62,9 +62,9 @@ def test_service_state_change_if_not_comprised(operating_state, expected_state):
"RESTORING", "RESTORING",
1, 1,
) )
service = Service("TCP", 80, SOFTWARE_STATE.GOOD) service = Service("TCP", 80, SoftwareState.GOOD)
service_node.add_service(service) 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 assert service_node.get_service_state("TCP") == expected_state