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

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

@@ -1,19 +1,26 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
An Active Node (i.e. not an actuator)
"""
"""An Active Node (i.e. not an actuator)."""
from primaite.common.enums import FILE_SYSTEM_STATE, SOFTWARE_STATE
from primaite.nodes.node import Node
from primaite.common.enums import *
class ActiveNode(Node):
"""
Active Node class
"""
"""Active Node class."""
def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values):
def __init__(
self,
_id,
_name,
_type,
_priority,
_state,
_ip_address,
_os_state,
_file_system_state,
_config_values,
):
"""
Init
Init.
Args:
_id: The node ID
@@ -26,7 +33,6 @@ class ActiveNode(Node):
_file_system_state: The node file system state
_config_values: The config values
"""
super().__init__(_id, _name, _type, _priority, _state, _config_values)
self.ip_address = _ip_address
# Related to O/S
@@ -39,20 +45,18 @@ class ActiveNode(Node):
self.file_system_scanning_count = 0
self.file_system_action_count = 0
def set_ip_address(self, _ip_address):
"""
Sets IP address
Sets IP address.
Args:
_ip_address: IP address
"""
self.ip_address = _ip_address
def get_ip_address(self):
"""
Gets IP address
Gets IP address.
Returns:
IP address
@@ -61,24 +65,22 @@ class ActiveNode(Node):
def set_os_state(self, _os_state):
"""
Sets operating system state
Sets operating system state.
Args:
_os_state: Operating system state
"""
self.os_state = _os_state
if _os_state == SOFTWARE_STATE.PATCHING:
self.patching_count = self.config_values.os_patching_duration
def set_os_state_if_not_compromised(self, _os_state):
"""
Sets operating system state if the node is not compromised
Sets operating system state if the node is not compromised.
Args:
_os_state: Operating system state
"""
if self.os_state != SOFTWARE_STATE.COMPROMISED:
self.os_state = _os_state
if _os_state == SOFTWARE_STATE.PATCHING:
@@ -86,19 +88,15 @@ class ActiveNode(Node):
def get_os_state(self):
"""
Gets operating system state
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
"""
"""Updates operating system status based on patching cycle."""
self.patching_count -= 1
if self.patching_count <= 0:
self.patching_count = 0
@@ -106,87 +104,88 @@ class ActiveNode(Node):
def set_file_system_state(self, _file_system_state):
"""
Sets the file system state (actual and observed)
Sets the file system state (actual and observed).
Args:
_file_system_state: File system state
"""
self.file_system_state_actual = _file_system_state
if _file_system_state == FILE_SYSTEM_STATE.REPAIRING:
self.file_system_action_count = self.config_values.file_system_repairing_limit
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_action_count = self.config_values.file_system_restoring_limit
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
def set_file_system_state_if_not_compromised(self, _file_system_state):
"""
Sets the file system state (actual and observed) if not in a compromised state
Sets the file system state (actual and observed) if not in a compromised state.
Use for green PoL to prevent it overturning a compromised state
Args:
_file_system_state: File system state
"""
if self.file_system_state_actual != FILE_SYSTEM_STATE.CORRUPT and self.file_system_state_actual != FILE_SYSTEM_STATE.DESTROYED:
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 = _file_system_state
if _file_system_state == FILE_SYSTEM_STATE.REPAIRING:
self.file_system_action_count = self.config_values.file_system_repairing_limit
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_action_count = self.config_values.file_system_restoring_limit
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
def get_file_system_state_actual(self):
"""
Gets file system state (actual)
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)
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
"""
"""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
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
"""
"""Updates file system status based on scanning / restore / repair cycle."""
# Deprecate both the action count (for restoring or reparing) and the scanning count
self.file_system_action_count -= 1
self.file_system_scanning_count -= 1
@@ -194,7 +193,10 @@ class ActiveNode(Node):
# Reparing / Restoring updates
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:
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 = FILE_SYSTEM_STATE.GOOD
self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD

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

@@ -1,18 +1,14 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The base Node class
"""
"""The base Node class."""
from primaite.common.enums import HARDWARE_STATE
from primaite.common.enums import *
class Node:
"""
Node class
"""
"""Node class."""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
"""
Init
Init.
Args:
_id: The node id
@@ -21,156 +17,124 @@ class Node:
_priority: The priority of the node
_state: The state of the node
"""
self.id = _id
self.name = _name
self.type = _type
self.priority = _priority
self.operating_state = _state
self.resetting_count = 0
self.resetting_count = 0
self.config_values = _config_values
def __repr__(self):
"""
Returns the name of the node
"""
"""Returns the name of the node."""
return self.name
def set_id(self, _id):
"""
Sets the node ID
Sets the node ID.
Args:
_id: The node ID
"""
self.id = _id
def get_id(self):
"""
Gets the node ID
Gets the node ID.
Returns:
The node ID
"""
return self.id
def set_name(self, _name):
"""
Sets the node name
Sets the node name.
Args:
_name: The node name
"""
self.name = _name
def get_name(self):
"""
Gets the node name
Gets the node name.
Returns:
The node name
"""
return self.name
def set_type(self, _type):
"""
Sets the node type
Sets the node type.
Args:
_type: The node type
"""
self.type = _type
def get_type(self):
"""
Gets the node type
Gets the node type.
Returns:
The node type
"""
return self.type
def set_priority(self, _priority):
"""
Sets the node priority
Sets the node priority.
Args:
_priority: The node priority
"""
self.priority = _priority
def get_priority(self):
"""
Gets the node priority
Gets the node priority.
Returns:
The node priority
"""
return self.priority
def set_state(self, _state):
"""
Sets the node state
Sets the node state.
Args:
_state: The node state
"""
self.operating_state = _state
def get_state(self):
"""
Gets the node operating state
Gets the node operating state.
Returns:
The node operating state
"""
return self.operating_state
def turn_on(self):
"""
Sets the node state to ON
"""
"""Sets the node state to ON."""
self.operating_state = HARDWARE_STATE.ON
def turn_off(self):
"""
Sets the node state to OFF
"""
"""Sets the node state to OFF."""
self.operating_state = HARDWARE_STATE.OFF
def reset(self):
"""
Sets the node state to Resetting and starts the reset count
"""
"""Sets the node state to Resetting and starts the reset count."""
self.operating_state = HARDWARE_STATE.RESETTING
self.resetting_count = self.config_values.node_reset_duration
def update_resetting_status(self):
"""
Updates the resetting count
"""
"""Updates the resetting count."""
self.resetting_count -= 1
if self.resetting_count <= 0:
self.resetting_count = 0
self.operating_state = HARDWARE_STATE.ON

46
src/primaite/nodes/node_state_instruction_green.py
View File

@@ -1,16 +1,22 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Defines node behaviour for Green PoL
"""
"""Defines node behaviour for Green PoL."""
class NodeStateInstructionGreen(object):
"""
The Node State Instruction class
"""
"""The Node State Instruction class."""
def __init__(self, _id, _start_step, _end_step, _node_id, _node_pol_type, _service_name, _state):
def __init__(
self,
_id,
_start_step,
_end_step,
_node_id,
_node_pol_type,
_service_name,
_state,
):
"""
Init
Init.
Args:
_id: The node state instruction id
@@ -21,72 +27,64 @@ class NodeStateInstructionGreen(object):
_service_name: The service name
_state: The state (node or service)
"""
self.id = _id
self.start_step = _start_step
self.end_step = _end_step
self.node_id = _node_id
self.node_pol_type = _node_pol_type
self.service_name = _service_name # Not used when not a service instruction
self.state = _state
self.service_name = _service_name # Not used when not a service instruction
self.state = _state
def get_start_step(self):
"""
Gets the start step
Gets the start step.
Returns:
The start step
"""
return self.start_step
def get_end_step(self):
"""
Gets the end step
Gets the end step.
Returns:
The end step
"""
return self.end_step
def get_node_id(self):
"""
Gets the node ID
Gets the node ID.
Returns:
The node ID
"""
return self.node_id
def get_node_pol_type(self):
"""
Gets the node pattern of life type (enum)
Gets the node pattern of life type (enum).
Returns:
The node pattern of life type (enum)
"""
return self.node_pol_type
def get_service_name(self):
"""
Gets the service name
Gets the service name.
Returns:
The service name
"""
return self.service_name
def get_state(self):
"""
Gets the state (node or service)
Gets the state (node or service).
Returns:
The state (node or service)
"""
return self.state

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

@@ -1,16 +1,26 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Defines node behaviour for Green PoL
"""
"""Defines node behaviour for Green PoL."""
class NodeStateInstructionRed(object):
"""
The Node State Instruction class
"""
"""The Node State Instruction class."""
def __init__(self, _id, _start_step, _end_step, _target_node_id, _pol_initiator, _pol_type, pol_protocol, _pol_state, _pol_source_node_id, _pol_source_node_service, _pol_source_node_service_state):
def __init__(
self,
_id,
_start_step,
_end_step,
_target_node_id,
_pol_initiator,
_pol_type,
pol_protocol,
_pol_state,
_pol_source_node_id,
_pol_source_node_service,
_pol_source_node_service_state,
):
"""
Init
Init.
Args:
_id: The node state instruction id
@@ -25,14 +35,13 @@ class NodeStateInstructionRed(object):
_pol_source_node_service: The source node service (used for initiator type SERVICE)
_pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
"""
self.id = _id
self.start_step = _start_step
self.end_step = _end_step
self.target_node_id = _target_node_id
self.initiator = _pol_initiator
self.pol_type = _pol_type
self.service_name = pol_protocol # Not used when not a service instruction
self.service_name = pol_protocol # Not used when not a service instruction
self.state = _pol_state
self.source_node_id = _pol_source_node_id
self.source_node_service = _pol_source_node_service
@@ -40,101 +49,90 @@ class NodeStateInstructionRed(object):
def get_start_step(self):
"""
Gets the start step
Gets the start step.
Returns:
The start step
"""
return self.start_step
def get_end_step(self):
"""
Gets the end step
Gets the end step.
Returns:
The end step
"""
return self.end_step
def get_target_node_id(self):
"""
Gets the node ID
Gets the node ID.
Returns:
The node ID
"""
return self.target_node_id
def get_initiator(self):
"""
Gets the initiator
Gets the initiator.
Returns:
The initiator
"""
return self.initiator
def get_pol_type(self):
"""
Gets the node pattern of life type (enum)
Gets the node pattern of life type (enum).
Returns:
The node pattern of life type (enum)
"""
return self.pol_type
def get_service_name(self):
"""
Gets the service name
Gets the service name.
Returns:
The service name
"""
return self.service_name
def get_state(self):
"""
Gets the state (node or service)
Gets the state (node or service).
Returns:
The state (node or service)
"""
return self.state
def get_source_node_id(self):
"""
Gets the source node id (used for initiator type SERVICE)
Gets the source node id (used for initiator type SERVICE).
Returns:
The source node id
"""
return self.source_node_id
def get_source_node_service(self):
"""
Gets the source node service (used for initiator type SERVICE)
Gets the source node service (used for initiator type SERVICE).
Returns:
The source node service
"""
return self.source_node_service
def get_source_node_service_state(self):
"""
Gets the source node service state (used for initiator type SERVICE)
Gets the source node service state (used for initiator type SERVICE).
Returns:
The source node service state
"""
return self.source_node_service_state

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

@@ -1,18 +1,15 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
The Passive Node class (i.e. an actuator)
"""
"""The Passive Node class (i.e. an actuator)."""
from primaite.nodes.node import Node
class PassiveNode(Node):
"""
The Passive Node class
"""
"""The Passive Node class."""
def __init__(self, _id, _name, _type, _priority, _state, _config_values):
"""
Init
Init.
Args:
_id: The node id
@@ -21,17 +18,15 @@ class PassiveNode(Node):
_priority: The priority of the node
_state: The state of the node
"""
# Pass through to Super for now
super().__init__(_id, _name, _type, _priority, _state, _config_values)
def get_ip_address(self):
"""
Gets the node IP address
Gets the node IP address.
Returns:
The node IP address
"""
# No concept of IP address for passive nodes for now
return ""
return ""

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

@@ -1,19 +1,26 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
A Service Node (i.e. not an actuator)
"""
"""A Service Node (i.e. not an actuator)."""
from primaite.common.enums import SOFTWARE_STATE
from primaite.nodes.active_node import ActiveNode
from primaite.common.enums import *
class ServiceNode(ActiveNode):
"""
ServiceNode class
"""
"""ServiceNode class."""
def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values):
def __init__(
self,
_id,
_name,
_type,
_priority,
_state,
_ip_address,
_os_state,
_file_system_state,
_config_values,
):
"""
Init
Init.
Args:
_id: The node id
@@ -25,38 +32,44 @@ class ServiceNode(ActiveNode):
_osState: The operating system state of the node
_file_system_state: The file system state of the node
"""
super().__init__(_id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values)
super().__init__(
_id,
_name,
_type,
_priority,
_state,
_ip_address,
_os_state,
_file_system_state,
_config_values,
)
self.services = {}
def add_service(self, _service):
"""
Adds a service to the node
Adds a service to the node.
Args:
_service: The service to add
"""
self.services[_service.get_name()] = _service
def get_services(self):
"""
Gets the dictionary of services on this node
Gets the dictionary of services on this node.
Returns:
Dictionary of services on this node
"""
return self.services
def has_service(self, _protocol):
"""
Indicates whether a service is on a node
Indicates whether a service is on a node.
Returns:
True if service (protocol) is on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
return True
@@ -66,12 +79,11 @@ class ServiceNode(ActiveNode):
def service_running(self, _protocol):
"""
Indicates whether a service is in a running state on the node
Indicates whether a service is in a running state on the node.
Returns:
True if service (protocol) is in a running state on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
if service_value.get_state() != SOFTWARE_STATE.PATCHING:
@@ -84,12 +96,11 @@ class ServiceNode(ActiveNode):
def service_is_overwhelmed(self, _protocol):
"""
Indicates whether a service is in an overwhelmed state on the node
Indicates whether a service is in an overwhelmed state on the node.
Returns:
True if service (protocol) is in an overwhelmed state on the node
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
if service_value.get_state() == SOFTWARE_STATE.OVERWHELMED:
@@ -102,61 +113,61 @@ class ServiceNode(ActiveNode):
def set_service_state(self, _protocol, _state):
"""
Sets the state of a service (protocol) on the node
Sets the state of a service (protocol) on the node.
Args:
_protocol: The service (protocol)
_state: The state value
"""
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:
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:
service_value.patching_count = self.config_values.service_patching_duration
service_value.patching_count = (
self.config_values.service_patching_duration
)
else:
# Do nothing
pass
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"
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
"""
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
service_value.patching_count = (
self.config_values.service_patching_duration
)
def get_service_state(self, _protocol):
"""
Gets the state of a service
Gets the state of a service.
Returns:
The state of the service
"""
for service_key, service_value in self.services.items():
if service_key == _protocol:
return service_value.get_state()
def update_services_patching_status(self):
"""
Updates the patching counter for any service that are patching
"""
"""Updates the patching counter for any service that are patching."""
for service_key, service_value in self.services.items():
if service_value.get_state() == SOFTWARE_STATE.PATCHING:
service_value.reduce_patching_count()