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

Ran pre-commit hook on all files and performed changes to fix flake8 failures

Browse Source
This commit is contained in:
Chris McCarthy
2023-05-25 11:42:19 +01:00
parent aa8284897a
commit 4eb1658966
42 changed files with 1371 additions and 1170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/primaite/common/__init__.py
View File

@@ -1,2 +1 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

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

@@ -1,39 +1,35 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The config class
"""
"""The config class."""
class config_values_main(object):
"""
Class to hold main config values
"""
"""Class to hold main config values."""
def __init__(self):
"""
Init
"""
"""Init."""
# Generic
self.agent_identifier = "" # the agent in use
self.num_episodes = 0 # number of episodes to train over
self.num_steps = 0 # number of steps in an episode
self.time_delay = 0 # delay between steps (ms) - applies to generic agents only
self.config_filename_use_case = "" # the filename for the Use Case config file
self.session_type = "" # the session type to run (TRAINING or EVALUATION)
self.agent_identifier = "" # the agent in use
self.num_episodes = 0 # number of episodes to train over
self.num_steps = 0 # number of steps in an episode
self.time_delay = 0 # delay between steps (ms) - applies to generic agents only
self.config_filename_use_case = "" # the filename for the Use Case config file
self.session_type = "" # the session type to run (TRAINING or EVALUATION)
# Environment
self.observation_space_high_value = 0 # The high value for the observation space
self.observation_space_high_value = (
0 # The high value for the observation space
)
# Reward values
# Generic
self.all_ok = 0
self.all_ok = 0
# Node Operating State
self.off_should_be_on = 0
self.off_should_be_resetting = 0
self.on_should_be_off = 0
self.on_should_be_resetting = 0
self.resetting_should_be_on = 0
self.resetting_should_be_off = 0
self.resetting_should_be_off = 0
self.resetting = 0
# Node O/S or Service State
self.good_should_be_patching = 0
@@ -46,7 +42,7 @@ class config_values_main(object):
self.compromised_should_be_good = 0
self.compromised_should_be_patching = 0
self.compromised_should_be_overwhelmed = 0
self.compromised = 0
self.compromised = 0
self.overwhelmed_should_be_good = 0
self.overwhelmed_should_be_patching = 0
self.overwhelmed_should_be_compromised = 0
@@ -59,11 +55,15 @@ class config_values_main(object):
self.repairing_should_be_good = 0
self.repairing_should_be_restoring = 0
self.repairing_should_be_corrupt = 0
self.repairing_should_be_destroyed = 0 # Repairing does not fix destroyed state - you need to restore
self.repairing_should_be_destroyed = (
0 # Repairing does not fix destroyed state - you need to restore
)
self.repairing = 0
self.restoring_should_be_good = 0
self.restoring_should_be_repairing = 0
self.restoring_should_be_corrupt = 0 # Not the optimal method (as repair will fix corruption)
self.restoring_should_be_corrupt = (
0 # Not the optimal method (as repair will fix corruption)
)
self.restoring_should_be_destroyed = 0
self.restoring = 0
self.corrupt_should_be_good = 0
@@ -82,10 +82,9 @@ class config_values_main(object):
self.green_ier_blocked = 0
# Patching / Reset
self.os_patching_duration = 0 # The time taken to patch the OS
self.node_reset_duration = 0 # The time taken to reset a node (hardware)
self.service_patching_duration = 0 # The time taken to patch a service
self.file_system_repairing_limit = 0 # The time take to repair a file
self.file_system_restoring_limit = 0 # The time take to restore a file
self.file_system_scanning_limit = 0 # The time taken to scan the file system
self.os_patching_duration = 0 # The time taken to patch the OS
self.node_reset_duration = 0 # The time taken to reset a node (hardware)
self.service_patching_duration = 0 # The time taken to patch a service
self.file_system_repairing_limit = 0 # The time take to repair a file
self.file_system_restoring_limit = 0 # The time take to restore a file
self.file_system_scanning_limit = 0 # The time taken to scan the file system

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

@@ -1,14 +1,11 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Enumerations for APE
"""
"""Enumerations for APE."""
from enum import Enum
class TYPE(Enum):
"""
Node type enumeration
"""
"""Node type enumeration."""
CCTV = 1
SWITCH = 2
@@ -21,10 +18,9 @@ class TYPE(Enum):
ACTUATOR = 9
SERVER = 10
class PRIORITY(Enum):
"""
Node priority enumeration
"""
"""Node priority enumeration."""
P1 = 1
P2 = 2
@@ -32,48 +28,43 @@ class PRIORITY(Enum):
P4 = 4
P5 = 5
class HARDWARE_STATE(Enum):
"""
Node hardware state enumeration
"""
"""Node hardware state enumeration."""
ON = 1
OFF = 2
RESETTING = 3
class SOFTWARE_STATE(Enum):
"""
O/S or Service state enumeration
"""
"""O/S or Service state enumeration."""
GOOD = 1
PATCHING = 2
COMPROMISED = 3
OVERWHELMED = 4
class NODE_POL_TYPE(Enum):
"""
Node Pattern of Life type enumeration
"""
"""Node Pattern of Life type enumeration."""
OPERATING = 1
OS = 2
SERVICE = 3
FILE = 4
class NODE_POL_INITIATOR(Enum):
"""
Node Pattern of Life initiator enumeration
"""
"""Node Pattern of Life initiator enumeration."""
DIRECT = 1
IER = 2
SERVICE = 3
class PROTOCOL(Enum):
"""
Service protocol enumeration
"""
"""Service protocol enumeration."""
LDAP = 0
FTP = 1
@@ -84,18 +75,16 @@ class PROTOCOL(Enum):
TCP = 6
NONE = 7
class ACTION_TYPE(Enum):
"""
Action type enumeration
"""
"""Action type enumeration."""
NODE = 0
ACL = 1
class FILE_SYSTEM_STATE(Enum):
"""
File System State
"""
"""File System State."""
GOOD = 1
CORRUPT = 2

30
src/primaite/common/protocol.py
View File

@@ -1,59 +1,47 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The protocol class
"""
"""The protocol class."""
class Protocol(object):
"""
Protocol class
"""
"""Protocol class."""
def __init__(self, _name):
"""
Init
Init.
Args:
_name: The protocol name
"""
self.name = _name
self.load = 0 # bps
self.load = 0 # bps
def get_name(self):
"""
Gets the protocol name
Gets the protocol name.
Returns:
The protocol name
"""
return self.name
def get_load(self):
"""
Gets the protocol load
Gets the protocol load.
Returns:
The protocol load (bps)
"""
return self.load
def add_load(self, _load):
"""
Adds load to the protocol
Adds load to the protocol.
Args:
_load: The load to add
"""
self.load += _load
def clear_load(self):
"""
Clears the load on this protocol
"""
"""Clears the load on this protocol."""
self.load = 0

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

@@ -1,25 +1,21 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The Service class
"""
"""The Service class."""
from primaite.common.enums import SOFTWARE_STATE
class Service(object):
"""
Service class
"""
"""Service class."""
def __init__(self, _name, _port, _state):
"""
Init
Init.
Args:
_name: The service name
_port: The service port
_state: The service state
"""
self.name = _name
self.port = _port
self.state = _state
@@ -27,74 +23,61 @@ class Service(object):
def set_name(self, _name):
"""
Sets the service name
Sets the service name.
Args:
_name: The service name
"""
self.name = _name
def get_name(self):
"""
Gets the service name
Gets the service name.
Returns:
The service name
"""
return self.name
def set_port(self, _port):
"""
Sets the service port
Sets the service port.
Args:
_port: The service port
"""
self.port = _port
def get_port(self):
"""
Gets the service port
Gets the service port.
Returns:
The service port
"""
return self.port
def set_state(self, _state):
"""
Sets the service state
Sets the service state.
Args:
_state: The service state
"""
self.state = _state
def get_state(self):
"""
Gets the service state
Gets the service state.
Returns:
The service state
"""
return self.state
def reduce_patching_count(self):
"""
Reduces the patching count for the service
"""
"""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