Merge remote-tracking branch 'origin/dev' into feature/2137-refactor-request-api

This commit is contained in:
Marek Wolan
2024-01-31 10:05:09 +00:00
63 changed files with 2392 additions and 488 deletions

View File

@@ -38,9 +38,6 @@ class Application(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.health_state_visible = SoftwareHealthState.UNUSED
self.health_state_actual = SoftwareHealthState.UNUSED
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
@@ -95,6 +92,9 @@ class Application(IOSoftware):
if self.operating_state == ApplicationOperatingState.CLOSED:
self.sys_log.info(f"Running Application {self.name}")
self.operating_state = ApplicationOperatingState.RUNNING
# set software health state to GOOD if initially set to UNUSED
if self.health_state_actual == SoftwareHealthState.UNUSED:
self.set_health_state(SoftwareHealthState.GOOD)
def _application_loop(self):
"""The main application loop."""

View File

@@ -72,7 +72,7 @@ class DataManipulationBot(DatabaseClient):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.attack()))
return rm
@@ -83,7 +83,7 @@ class DataManipulationBot(DatabaseClient):
payload: Optional[str] = None,
port_scan_p_of_success: float = 0.1,
data_manipulation_p_of_success: float = 0.1,
repeat: bool = False,
repeat: bool = True,
):
"""
Configure the DataManipulatorBot to communicate with a DatabaseService.
@@ -168,6 +168,12 @@ class DataManipulationBot(DatabaseClient):
Calls the parent classes execute method before starting the application loop.
"""
super().run()
def attack(self):
"""Perform the attack steps after opening the application."""
if not self._can_perform_action():
_LOGGER.debug("Data manipulation application attempted to execute but it cannot perform actions right now.")
self.run()
self._application_loop()
def _application_loop(self):
@@ -198,4 +204,4 @@ class DataManipulationBot(DatabaseClient):
:param timestep: The timestep value to update the bot's state.
"""
self._application_loop()
pass

View File

@@ -41,6 +41,9 @@ class PacketCapture:
def setup_logger(self):
"""Set up the logger configuration."""
if not SIM_OUTPUT.save_pcap_logs:
return
log_path = self._get_log_path()
file_handler = logging.FileHandler(filename=log_path)
@@ -88,5 +91,6 @@ class PacketCapture:
:param frame: The PCAP frame to capture.
"""
msg = frame.model_dump_json()
self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL
if SIM_OUTPUT.save_pcap_logs:
msg = frame.model_dump_json()
self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL

View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -141,41 +141,76 @@ class SessionManager:
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[IPv4Address] = None,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager.
Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission.
If no session_id, a Session is established. Once established, the payload is sent to ``send_payload_to_nic``.
This method supports both unicast and Layer 3 broadcast transmissions. If `dst_ip_address` is an
IPv4Network, a broadcast is initiated. For unicast, the destination MAC address is resolved via ARP.
A new session is established if `session_id` is not provided, and an existing session is used otherwise.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
:param dst_ip_address: The destination IP address or network for broadcast. Optional.
:param dst_port: The destination port for the TCP packet. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:param is_reattempt: Flag to indicate if this is a reattempt after an ARP request. Default is False.
:return: The outcome of sending the frame, or None if sending was unsuccessful.
"""
is_broadcast = False
outbound_nic = None
dst_mac_address = None
# Use session details if session_id is provided
if session_id:
session = self.sessions_by_uuid[session_id]
dst_ip_address = self.sessions_by_uuid[session_id].with_ip_address
dst_port = self.sessions_by_uuid[session_id].dst_port
dst_ip_address = session.with_ip_address
dst_port = session.dst_port
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Determine if the payload is for broadcast or unicast
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# Handle broadcast transmission
if isinstance(dst_ip_address, IPv4Network):
is_broadcast = True
dst_ip_address = dst_ip_address.broadcast_address
if dst_ip_address:
# Find a suitable NIC for the broadcast
for nic in self.arp_cache.nics.values():
if dst_ip_address in nic.ip_network and nic.enabled:
dst_mac_address = "ff:ff:ff:ff:ff:ff"
outbound_nic = nic
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
return
# Resolve MAC address for unicast transmission
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Resolve outbound NIC for unicast transmission
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# If MAC address not found, initiate ARP request
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
# Reattempt payload transmission after ARP request
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
# Return None if reattempt fails
return
# Check if outbound NIC and destination MAC address are resolved
if not outbound_nic or not dst_mac_address:
return False
# Construct the frame for transmission
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
@@ -189,15 +224,17 @@ class SessionManager:
payload=payload,
)
if not session_id:
# Manage session for unicast transmission
if not (is_broadcast and session_id):
session_key = self._get_session_key(frame, inbound_frame=False)
session = self.sessions_by_key.get(session_key)
if not session:
# Create new session
# Create a new session if it doesn't exist
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
# Send the frame through the NIC
return outbound_nic.send_frame(frame)
def receive_frame(self, frame: Frame):

View File

@@ -1,4 +1,4 @@
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -130,20 +130,28 @@ class SoftwareManager:
def send_payload_to_session_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
session_id: Optional[str] = None,
) -> bool:
"""
Send a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
return self.session_manager.receive_payload_from_software_manager(
payload=payload, dst_ip_address=dest_ip_address, dst_port=dest_port, session_id=session_id
payload=payload,
dst_ip_address=dest_ip_address,
dst_port=dest_port,
session_id=session_id,
)
def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str):

View File

@@ -41,6 +41,9 @@ class SysLog:
The logger is set to the DEBUG level, and is equipped with a handler that writes to a file and filters out
JSON-like messages.
"""
if not SIM_OUTPUT.save_sys_logs:
return
log_path = self._get_log_path()
file_handler = logging.FileHandler(filename=log_path)
file_handler.setLevel(logging.DEBUG)
@@ -91,7 +94,8 @@ class SysLog:
:param msg: The message to be logged.
"""
self.logger.debug(msg)
if SIM_OUTPUT.save_sys_logs:
self.logger.debug(msg)
def info(self, msg: str):
"""
@@ -99,7 +103,8 @@ class SysLog:
:param msg: The message to be logged.
"""
self.logger.info(msg)
if SIM_OUTPUT.save_sys_logs:
self.logger.info(msg)
def warning(self, msg: str):
"""
@@ -107,7 +112,8 @@ class SysLog:
:param msg: The message to be logged.
"""
self.logger.warning(msg)
if SIM_OUTPUT.save_sys_logs:
self.logger.warning(msg)
def error(self, msg: str):
"""
@@ -115,7 +121,8 @@ class SysLog:
:param msg: The message to be logged.
"""
self.logger.error(msg)
if SIM_OUTPUT.save_sys_logs:
self.logger.error(msg)
def critical(self, msg: str):
"""
@@ -123,4 +130,5 @@ class SysLog:
:param msg: The message to be logged.
"""
self.logger.critical(msg)
if SIM_OUTPUT.save_sys_logs:
self.logger.critical(msg)

View File

@@ -3,6 +3,8 @@ from typing import Any, Dict, List, Literal, Optional, Union
from primaite import getLogger
from primaite.simulator.file_system.file_system import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.file_system.folder import Folder
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.software_manager import SoftwareManager
@@ -22,7 +24,7 @@ class DatabaseService(Service):
password: Optional[str] = None
backup_server: IPv4Address = None
backup_server_ip: IPv4Address = None
"""IP address of the backup server."""
latest_backup_directory: str = None
@@ -36,7 +38,6 @@ class DatabaseService(Service):
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
self._db_file: File
self._create_db_file()
def set_original_state(self):
@@ -45,8 +46,8 @@ class DatabaseService(Service):
super().set_original_state()
vals_to_include = {
"password",
"_connections",
"backup_server",
"connections",
"backup_server_ip",
"latest_backup_directory",
"latest_backup_file_name",
}
@@ -64,7 +65,7 @@ class DatabaseService(Service):
:param: backup_server_ip: The IP address of the backup server
"""
self.backup_server = backup_server
self.backup_server_ip = backup_server
def backup_database(self) -> bool:
"""Create a backup of the database to the configured backup server."""
@@ -73,7 +74,7 @@ class DatabaseService(Service):
return False
# check if the backup server was configured
if self.backup_server is None:
if self.backup_server_ip is None:
self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.")
return False
@@ -81,10 +82,14 @@ class DatabaseService(Service):
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
# send backup copy of database file to FTP server
if not self.db_file:
self.sys_log.error("Attempted to backup database file but it doesn't exist.")
return False
response = ftp_client_service.send_file(
dest_ip_address=self.backup_server,
src_file_name=self._db_file.name,
src_folder_name=self.folder.name,
dest_ip_address=self.backup_server_ip,
src_file_name=self.db_file.name,
src_folder_name="database",
dest_folder_name=str(self.uuid),
dest_file_name="database.db",
)
@@ -110,7 +115,7 @@ class DatabaseService(Service):
src_file_name="database.db",
dest_folder_name="downloads",
dest_file_name="database.db",
dest_ip_address=self.backup_server,
dest_ip_address=self.backup_server_ip,
)
if not response:
@@ -118,13 +123,10 @@ class DatabaseService(Service):
return False
# replace db file
self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db")
self.file_system.copy_file(
src_folder_name="downloads", src_file_name="database.db", dst_folder_name=self.folder.name
)
self._db_file = self.file_system.get_file(folder_name=self.folder.name, file_name="database.db")
self.file_system.delete_file(folder_name="database", file_name="database.db")
self.file_system.copy_file(src_folder_name="downloads", src_file_name="database.db", dst_folder_name="database")
if self._db_file is None:
if self.db_file is None:
self.sys_log.error("Copying database backup failed.")
return False
@@ -134,12 +136,30 @@ class DatabaseService(Service):
def _create_db_file(self):
"""Creates the Simulation File and sqlite file in the file system."""
self._db_file: File = self.file_system.create_file(folder_name="database", file_name="database.db")
self.folder = self.file_system.get_folder_by_id(self._db_file.folder_id)
self.file_system.create_file(folder_name="database", file_name="database.db")
@property
def db_file(self) -> File:
"""Returns the database file."""
return self.file_system.get_file(folder_name="database", file_name="database.db")
@property
def folder(self) -> Folder:
"""Returns the database folder."""
return self.file_system.get_folder_by_id(self.db_file.folder_id)
def _process_connect(
self, connection_id: str, password: Optional[str] = None
) -> Dict[str, Union[int, Dict[str, bool]]]:
"""Process an incoming connection request.
:param connection_id: A unique identifier for the connection
:type connection_id: str
:param password: Supplied password. It must match self.password for connection success, defaults to None
:type password: Optional[str], optional
:return: Response to connection request containing success info.
:rtype: Dict[str, Union[int, Dict[str, bool]]]
"""
status_code = 500 # Default internal server error
if self.operating_state == ServiceOperatingState.RUNNING:
status_code = 503 # service unavailable
@@ -184,7 +204,7 @@ class DatabaseService(Service):
self.sys_log.info(f"{self.name}: Running {query}")
if query == "SELECT":
if self.health_state_actual == SoftwareHealthState.GOOD:
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {
"status_code": 200,
"type": "sql",
@@ -195,17 +215,8 @@ class DatabaseService(Service):
else:
return {"status_code": 404, "data": False}
elif query == "DELETE":
if self.health_state_actual == SoftwareHealthState.GOOD:
self.health_state_actual = SoftwareHealthState.COMPROMISED
return {
"status_code": 200,
"type": "sql",
"data": False,
"uuid": query_id,
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
return {"status_code": 200, "type": "sql", "data": False, "uuid": query_id, "connection_id": connection_id}
else:
# Invalid query
return {"status_code": 500, "data": False}
@@ -265,3 +276,19 @@ class DatabaseService(Service):
software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
return payload["status_code"] == 200
def apply_timestep(self, timestep: int) -> None:
"""
Apply a single timestep of simulation dynamics to this service.
Here at the first step, the database backup is created, in addition to normal service update logic.
"""
if timestep == 1:
self.backup_database()
return super().apply_timestep(timestep)
def _update_patch_status(self) -> None:
"""Perform a database restore when the patching countdown is finished."""
super()._update_patch_status()
if self._patching_countdown is None:
self.restore_backup()

View File

@@ -89,6 +89,7 @@ class FTPClient(FTPServiceABC):
f"{self.name}: Successfully connected to FTP Server "
f"{dest_ip_address} via port {payload.ftp_command_args.value}"
)
self.add_connection(connection_id="server_connection", session_id=session_id)
return True
else:
if is_reattempt:

View File

@@ -99,5 +99,5 @@ class FTPServer(FTPServiceABC):
if payload.status_code is not None:
return False
self.send(self._process_ftp_command(payload=payload, session_id=session_id), session_id)
self._process_ftp_command(payload=payload, session_id=session_id)
return True

View File

@@ -1,7 +1,7 @@
import shutil
from abc import ABC
from ipaddress import IPv4Address
from typing import Optional
from typing import Dict, Optional
from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
@@ -16,6 +16,10 @@ class FTPServiceABC(Service, ABC):
Contains shared methods between both classes.
"""
def describe_state(self) -> Dict:
"""Returns a Dict of the FTPService state."""
return super().describe_state()
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
"""
Process the command in the FTP Packet.
@@ -52,10 +56,12 @@ class FTPServiceABC(Service, ABC):
folder_name = payload.ftp_command_args["dest_folder_name"]
file_size = payload.ftp_command_args["file_size"]
real_file_path = payload.ftp_command_args.get("real_file_path")
health_status = payload.ftp_command_args["health_status"]
is_real = real_file_path is not None
file = self.file_system.create_file(
file_name=file_name, folder_name=folder_name, size=file_size, real=is_real
)
file.health_status = health_status
self.sys_log.info(
f"{self.name}: Created item in {self.sys_log.hostname}: {payload.ftp_command_args['dest_folder_name']}/"
f"{payload.ftp_command_args['dest_file_name']}"
@@ -110,6 +116,7 @@ class FTPServiceABC(Service, ABC):
"dest_file_name": dest_file_name,
"file_size": file.sim_size,
"real_file_path": file.sim_path if file.real else None,
"health_status": file.health_status,
},
packet_payload_size=file.sim_size,
status_code=FTPStatusCode.OK if is_response else None,

View File

@@ -1,3 +1,4 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, Optional
@@ -43,9 +44,6 @@ class Service(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.health_state_visible = SoftwareHealthState.UNUSED
self.health_state_actual = SoftwareHealthState.UNUSED
def _can_perform_action(self) -> bool:
"""
Checks if the service can perform actions.
@@ -98,6 +96,7 @@ class Service(IOSoftware):
rm.add_request("enable", RequestType(func=lambda request, context: self.enable()))
return rm
@abstractmethod
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -118,7 +117,6 @@ class Service(IOSoftware):
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Stopping service {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
self.health_state_actual = SoftwareHealthState.UNUSED
def start(self, **kwargs) -> None:
"""Start the service."""
@@ -129,42 +127,39 @@ class Service(IOSoftware):
if self.operating_state == ServiceOperatingState.STOPPED:
self.sys_log.info(f"Starting service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
self.health_state_actual = SoftwareHealthState.GOOD
# set software health state to GOOD if initially set to UNUSED
if self.health_state_actual == SoftwareHealthState.UNUSED:
self.set_health_state(SoftwareHealthState.GOOD)
def pause(self) -> None:
"""Pause the service."""
if self.operating_state == ServiceOperatingState.RUNNING:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.PAUSED
self.health_state_actual = SoftwareHealthState.OVERWHELMED
def resume(self) -> None:
"""Resume paused service."""
if self.operating_state == ServiceOperatingState.PAUSED:
self.sys_log.info(f"Resuming service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
self.health_state_actual = SoftwareHealthState.GOOD
def restart(self) -> None:
"""Restart running service."""
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.RESTARTING
self.health_state_actual = SoftwareHealthState.OVERWHELMED
self.restart_countdown = self.restart_duration
def disable(self) -> None:
"""Disable the service."""
self.sys_log.info(f"Disabling Application {self.name}")
self.operating_state = ServiceOperatingState.DISABLED
self.health_state_actual = SoftwareHealthState.OVERWHELMED
def enable(self) -> None:
"""Enable the disabled service."""
if self.operating_state == ServiceOperatingState.DISABLED:
self.sys_log.info(f"Enabling Application {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
self.health_state_actual = SoftwareHealthState.OVERWHELMED
def apply_timestep(self, timestep: int) -> None:
"""
@@ -181,5 +176,4 @@ class Service(IOSoftware):
if self.restart_countdown <= 0:
_LOGGER.debug(f"Restarting finished for service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
self.health_state_actual = SoftwareHealthState.GOOD
self.restart_countdown -= 1

View File

@@ -13,6 +13,7 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import SoftwareHealthState
_LOGGER = getLogger(__name__)
@@ -123,7 +124,10 @@ class WebServer(Service):
# get all users
if db_client.query("SELECT"):
# query succeeded
self.set_health_state(SoftwareHealthState.GOOD)
response.status_code = HttpStatusCode.OK
else:
self.set_health_state(SoftwareHealthState.COMPROMISED)
return response
except Exception:

View File

@@ -2,8 +2,8 @@ import copy
from abc import abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Union
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -38,12 +38,12 @@ class SoftwareHealthState(Enum):
"Unused state."
GOOD = 1
"The software is in a good and healthy condition."
COMPROMISED = 2
"The software's security has been compromised."
OVERWHELMED = 3
"he software is overwhelmed and not functioning properly."
PATCHING = 4
PATCHING = 2
"The software is undergoing patching or updates."
COMPROMISED = 3
"The software's security has been compromised."
OVERWHELMED = 4
"he software is overwhelmed and not functioning properly."
class SoftwareCriticality(Enum):
@@ -71,9 +71,9 @@ class Software(SimComponent):
name: str
"The name of the software."
health_state_actual: SoftwareHealthState = SoftwareHealthState.GOOD
health_state_actual: SoftwareHealthState = SoftwareHealthState.UNUSED
"The actual health state of the software."
health_state_visible: SoftwareHealthState = SoftwareHealthState.GOOD
health_state_visible: SoftwareHealthState = SoftwareHealthState.UNUSED
"The health state of the software visible to the red agent."
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
"The criticality level of the software."
@@ -195,8 +195,9 @@ class Software(SimComponent):
def patch(self) -> None:
"""Perform a patch on the software."""
self._patching_countdown = self.patching_duration
self.set_health_state(SoftwareHealthState.PATCHING)
if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD):
self._patching_countdown = self.patching_duration
self.set_health_state(SoftwareHealthState.PATCHING)
def _update_patch_status(self) -> None:
"""Update the patch status of the software."""
@@ -282,7 +283,7 @@ class IOSoftware(Software):
Returns true if the software can perform actions.
"""
if self.software_manager and self.software_manager.node.operating_state is NodeOperatingState.OFF:
if self.software_manager and self.software_manager.node.operating_state != NodeOperatingState.ON:
_LOGGER.debug(f"{self.name} Error: {self.software_manager.node.hostname} is not online.")
return False
return True
@@ -303,13 +304,13 @@ class IOSoftware(Software):
"""
# if over or at capacity, set to overwhelmed
if len(self._connections) >= self.max_sessions:
self.health_state_actual = SoftwareHealthState.OVERWHELMED
self.set_health_state(SoftwareHealthState.OVERWHELMED)
self.sys_log.error(f"{self.name}: Connect request for {connection_id=} declined. Service is at capacity.")
return False
else:
# if service was previously overwhelmed, set to good because there is enough space for connections
if self.health_state_actual == SoftwareHealthState.OVERWHELMED:
self.health_state_actual = SoftwareHealthState.GOOD
self.set_health_state(SoftwareHealthState.GOOD)
# check that connection already doesn't exist
if not self._connections.get(connection_id):
@@ -350,19 +351,22 @@ class IOSoftware(Software):
self,
payload: Any,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
**kwargs,
) -> bool:
"""
Sends a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network. It delegates the actual sending process to the SoftwareManager.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:return: True if successful, False otherwise.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
if not self._can_perform_action():
return False