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

Merged PR 60: #1355 - Carried out full renaming in node.py, active_node.py, passive_node.py...

Browse Source 
**The following changes are made to constructor params in the Node class and its children (ActiveNode, PassiveNode, and ServiceNode):**
- _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
- Add type hints to all params.

(node_id, name, and ip_address are str, states and other defines types are the respective enums, leave config_values without a type for now.)

**The following changes are made to instance variables in the Node class and its children:**
- self.type -> self.node_type
- self.operating_state -> self.hardware_state
- self.os_state -> self.software_state
- Add type hints to all instance variables.

(node_id, name, and ip_address are str, states and other defines types are the respective enums, leave config_values without a type for now.)

**The following changes are made to the config files where itemType is NODE:**
- itemType -> item_type
- id -> node_id
- portsList -> ports_list
- serviceList -> service_list
- baseType -> base_type
- nodeType -> node_type
- hardwareState -> hardware_state
- ipAddress -> ip_address
- softwareState -> software_state
- fileSystemState -> file_system_state

**The following changes are made in the primaite/environment/primaite_env.py module:
In the create_node function, the id of the node needs to be retrieved using the new "node_id" key.**
- _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

**Few other cosmetic/code style changes too:**
- Enum classes renamed to use CamelCase.
Started refactoring out unnescessary getters and setters by using `@property` and `@<property name>.setter`.
- Have started to add Type Hints.
- Have started to move docstrings over to the Sphinx ReStructured text format.

Related work items: #1355
This commit is contained in:
Christopher McCarthy
2023-05-26 09:01:21 +00:00
parent 00e9d1f88d b62ea525e9
commit f7638ddb0c
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
E712
D401
F811
exclude =
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​
* Nodes and links (edges) host Python classes in order to present attributes and methods (and hence, a more representative model of a platform / system)​
* A ‘green agent’ Information Exchange Requirement (IER) function allows the representation of traffic (protocols and loading) on any / all links. Application of IERs is based on the status of node operating systems and services​
* A ‘green agent’ node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. Operating state, Operating System state, Service state, File System state)​
* A ‘green agent’ node Pattern-of-Life (PoL) function allows the representation of core behaviours on nodes (e.g. Hardware state, Software State, Service state, File System state)​
* An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP, destination IP, protocol and port). Application of IERs adheres to any ACL restrictions​
* Presents an OpenAI Gym interface to the environment, allowing integration with any OpenAI Gym compliant defensive agents ​
* Red agent activity based on ‘red’ IERs and ‘red’ PoL
@@ -31,12 +31,12 @@ An inheritance model has been adopted in order to model nodes. All nodes have th
* Name
* Type (e.g. computer, switch, RTU - enumeration)​
* Priority (P1, P2, P3, P4 or P5 - enumeration)​
* Operating State (ON, OFF, RESETTING - enumeration)​
* Hardware State (ON, OFF, RESETTING - enumeration)​
Active Nodes also have the following attributes (Class: Active Node):​
* IP Address​
* Operating System State (GOOD, PATCHING, COMPROMISED - enumeration)​
* Software State (GOOD, PATCHING, COMPROMISED - enumeration)​
* File System State (GOOD, CORRUPT, DESTROYED, REPAIRING, RESTORING - enumeration)
Service Nodes also have the following attributes (Class: Service Node)​:
@@ -101,7 +101,7 @@ The status changes that can be made to a node are as follows:
* All Nodes:
* Operating State:
* Hardware State:
* ON
* OFF
@@ -109,7 +109,7 @@ The status changes that can be made to a node are as follows:
* Active Nodes and Service Nodes:
* Operating System State:
* Software State:
* GOOD
* PATCHING - when a status of patching is entered, the node will automatically exit this state after a number of steps (as defined by the osPatchingDuration configuration item) after which it returns to a GOOD state
@@ -185,7 +185,7 @@ Observation Spaces
The OpenAI Gym observation space provides the status of all nodes and links across the whole system:​
* Nodes (in terms of operating state, operating system state, file system state and services state) ​
* Nodes (in terms of hardware state, Software State, file system state and services state) ​
* Links (in terms of current loading for each service/protocol)
An example observation space is provided below:
@@ -196,8 +196,8 @@ An example observation space is provided below:
* -
- ID
- Operating State
- O/S State
- Hardware State
- SoftwareState
- File System State
- Service / Protocol A
- Service / Protocol B
@@ -249,13 +249,13 @@ The observation space is a 6 x 6 Box type (OpenAI Gym Space) in this example. Th
For the nodes, the following values are represented:
* ID
* Operating State:
* Hardware State:
* 1 = ON
* 2 = OFF
* 3 = RESETTING
* O/S State:
* SoftwareState:
* 1 = GOOD
* 2 = PATCHING
@@ -281,8 +281,8 @@ For the nodes, the following values are represented:
For the links, the following statuses are represented:
* ID
* Operating State = N/A
* O/S State = N/A
* Hardware State = N/A
* SoftwareState = N/A
* Protocol = loading in bits/s
Action Spaces
@@ -300,7 +300,7 @@ The choice of action space used during a training session is determined in the c
The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym multidiscrete type, as follows:
* [0, num nodes] - Node ID (0 = nothing, node ID)
* [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = O/S state, 3 = service state, 4 = file system state)
* [0, 4] - What property it's acting on (0 = nothing, 1 = state, 2 = SoftwareState, 3 = service state, 4 = file system state)
* [0, 3] - Action on property (0 = nothing, 1 = on / scan, 2 = off / repair, 3 = reset / patch / restore)
* [0, num services] - Resolves to service ID (0 = nothing, resolves to service)

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)
* **Node Operating State [offShouldBeOn]** [int]
* **Node Hardware State [offShouldBeOn]** [int]
The score to give when the node should be on, but is off
* **Node Operating State [offShouldBeResetting]** [int]
* **Node Hardware State [offShouldBeResetting]** [int]
The score to give when the node should be resetting, but is off
* **Node Operating State [onShouldBeOff]** [int]
* **Node Hardware State [onShouldBeOff]** [int]
The score to give when the node should be off, but is on
* **Node Operating State [onShouldBeResetting]** [int]
* **Node Hardware State [onShouldBeResetting]** [int]
The score to give when the node should be resetting, but is on
* **Node Operating State [resettingShouldBeOn]** [int]
* **Node Hardware State [resettingShouldBeOn]** [int]
The score to give when the node should be on, but is resetting
* **Node Operating State [resettingShouldBeOff]** [int]
* **Node Hardware State [resettingShouldBeOff]** [int]
The score to give when the node should be off, but is resetting
* **Node Operating State [resetting]** [int]
* **Node Hardware State [resetting]** [int]
The score to give when the node is resetting
@@ -261,7 +261,7 @@ The config_main.yaml file consists of the following attributes:
* **nodeResetDuration** [int]
The number of steps to take when resetting a node's operating state
The number of steps to take when resetting a node's hardware state
* **servicePatchingDuration** [int]
@@ -306,13 +306,13 @@ The config_[name].yaml file consists of the following attributes:
* **id** [int]: Unique ID for this YAML item
* **name** [freetext]: Human-readable name of the component
* **baseType** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services
* **nodeType** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER
* **node_class** [enum]: Relates to the base type of the node. Can be SERVICE, ACTIVE or PASSIVE. PASSIVE nodes do not have an operating system or services. ACTIVE nodes have an operating system, but no services. SERVICE nodes have both an operating system and one or more services
* **node_type** [enum]: Relates to the component type. Can be one of CCTV, SWITCH, COMPUTER, LINK, MONITOR, PRINTER, LOP, RTU, ACTUATOR or SERVER
* **priority** [enum]: Provides a priority for each node. Can be one of P1, P2, P3, P4 or P5 (which P1 being the highest)
* **hardwareState** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING
* **ipAddress** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx
* **softwareState** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED
* **fileSystemState** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING
* **hardware_state** [enum]: The initial hardware state of the node. Can be one of ON, OFF or RESETTING
* **ip_address** [IP address]: The IP address of the component in format xxx.xxx.xxx.xxx
* **software_state** [enum]: The intial state of the node operating system. Can be GOOD, PATCHING or COMPROMISED
* **file_system_state** [enum]: The initial state of the node file system. Can be GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING
* **services**: For each service associated with the node:
* **name** [freetext]: Free-text name of the service, but must match one of the services defined for the system in the services list
@@ -367,7 +367,7 @@ The config_[name].yaml file consists of the following attributes:
* **nodeId** [int]: The ID of the node to apply the PoL to
* **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE
* **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list
* **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state)
* **state** [enuum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state)
* **itemType: RED_POL**
@@ -380,7 +380,7 @@ The config_[name].yaml file consists of the following attributes:
* **initiator** [enum]: What initiates the PoL. Can be DIRECT, IER or SERVICE
* **type** [enum]: The type of PoL to apply. Can be one of OPERATING, OS or SERVICE
* **protocol** [freetext]: The protocol to be affected if SERVICE type is chosen. Must match a value in the services list
* **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for operating system state) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state)
* **state** [enum]: The state to apply to the node (which represents the PoL change). Can be one of ON, OFF or RESETTING (for node state) or GOOD, PATCHING or COMPROMISED (for Software State) or GOOD, PATCHING, COMPROMISED or OVERWHELMED (for service state) or GOOD, CORRUPT, DESTROYED, REPAIRING or RESTORING (for file system state)
* **sourceNodeId** [int] The ID of the source node containing the service to check (used for SERVICE initiator)
* **sourceNodeService** [freetext]: The service on the source node to check (used for SERVICE initiator). Must match a value in the services list for this node
* **sourceNodeServiceState** [enum]: The state of the source node service to check (used for SERVICE initiator). Can be one of GOOD, PATCHING, COMPROMISED or OVERWHELMED

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

@@ -1,5 +1,6 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""A class that implements the access control list implementation for the network."""
from typing import Dict
from primaite.acl.acl_rule import ACLRule
@@ -9,7 +10,7 @@ class AccessControlList:
def __init__(self):
"""Init."""
self.acl = {} # A dictionary of ACL Rules
self.acl: Dict[str, AccessControlList] = {} # A dictionary of ACL Rules
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
"""

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

@@ -2,7 +2,7 @@
"""The config class."""
class config_values_main(object):
class ConfigValuesMain(object):
"""Class to hold main config values."""
def __init__(self):
@@ -23,7 +23,7 @@ class config_values_main(object):
# Reward values
# Generic
self.all_ok = 0
# Node Operating State
# Node Hardware State
self.off_should_be_on = 0
self.off_should_be_resetting = 0
self.on_should_be_off = 0
@@ -31,7 +31,7 @@ class config_values_main(object):
self.resetting_should_be_on = 0
self.resetting_should_be_off = 0
self.resetting = 0
# Node O/S or Service State
# Node Software or Service State
self.good_should_be_patching = 0
self.good_should_be_compromised = 0
self.good_should_be_overwhelmed = 0

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

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

@@ -1,83 +1,28 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The Service class."""
from primaite.common.enums import SOFTWARE_STATE
from primaite.common.enums import SoftwareState
class Service(object):
"""Service class."""
def __init__(self, _name, _port, _state):
def __init__(self, name: str, port: str, software_state: SoftwareState):
"""
Init.
Args:
_name: The service name
_port: The service port
_state: The service state
:param name: The service name.
:param port: The service port.
:param software_state: The service SoftwareState.
"""
self.name = _name
self.port = _port
self.state = _state
self.name = name
self.port = port
self.software_state = software_state
self.patching_count = 0
def set_name(self, _name):
"""
Sets the service name.
Args:
_name: The service name
"""
self.name = _name
def get_name(self):
"""
Gets the service name.
Returns:
The service name
"""
return self.name
def set_port(self, _port):
"""
Sets the service port.
Args:
_port: The service port
"""
self.port = _port
def get_port(self):
"""
Gets the service port.
Returns:
The service port
"""
return self.port
def set_state(self, _state):
"""
Sets the service state.
Args:
_state: The service state
"""
self.state = _state
def get_state(self):
"""
Gets the service state.
Returns:
The service state
"""
return self.state
def reduce_patching_count(self):
"""Reduces the patching count for the service."""
self.patching_count -= 1
if self.patching_count <= 0:
self.patching_count = 0
self.state = SOFTWARE_STATE.GOOD
self.software_state = SoftwareState.GOOD

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,6 +1,9 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements reward function."""
from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE
from typing import Dict
from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
from primaite.common.service import Service
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode
@@ -28,17 +31,17 @@ def calculate_reward_function(
"""
reward_value = 0
# For each node, compare operating state, o/s operating state, service states
# For each node, compare hardware state, SoftwareState, service states
for node_key, final_node in final_nodes.items():
initial_node = initial_nodes[node_key]
reference_node = reference_nodes[node_key]
# Operating State
# Hardware State
reward_value += score_node_operating_state(
final_node, initial_node, reference_node, config_values
)
# Operating System State
# Software State
if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode):
reward_value += score_node_os_state(
final_node, initial_node, reference_node, config_values
@@ -80,7 +83,7 @@ def calculate_reward_function(
def score_node_operating_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating state of a node.
Calculates score relating to the hardware state of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -89,9 +92,9 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
config_values: Config values
"""
score = 0
final_node_operating_state = final_node.get_state()
initial_node_operating_state = initial_node.get_state()
reference_node_operating_state = reference_node.get_state()
final_node_operating_state = final_node.hardware_state
initial_node_operating_state = initial_node.hardware_state
reference_node_operating_state = reference_node.hardware_state
if final_node_operating_state == reference_node_operating_state:
# All is well - we're no different from the reference situation
@@ -99,26 +102,26 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
else:
# We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_operating_state == HARDWARE_STATE.ON:
if reference_node_operating_state == HARDWARE_STATE.OFF:
if initial_node_operating_state == HardwareState.ON:
if reference_node_operating_state == HardwareState.OFF:
score += config_values.off_should_be_on
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting_should_be_on
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.OFF:
if reference_node_operating_state == HARDWARE_STATE.ON:
elif initial_node_operating_state == HardwareState.OFF:
if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_off
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting_should_be_off
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
if reference_node_operating_state == HARDWARE_STATE.ON:
elif initial_node_operating_state == HardwareState.RESETTING:
if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_resetting
elif reference_node_operating_state == HARDWARE_STATE.OFF:
elif reference_node_operating_state == HardwareState.OFF:
score += config_values.off_should_be_resetting
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting
else:
pass
@@ -130,7 +133,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
def score_node_os_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating system state of a node.
Calculates score relating to the Software State of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -139,9 +142,9 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
config_values: Config values
"""
score = 0
final_node_os_state = final_node.get_os_state()
initial_node_os_state = initial_node.get_os_state()
reference_node_os_state = reference_node.get_os_state()
final_node_os_state = final_node.software_state
initial_node_os_state = initial_node.software_state
reference_node_os_state = reference_node.software_state
if final_node_os_state == reference_node_os_state:
# All is well - we're no different from the reference situation
@@ -149,28 +152,28 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
else:
# We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_os_state == SOFTWARE_STATE.GOOD:
if reference_node_os_state == SOFTWARE_STATE.PATCHING:
if initial_node_os_state == SoftwareState.GOOD:
if reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_good
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
elif initial_node_os_state == SoftwareState.PATCHING:
if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_patching
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
elif initial_node_os_state == SoftwareState.COMPROMISED:
if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_compromised
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised
else:
pass
@@ -191,59 +194,59 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
config_values: Config values
"""
score = 0
final_node_services = final_node.get_services()
initial_node_services = initial_node.get_services()
reference_node_services = reference_node.get_services()
final_node_services: Dict[str, Service] = final_node.services
initial_node_services: Dict[str, Service] = initial_node.services
reference_node_services: Dict[str, Service] = reference_node.services
for service_key, final_service in final_node_services.items():
reference_service = reference_node_services[service_key]
initial_service = initial_node_services[service_key]
if final_service.get_state() == reference_service.get_state():
if final_service.software_state == reference_service.software_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step)
if initial_service.get_state() == SOFTWARE_STATE.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.PATCHING:
if initial_service.software_state == SoftwareState.GOOD:
if reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_good
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_good
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.PATCHING:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.COMPROMISED:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_compromised
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.OVERWHELMED:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed
else:
pass
@@ -263,12 +266,12 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
reference_node: The node if there had been no red or blue effect
"""
score = 0
final_node_file_system_state = final_node.get_file_system_state_actual()
initial_node_file_system_state = initial_node.get_file_system_state_actual()
reference_node_file_system_state = reference_node.get_file_system_state_actual()
final_node_file_system_state = final_node.file_system_state_actual
initial_node_file_system_state = initial_node.file_system_state_actual
reference_node_file_system_state = reference_node.file_system_state_actual
final_node_scanning_state = final_node.is_scanning_file_system()
reference_node_scanning_state = reference_node.is_scanning_file_system()
final_node_scanning_state = final_node.file_system_scanning
reference_node_scanning_state = reference_node.file_system_scanning
# File System State
if final_node_file_system_state == reference_node_file_system_state:
@@ -277,66 +280,66 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
else:
# We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step)
if initial_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if initial_node_file_system_state == FileSystemState.GOOD:
if reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_good
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.REPAIRING:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.RESTORING:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.CORRUPT:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.DESTROYED:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed
else:
pass

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

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

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.ppo import MlpPolicy as PPOMlp
from primaite.common.config_values_main import config_values_main
from primaite.common.config_values_main import ConfigValuesMain
from primaite.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import write_transaction_to_file
@@ -178,7 +178,7 @@ def load_config_values():
# Reward values
# Generic
config_values.all_ok = int(config_data["allOk"])
# Node Operating State
# Node Hardware State
config_values.off_should_be_on = int(config_data["offShouldBeOn"])
config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
config_values.on_should_be_off = int(config_data["onShouldBeOff"])
@@ -186,7 +186,7 @@ def load_config_values():
config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"])
config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
config_values.resetting = int(config_data["resetting"])
# Node O/S or Service State
# Node Software or Service State
config_values.good_should_be_patching = int(config_data["goodShouldBePatching"])
config_values.good_should_be_compromised = int(
config_data["goodShouldBeCompromised"]
@@ -331,7 +331,7 @@ try:
config_file_main = open("config/config_main.yaml", "r")
config_data = yaml.safe_load(config_file_main)
# Create a config class
config_values = config_values_main()
config_values = ConfigValuesMain()
# Load in config data
load_config_values()
except Exception:

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

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

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

@@ -1,135 +1,56 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The base Node class."""
from primaite.common.enums import HARDWARE_STATE
from typing import Final
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import HardwareState, NodeType, Priority
class Node:
"""Node class."""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
def __init__(
self,
node_id: str,
name: str,
node_type: NodeType,
priority: Priority,
hardware_state: HardwareState,
config_values: ConfigValuesMain,
):
"""
Init.
Args:
_id: The node id
_name: The name of the node
_type: The type of the node
_priority: The priority of the node
_state: The state of the node
:param node_id: The node id.
:param name: The name of the node.
:param node_type: The type of the node.
:param priority: The priority of the node.
:param hardware_state: The state of the node.
:param config_values: Config values.
"""
self.id = _id
self.name = _name
self.type = _type
self.priority = _priority
self.operating_state = _state
self.resetting_count = 0
self.config_values = _config_values
self.node_id: Final[str] = node_id
self.name: Final[str] = name
self.node_type: Final[NodeType] = node_type
self.priority = priority
self.hardware_state: HardwareState = hardware_state
self.resetting_count: int = 0
self.config_values: ConfigValuesMain = config_values
def __repr__(self):
"""Returns the name of the node."""
return self.name
def set_id(self, _id):
"""
Sets the node ID.
Args:
_id: The node ID
"""
self.id = _id
def get_id(self):
"""
Gets the node ID.
Returns:
The node ID
"""
return self.id
def set_name(self, _name):
"""
Sets the node name.
Args:
_name: The node name
"""
self.name = _name
def get_name(self):
"""
Gets the node name.
Returns:
The node name
"""
return self.name
def set_type(self, _type):
"""
Sets the node type.
Args:
_type: The node type
"""
self.type = _type
def get_type(self):
"""
Gets the node type.
Returns:
The node type
"""
return self.type
def set_priority(self, _priority):
"""
Sets the node priority.
Args:
_priority: The node priority
"""
self.priority = _priority
def get_priority(self):
"""
Gets the node priority.
Returns:
The node priority
"""
return self.priority
def set_state(self, _state):
"""
Sets the node state.
Args:
_state: The node state
"""
self.operating_state = _state
def get_state(self):
"""
Gets the node operating state.
Returns:
The node operating state
"""
return self.operating_state
def turn_on(self):
"""Sets the node state to ON."""
self.operating_state = HARDWARE_STATE.ON
self.hardware_state = HardwareState.ON
def turn_off(self):
"""Sets the node state to OFF."""
self.operating_state = HARDWARE_STATE.OFF
self.hardware_state = HardwareState.OFF
def reset(self):
"""Sets the node state to Resetting and starts the reset count."""
self.operating_state = HARDWARE_STATE.RESETTING
self.hardware_state = HardwareState.RESETTING
self.resetting_count = self.config_values.node_reset_duration
def update_resetting_status(self):
@@ -137,4 +58,4 @@ class Node:
self.resetting_count -= 1
if self.resetting_count <= 0:
self.resetting_count = 0
self.operating_state = HARDWARE_STATE.ON
self.hardware_state = HardwareState.ON

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

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

@@ -1,32 +1,44 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""The Passive Node class (i.e. an actuator)."""
from primaite.common.config_values_main import ConfigValuesMain
from primaite.common.enums import HardwareState, NodeType, Priority
from primaite.nodes.node import Node
class PassiveNode(Node):
"""The Passive Node class."""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
def __init__(
self,
node_id: str,
name: str,
node_type: NodeType,
priority: Priority,
hardware_state: HardwareState,
config_values: ConfigValuesMain,
):
"""
Init.
Args:
_id: The node id
_name: The name of the node
_type: The type of the node
_priority: The priority of the node
_state: The state of the node
:param node_id: The node id.
:param name: The name of the node.
:param node_type: The type of the node.
:param priority: The priority of the node.
:param hardware_state: The state of the node.
:param config_values: Config values.
"""
# Pass through to Super for now
super().__init__(_id, _name, _type, _priority, _state, _config_values)
super().__init__(
node_id, name, node_type, priority, hardware_state, config_values
)
def get_ip_address(self):
@property
def ip_address(self) -> str:
"""
Gets the node IP address.
Gets the node IP address as an empty string.
Returns:
The node IP address
No concept of IP address for passive nodes for now.
:return: The node IP address.
"""
# No concept of IP address for passive nodes for now
return ""

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

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

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

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

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

@@ -1,22 +1,35 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements POL on the network (nodes and links) resulting from the red agent attack."""
from typing import Dict
from networkx import shortest_path
from networkx import MultiGraph, shortest_path
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import (
HARDWARE_STATE,
NODE_POL_INITIATOR,
NODE_POL_TYPE,
SOFTWARE_STATE,
TYPE,
HardwareState,
NodePOLInitiator,
NodePOLType,
NodeType,
SoftwareState,
)
from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from primaite.nodes.service_node import ServiceNode
from primaite.pol.ier import IER
_VERBOSE = False
def apply_red_agent_iers(network, nodes, links, iers, acl, step):
def apply_red_agent_iers(
network: MultiGraph,
nodes: Dict[str, NodeUnion],
links: Dict[str, Link],
iers: Dict[str, IER],
acl: AccessControlList,
step: int,
):
"""
Applies IERs to the links (link POL) resulting from red agent attack.
@@ -54,25 +67,25 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
dest_node = nodes[dest_node_id]
# 1. Check the source node situation
if source_node.get_type() == TYPE.SWITCH:
if source_node.node_type == NodeType.SWITCH:
# It's a switch
if source_node.get_state() == HARDWARE_STATE.ON:
if source_node.hardware_state == HardwareState.ON:
source_valid = True
else:
# IER no longer valid
source_valid = False
elif source_node.get_type() == TYPE.ACTUATOR:
elif source_node.node_type == NodeType.ACTUATOR:
# It's an actuator
# TO DO
pass
else:
# It's not a switch or an actuator (so active node)
if source_node.get_state() == HARDWARE_STATE.ON:
if source_node.hardware_state == HardwareState.ON:
if source_node.has_service(protocol):
# Red agents IERs can only be valid if the source service is in a compromised state
if (
source_node.get_service_state(protocol)
== SOFTWARE_STATE.COMPROMISED
== SoftwareState.COMPROMISED
):
source_valid = True
else:
@@ -86,19 +99,19 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
source_valid = False
# 2. Check the dest node situation
if dest_node.get_type() == TYPE.SWITCH:
if dest_node.node_type == NodeType.SWITCH:
# It's a switch
if dest_node.get_state() == HARDWARE_STATE.ON:
if dest_node.hardware_state == HardwareState.ON:
dest_valid = True
else:
# IER no longer valid
dest_valid = False
elif dest_node.get_type() == TYPE.ACTUATOR:
elif dest_node.node_type == NodeType.ACTUATOR:
# It's an actuator
pass
else:
# It's not a switch or an actuator (so active node)
if dest_node.get_state() == HARDWARE_STATE.ON:
if dest_node.hardware_state == HardwareState.ON:
if dest_node.has_service(protocol):
# We don't care what state the destination service is in for an IER
dest_valid = True
@@ -112,15 +125,15 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
# 3. Check that the ACL doesn't block it
acl_block = acl.is_blocked(
source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port
source_node.ip_address, dest_node.ip_address, protocol, port
)
if acl_block:
if _VERBOSE:
print(
"ACL block on source: "
+ source_node.get_ip_address()
+ source_node.ip_address
+ ", dest: "
+ dest_node.get_ip_address()
+ dest_node.ip_address
+ ", protocol: "
+ protocol
+ ", port: "
@@ -145,7 +158,7 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
# We might have a switch in the path, so check all nodes are operational
# We're assuming here that red agents can get past switches that are patching
for node in path_node_list:
if node.get_state() != HARDWARE_STATE.ON:
if node.hardware_state != HardwareState.ON:
path_valid = False
if path_valid:
@@ -207,7 +220,12 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
pass
def apply_red_agent_node_pol(nodes, iers, node_pol, step):
def apply_red_agent_node_pol(
nodes: Dict[str, NodeUnion],
iers: Dict[str, IER],
node_pol: Dict[str, NodeStateInstructionRed],
step: int,
):
"""
Applies node pattern of life.
@@ -238,22 +256,22 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
if step >= start_step and step <= stop_step:
# continue --------------------------
target_node = nodes[target_node_id]
target_node: NodeUnion = nodes[target_node_id]
# Based the action taken on the initiator type
if initiator == NODE_POL_INITIATOR.DIRECT:
if initiator == NodePOLInitiator.DIRECT:
# No conditions required, just apply the change
passed_checks = True
elif initiator == NODE_POL_INITIATOR.IER:
elif initiator == NodePOLInitiator.IER:
# Need to check there is a red IER incoming
passed_checks = is_red_ier_incoming(target_node, iers, pol_type)
elif initiator == NODE_POL_INITIATOR.SERVICE:
elif initiator == NodePOLInitiator.SERVICE:
# Need to check the condition of a service on another node
source_node = nodes[source_node_id]
if source_node.has_service(source_node_service_name):
if (
source_node.get_service_state(source_node_service_name)
== SOFTWARE_STATE[source_node_service_state_value]
== SoftwareState[source_node_service_state_value]
):
passed_checks = True
else:
@@ -269,16 +287,16 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
# Only apply the PoL if the checks have passed (based on the initiator type)
if passed_checks:
# Apply the change
if pol_type == NODE_POL_TYPE.OPERATING:
# Change operating state
target_node.set_state(state)
elif pol_type == NODE_POL_TYPE.OS:
if pol_type == NodePOLType.OPERATING:
# Change hardware state
target_node.hardware_state = state
elif pol_type == NodePOLType.OS:
# Change OS state
if isinstance(target_node, ActiveNode) or isinstance(
target_node, ServiceNode
):
target_node.set_os_state(state)
elif pol_type == NODE_POL_TYPE.SERVICE:
target_node.software_state = state
elif pol_type == NodePOLType.SERVICE:
# Change a service state
if isinstance(target_node, ServiceNode):
target_node.set_service_state(service_name, state)
@@ -302,18 +320,18 @@ def is_red_ier_incoming(node, iers, node_pol_type):
TODO: Write more descriptive docstring with params and returns.
"""
node_id = node.get_id()
node_id = node.node_id
for ier_key, ier_value in iers.items():
if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id:
if (
node_pol_type == NODE_POL_TYPE.OPERATING
or node_pol_type == NODE_POL_TYPE.OS
or node_pol_type == NODE_POL_TYPE.FILE
node_pol_type == NodePOLType.OPERATING
or node_pol_type == NodePOLType.OS
or node_pol_type == NodePOLType.FILE
):
# It's looking to change operating state, file system or O/S state, so valid
# It's looking to change hardware state, file system or SoftwareState, so valid
return True
elif node_pol_type == NODE_POL_TYPE.SERVICE:
elif node_pol_type == NodePOLType.SERVICE:
# Check if the service is present on the node and running
ier_protocol = ier_value.get_protocol()
if isinstance(node, ServiceNode):

14
tests/config/one_node_states_on_off_lay_down_config.yaml
View File

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

4
tests/config/one_node_states_on_off_main_config.yaml
View File

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

8
tests/conftest.py
View File

@@ -5,7 +5,7 @@ from typing import Union
import yaml
from primaite.common.config_values_main import config_values_main
from primaite.common.config_values_main import ConfigValuesMain
from primaite.environment.primaite_env import Primaite
ACTION_SPACE_NODE_VALUES = 1
@@ -32,7 +32,7 @@ def _get_primaite_env_from_config(
# Reward values
# Generic
config_values.all_ok = int(config_data["allOk"])
# Node Operating State
# Node Hardware State
config_values.off_should_be_on = int(config_data["offShouldBeOn"])
config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
config_values.on_should_be_off = int(config_data["onShouldBeOff"])
@@ -40,7 +40,7 @@ def _get_primaite_env_from_config(
config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"])
config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
config_values.resetting = int(config_data["resetting"])
# Node O/S or Service State
# Node Software or Service State
config_values.good_should_be_patching = int(config_data["goodShouldBePatching"])
config_values.good_should_be_compromised = int(
config_data["goodShouldBeCompromised"]
@@ -160,7 +160,7 @@ def _get_primaite_env_from_config(
config_file_main = open(main_config_path, "r")
config_data = yaml.safe_load(config_file_main)
# Create a config class
config_values = config_values_main()
config_values = ConfigValuesMain()
# Load in config data
load_config_values()
env = Primaite(config_values, [])

54
tests/test_active_node.py
View File

@@ -1,22 +1,22 @@
"""Used to test Active Node functions."""
import pytest
from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE
from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
from primaite.nodes.active_node import ActiveNode
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED),
(HardwareState.OFF, SoftwareState.GOOD),
(HardwareState.ON, SoftwareState.OVERWHELMED),
],
)
def test_os_state_change(operating_state, expected_state):
"""
Test that a node cannot change its operating system state.
Test that a node cannot change its Software State.
When its operating state is OFF.
When its hardware state is OFF.
"""
active_node = ActiveNode(
0,
@@ -25,28 +25,28 @@ def test_os_state_change(operating_state, expected_state):
"1",
operating_state,
"192.168.0.1",
SOFTWARE_STATE.GOOD,
SoftwareState.GOOD,
"GOOD",
1,
)
active_node.set_os_state(SOFTWARE_STATE.OVERWHELMED)
active_node.software_state = SoftwareState.OVERWHELMED
assert active_node.get_os_state() == expected_state
assert active_node.software_state == expected_state
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED),
(HardwareState.OFF, SoftwareState.GOOD),
(HardwareState.ON, SoftwareState.OVERWHELMED),
],
)
def test_os_state_change_if_not_compromised(operating_state, expected_state):
"""
Test that a node cannot change its operating system state.
Test that a node cannot change its Software State.
If not compromised) when its operating state is OFF.
If not compromised) when its hardware state is OFF.
"""
active_node = ActiveNode(
0,
@@ -55,25 +55,25 @@ def test_os_state_change_if_not_compromised(operating_state, expected_state):
"1",
operating_state,
"192.168.0.1",
SOFTWARE_STATE.GOOD,
SoftwareState.GOOD,
"GOOD",
1,
)
active_node.set_os_state_if_not_compromised(SOFTWARE_STATE.OVERWHELMED)
active_node.set_software_state_if_not_compromised(SoftwareState.OVERWHELMED)
assert active_node.get_os_state() == expected_state
assert active_node.software_state == expected_state
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD),
(HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT),
(HardwareState.OFF, FileSystemState.GOOD),
(HardwareState.ON, FileSystemState.CORRUPT),
],
)
def test_file_system_change(operating_state, expected_state):
"""Test that a node cannot change its file system state when its operating state is ON."""
"""Test that a node cannot change its file system state when its hardware state is ON."""
active_node = ActiveNode(
0,
"node",
@@ -82,27 +82,27 @@ def test_file_system_change(operating_state, expected_state):
operating_state,
"192.168.0.1",
"COMPROMISED",
FILE_SYSTEM_STATE.GOOD,
FileSystemState.GOOD,
1,
)
active_node.set_file_system_state(FILE_SYSTEM_STATE.CORRUPT)
active_node.set_file_system_state(FileSystemState.CORRUPT)
assert active_node.get_file_system_state_actual() == expected_state
assert active_node.file_system_state_actual == expected_state
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, FILE_SYSTEM_STATE.GOOD),
(HARDWARE_STATE.ON, FILE_SYSTEM_STATE.CORRUPT),
(HardwareState.OFF, FileSystemState.GOOD),
(HardwareState.ON, FileSystemState.CORRUPT),
],
)
def test_file_system_change_if_not_compromised(operating_state, expected_state):
"""
Test that a node cannot change its file system state.
If not compromised) when its operating state is OFF.
If not compromised) when its hardware state is OFF.
"""
active_node = ActiveNode(
0,
@@ -112,10 +112,10 @@ def test_file_system_change_if_not_compromised(operating_state, expected_state):
operating_state,
"192.168.0.1",
"GOOD",
FILE_SYSTEM_STATE.GOOD,
FileSystemState.GOOD,
1,
)
active_node.set_file_system_state_if_not_compromised(FILE_SYSTEM_STATE.CORRUPT)
active_node.set_file_system_state_if_not_compromised(FileSystemState.CORRUPT)
assert active_node.get_file_system_state_actual() == expected_state
assert active_node.file_system_state_actual == expected_state

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)
Hardware State: onShouldBeOff = -2 (between Steps 4 & 6)
Service State: goodShouldBeCompromised = 5 (between Steps 7 & 9)
Operating System State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12)
Software State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12)
Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26
Step Count: 13

22
tests/test_service_node.py
View File

@@ -1,7 +1,7 @@
"""Used to test Service Node functions."""
import pytest
from primaite.common.enums import HARDWARE_STATE, SOFTWARE_STATE
from primaite.common.enums import HardwareState, SoftwareState
from primaite.common.service import Service
from primaite.nodes.service_node import ServiceNode
@@ -9,15 +9,15 @@ from primaite.nodes.service_node import ServiceNode
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED),
(HardwareState.OFF, SoftwareState.GOOD),
(HardwareState.ON, SoftwareState.OVERWHELMED),
],
)
def test_service_state_change(operating_state, expected_state):
"""
Test that a node cannot change the state of a running service.
When its operating state is OFF.
When its hardware state is OFF.
"""
service_node = ServiceNode(
0,
@@ -30,10 +30,10 @@ def test_service_state_change(operating_state, expected_state):
"RESTORING",
1,
)
service = Service("TCP", 80, SOFTWARE_STATE.GOOD)
service = Service("TCP", 80, SoftwareState.GOOD)
service_node.add_service(service)
service_node.set_service_state("TCP", SOFTWARE_STATE.OVERWHELMED)
service_node.set_service_state("TCP", SoftwareState.OVERWHELMED)
assert service_node.get_service_state("TCP") == expected_state
@@ -41,15 +41,15 @@ def test_service_state_change(operating_state, expected_state):
@pytest.mark.parametrize(
"operating_state, expected_state",
[
(HARDWARE_STATE.OFF, SOFTWARE_STATE.GOOD),
(HARDWARE_STATE.ON, SOFTWARE_STATE.OVERWHELMED),
(HardwareState.OFF, SoftwareState.GOOD),
(HardwareState.ON, SoftwareState.OVERWHELMED),
],
)
def test_service_state_change_if_not_comprised(operating_state, expected_state):
"""
Test that a node cannot change the state of a running service.
If not compromised when its operating state is ON.
If not compromised when its hardware state is ON.
"""
service_node = ServiceNode(
0,
@@ -62,9 +62,9 @@ def test_service_state_change_if_not_comprised(operating_state, expected_state):
"RESTORING",
1,
)
service = Service("TCP", 80, SOFTWARE_STATE.GOOD)
service = Service("TCP", 80, SoftwareState.GOOD)
service_node.add_service(service)
service_node.set_service_state_if_not_compromised("TCP", SOFTWARE_STATE.OVERWHELMED)
service_node.set_service_state_if_not_compromised("TCP", SoftwareState.OVERWHELMED)
assert service_node.get_service_state("TCP") == expected_state