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 18f89faf03
commit 71f33ed44e
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