#1355 - Carried out full renaming in node.py, active_node.py, passive_node.py, and service_node.py to make params and variable names explicit.
- Made the same renaming in the yaml laydown config files. - Added Type hints wherever I've been. - Added a custom NodeType in custom_typing.py to encompass the Union of ActiveNode, PassiveNode, ServiceNode.
This commit is contained in:
Chris McCarthy
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user