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

Merge branch 'dev' into feature/1752-dns-server-and-client

Browse Source
This commit is contained in:
Czar Echavez
2023-09-06 13:42:07 +01:00
parent c1ba3b0850 6b9928c9b9
commit 2cb0c238c9
72 changed files with 4611 additions and 589 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/primaite/simulator/system/applications/application.py
View File

@@ -1,6 +1,6 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, List, Set
from typing import Any, Dict, Set
from primaite.simulator.system.software import IOSoftware
@@ -8,13 +8,12 @@ from primaite.simulator.system.software import IOSoftware
class ApplicationOperatingState(Enum):
"""Enumeration of Application Operating States."""
RUNNING = 1
"The application is running."
CLOSED = 2
"The application is closed or not running."
INSTALLING = 3
"The application is being installed or updated."
RUNNING = 1
"The application is running."
CLOSED = 2
"The application is closed or not running."
INSTALLING = 3
"The application is being installed or updated."
class Application(IOSoftware):
@@ -36,23 +35,23 @@ class Application(IOSoftware):
@abstractmethod
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
Produce a dictionary describing the current state of this object.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: A dictionary containing key-value pairs representing the current state of the software.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
def apply_action(self, action: List[str]) -> None:
"""
Applies a list of actions to the Application.
:param action: A list of actions to apply.
"""
pass
state = super().describe_state()
state.update(
{
"opearting_state": self.operating_state.name,
"execution_control_status": self.execution_control_status,
"num_executions": self.num_executions,
"groups": list(self.groups),
}
)
return state
def reset_component_for_episode(self, episode: int):
"""

133
src/primaite/simulator/system/core/session_manager.py
View File

@@ -1,12 +1,14 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.core import SimComponent
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader
if TYPE_CHECKING:
from primaite.simulator.network.hardware.base import ARPCache
@@ -22,16 +24,16 @@ class Session(SimComponent):
source and destination IPs and ports.
:param protocol: The IP protocol used in the session.
:param src_ip: The source IP address.
:param dst_ip: The destination IP address.
:param src_ip_address: The source IP address.
:param dst_ip_address: The destination IP address.
:param src_port: The source port number (optional).
:param dst_port: The destination port number (optional).
:param connected: A flag indicating whether the session is connected.
"""
protocol: IPProtocol
src_ip: IPv4Address
dst_ip: IPv4Address
src_ip_address: IPv4Address
dst_ip_address: IPv4Address
src_port: Optional[Port]
dst_port: Optional[Port]
connected: bool = False
@@ -46,14 +48,23 @@ class Session(SimComponent):
:param session_key: Tuple containing the session details.
:return: A Session instance.
"""
protocol, src_ip, dst_ip, src_port, dst_port = session_key
return Session(protocol=protocol, src_ip=src_ip, dst_ip=dst_ip, src_port=src_port, dst_port=dst_port)
protocol, src_ip_address, dst_ip_address, src_port, dst_port = session_key
return Session(
protocol=protocol,
src_ip_address=src_ip_address,
dst_ip_address=dst_ip_address,
src_port=src_port,
dst_port=dst_port,
)
def describe_state(self) -> Dict:
"""
Describes the current state of the session as a dictionary.
Produce a dictionary describing the current state of this object.
:return: A dictionary containing the current state of the session.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
@@ -77,9 +88,12 @@ class SessionManager:
def describe_state(self) -> Dict:
"""
Describes the current state of the session manager as a dictionary.
Produce a dictionary describing the current state of this object.
:return: A dictionary containing the current state of the session manager.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
@@ -102,8 +116,8 @@ class SessionManager:
:return: A tuple containing the session key.
"""
protocol = frame.ip.protocol
src_ip = frame.ip.src_ip
dst_ip = frame.ip.dst_ip
src_ip_address = frame.ip.src_ip_address
dst_ip_address = frame.ip.dst_ip_address
if protocol == IPProtocol.TCP:
if from_source:
src_port = frame.tcp.src_port
@@ -121,9 +135,16 @@ class SessionManager:
else:
src_port = None
dst_port = None
return protocol, src_ip, dst_ip, src_port, dst_port
return protocol, src_ip_address, dst_ip_address, src_port, dst_port
def receive_payload_from_software_manager(self, payload: Any, session_id: Optional[int] = None):
def receive_payload_from_software_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager.
@@ -132,9 +153,50 @@ class SessionManager:
: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.
"""
# TODO: Implement session creation and
if session_id:
dest_ip_address = self.sessions_by_uuid[session_id].dst_ip_address
dest_port = self.sessions_by_uuid[session_id].dst_port
self.send_payload_to_nic(payload, session_id)
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dest_ip_address)
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dest_ip_address)
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dest_ip_address)
return self.receive_payload_from_software_manager(
payload=payload,
dest_ip_address=dest_ip_address,
dest_port=dest_port,
session_id=session_id,
is_reattempt=True,
)
else:
return
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
src_ip_address=outbound_nic.ip_address,
dst_ip_address=dest_ip_address,
),
tcp=TCPHeader(
src_port=dest_port,
dst_port=dest_port,
),
payload=payload,
)
if not session_id:
session_key = self._get_session_key(frame, from_source=True)
session = self.sessions_by_key.get(session_key)
if not session:
# Create new session
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
outbound_nic.send_frame(frame)
def send_payload_to_software_manager(self, payload: Any, session_id: int):
"""
@@ -145,18 +207,6 @@ class SessionManager:
"""
self.software_manager.receive_payload_from_session_manger()
def send_payload_to_nic(self, payload: Any, session_id: int):
"""
Send a payload across the Network.
Takes a payload and a session_id. Builds a Frame and sends it across the network via a NIC.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload originates from
"""
# TODO: Implement frame construction and sent to NIC.
pass
def receive_payload_from_nic(self, frame: Frame):
"""
Receive a Frame from the NIC.
@@ -175,3 +225,22 @@ class SessionManager:
self.sessions_by_uuid[session.uuid] = session
self.software_manager.receive_payload_from_session_manger(payload=frame, session=session)
# TODO: Implement the frame deconstruction and send to SoftwareManager.
def show(self, markdown: bool = False):
"""
Print tables describing the SessionManager.
Generate and print PrettyTable instances that show details about
session's destination IP Address, destination Ports and the protocol to use.
Output can be in Markdown format.
:param markdown: Use Markdown style in table output. Defaults to False.
"""
table = PrettyTable(["Destination IP", "Port", "Protocol"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Session Manager"
for session in self.sessions_by_key.values():
table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
print(table)

51
src/primaite/simulator/system/core/software_manager.py
View File

@@ -1,5 +1,8 @@
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
@@ -12,6 +15,10 @@ if TYPE_CHECKING:
from primaite.simulator.system.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from typing import Type, TypeVar
ServiceClass = TypeVar("ServiceClass", bound=Service)
class SoftwareManager:
"""A class that manages all running Services and Applications on a Node and facilitates their communication."""
@@ -28,18 +35,17 @@ class SoftwareManager:
self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {}
self.sys_log: SysLog = sys_log
def add_service(self, name: str, service: Service, port: Port, protocol: IPProtocol):
def add_service(self, service_class: Type[ServiceClass]):
"""
Add a Service to the manager.
:param name: The name of the service.
:param service: The service instance.
:param port: The port used by the service.
:param protocol: The network protocol used by the service.
:param: service_class: The class of the service to add
"""
service = service_class(software_manager=self, sys_log=self.sys_log)
service.software_manager = self
self.services[name] = service
self.port_protocol_mapping[(port, protocol)] = service
self.services[service.name] = service
self.port_protocol_mapping[(service.port, service.protocol)] = service
def add_application(self, name: str, application: Application, port: Port, protocol: IPProtocol):
"""
@@ -75,14 +81,24 @@ class SoftwareManager:
else:
raise ValueError(f"No {target_software_type.name.lower()} found with the name {target_software}")
def send_payload_to_session_manger(self, payload: Any, session_id: Optional[int] = None):
def send_payload_to_session_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
session_id: Optional[int] = None,
):
"""
Send a payload to the SessionManager.
: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.
"""
self.session_manager.receive_payload_from_software_manager(payload, session_id)
self.session_manager.receive_payload_from_software_manager(
payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
)
def receive_payload_from_session_manger(self, payload: Any, session: Session):
"""
@@ -97,3 +113,20 @@ class SoftwareManager:
# else:
# raise ValueError(f"No service or application found for port {port} and protocol {protocol}")
pass
def show(self, markdown: bool = False):
"""
Prints a table of the SwitchPorts on the Switch.
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Name", "Operating State", "Health State", "Port"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Software Manager"
for service in self.services.values():
table.add_row(
[service.name, service.operating_state.name, service.health_state_actual.name, service.port.value]
)
print(table)

25
src/primaite/simulator/system/core/sys_log.py
View File

@@ -1,6 +1,8 @@
import logging
from pathlib import Path
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator import TEMP_SIM_OUTPUT
@@ -43,7 +45,7 @@ class SysLog:
file_handler = logging.FileHandler(filename=log_path)
file_handler.setLevel(logging.DEBUG)
log_format = "%(asctime)s %(levelname)s: %(message)s"
log_format = "%(asctime)s::%(levelname)s::%(message)s"
file_handler.setFormatter(logging.Formatter(log_format))
self.logger = logging.getLogger(f"{self.hostname}_sys_log")
@@ -52,6 +54,27 @@ class SysLog:
self.logger.addFilter(_NotJSONFilter())
def show(self, last_n: int = 10, markdown: bool = False):
"""
Print the Node Sys Log as a table.
Generate and print PrettyTable instance that shows the Nodes Sys Log, with columns Timestamp, Level,
and Massage.
:param markdown: Use Markdown style in table output. Defaults to False.
"""
table = PrettyTable(["Timestamp", "Level", "Message"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.hostname} Sys Log"
if self._get_log_path().exists():
with open(self._get_log_path()) as file:
lines = file.readlines()
for line in lines[-last_n:]:
table.add_row(line.strip().split("::"))
print(table)
def _get_log_path(self) -> Path:
"""
Constructs the path for the log file based on the hostname.

11
src/primaite/simulator/system/processes/process.py
View File

@@ -27,12 +27,13 @@ class Process(Software):
@abstractmethod
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
Produce a dictionary describing the current state of this object.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: A dictionary containing key-value pairs representing the current state of the software.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
state = super().describe_state()
state.update({"operating_state": self.operating_state.name})
return state

76
src/primaite/simulator/system/services/database.py Normal file
View File

@@ -0,0 +1,76 @@
from typing import Dict
from primaite.simulator.file_system.file_system_file_type import FileSystemFileType
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.system.services.service import Service
class DatabaseService(Service):
"""Service loosely modelled on Microsoft SQL Server."""
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: Current state of this object and child objects.
:rtype: Dict
"""
return super().describe_state()
def uninstall(self) -> None:
"""
Undo installation procedure.
This method deletes files created when installing the database, and the database folder if it is empty.
"""
super().uninstall()
node: Node = self.parent
node.file_system.delete_file(self.primary_store)
node.file_system.delete_file(self.transaction_log)
if self.secondary_store:
node.file_system.delete_file(self.secondary_store)
if len(self.folder.files) == 0:
node.file_system.delete_folder(self.folder)
def install(self) -> None:
"""Perform first time install on a node, creating necessary files."""
super().install()
assert isinstance(self.parent, Node), "Database install can only happen after the db service is added to a node"
self._setup_files()
def _setup_files(
self,
db_size: int = 1000,
use_secondary_db_file: bool = False,
secondary_db_size: int = 300,
folder_name: str = "database",
):
"""Set up files that are required by the database on the parent host.
:param db_size: Initial file size of the main database file, defaults to 1000
:type db_size: int, optional
:param use_secondary_db_file: Whether to use a secondary database file, defaults to False
:type use_secondary_db_file: bool, optional
:param secondary_db_size: Size of the secondary db file, defaults to None
:type secondary_db_size: int, optional
:param folder_name: Name of the folder which will be setup to hold the db files, defaults to "database"
:type folder_name: str, optional
"""
# note that this parent.file_system.create_folder call in the future will be authenticated by using permissions
# handler. This permission will be granted based on service account given to the database service.
self.parent: Node
self.folder = self.parent.file_system.create_folder(folder_name)
self.primary_store = self.parent.file_system.create_file(
"db_primary_store", db_size, FileSystemFileType.MDF, folder=self.folder
)
self.transaction_log = self.parent.file_system.create_file(
"db_transaction_log", "1", FileSystemFileType.LDF, folder=self.folder
)
if use_secondary_db_file:
self.secondary_store = self.parent.file_system.create_file(
"db_secondary_store", secondary_db_size, FileSystemFileType.NDF, folder=self.folder
)
else:
self.secondary_store = None

34
src/primaite/simulator/system/services/red_services/data_manipulator_service.py Normal file
View File

@@ -0,0 +1,34 @@
from ipaddress import IPv4Address
from typing import Any, Optional
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
class DataManipulatorService(Service):
"""
Red Agent Data Integration Service.
The Service represents a bot that causes files/folders in the File System to
become corrupted.
"""
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulatorBot"
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def start(self, target_ip_address: IPv4Address, payload: Optional[Any] = "DELETE TABLE users", **kwargs):
"""
Run the DataManipulatorService actions.
:param: target_ip_address: The IP address of the target machine to attack
:param: payload: The payload to send to the target machine
"""
super().start()
self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=target_ip_address, dest_port=self.port
)

110
src/primaite/simulator/system/services/service.py
View File

@@ -1,9 +1,12 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, List
from typing import Any, Dict, Optional
from primaite import getLogger
from primaite.simulator.core import Action, ActionManager
from primaite.simulator.system.software import IOSoftware
_LOGGER = getLogger(__name__)
class ServiceOperatingState(Enum):
"""Enumeration of Service Operating States."""
@@ -29,29 +32,36 @@ class Service(IOSoftware):
Services are programs that run in the background and may perform input/output operations.
"""
operating_state: ServiceOperatingState
operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED
"The current operating state of the Service."
restart_duration: int = 5
"How many timesteps does it take to restart this service."
_restart_countdown: Optional[int] = None
"If currently restarting, how many timesteps remain until the restart is finished."
def _init_action_manager(self) -> ActionManager:
am = super()._init_action_manager()
am.add_action("stop", Action(func=lambda request, context: self.stop()))
am.add_action("start", Action(func=lambda request, context: self.start()))
am.add_action("pause", Action(func=lambda request, context: self.pause()))
am.add_action("resume", Action(func=lambda request, context: self.resume()))
am.add_action("restart", Action(func=lambda request, context: self.restart()))
am.add_action("disable", Action(func=lambda request, context: self.disable()))
am.add_action("enable", Action(func=lambda request, context: self.enable()))
return am
@abstractmethod
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
Produce a dictionary describing the current state of this object.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: A dictionary containing key-value pairs representing the current state of the software.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
def apply_action(self, action: List[str]) -> None:
"""
Applies a list of actions to the Service.
:param action: A list of actions to apply.
"""
pass
state = super().describe_state()
state.update({"operating_state": self.operating_state.name})
return state
def reset_component_for_episode(self, episode: int):
"""
@@ -85,3 +95,69 @@ class Service(IOSoftware):
:return: True if successful, False otherwise.
"""
pass
def stop(self) -> None:
"""Stop the service."""
_LOGGER.debug(f"Stopping service {self.name}")
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Stopping service {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
def start(self, **kwargs) -> None:
"""Start the service."""
_LOGGER.debug(f"Starting service {self.name}")
if self.operating_state == ServiceOperatingState.STOPPED:
self.sys_log.info(f"Starting service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
def pause(self) -> None:
"""Pause the service."""
_LOGGER.debug(f"Pausing service {self.name}")
if self.operating_state == ServiceOperatingState.RUNNING:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.PAUSED
def resume(self) -> None:
"""Resume paused service."""
_LOGGER.debug(f"Resuming service {self.name}")
if self.operating_state == ServiceOperatingState.PAUSED:
self.sys_log.info(f"Resuming service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
def restart(self) -> None:
"""Restart running service."""
_LOGGER.debug(f"Restarting service {self.name}")
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.RESTARTING
self.restart_countdown = self.restarting_duration
def disable(self) -> None:
"""Disable the service."""
_LOGGER.debug(f"Disabling service {self.name}")
self.sys_log.info(f"Disabling Application {self.name}")
self.operating_state = ServiceOperatingState.DISABLED
def enable(self) -> None:
"""Enable the disabled service."""
_LOGGER.debug(f"Enabling service {self.name}")
if self.operating_state == ServiceOperatingState.DISABLED:
self.sys_log.info(f"Enabling Application {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
def apply_timestep(self, timestep: int) -> None:
"""
Apply a single timestep of simulation dynamics to this service.
In this instance, if any multi-timestep processes are currently occurring (such as restarting or installation),
then they are brought one step closer to being finished.
:param timestep: The current timestep number. (Amount of time since simulation episode began)
:type timestep: int
"""
super().apply_timestep(timestep)
if self.operating_state == ServiceOperatingState.RESTARTING:
if self.restart_countdown <= 0:
_LOGGER.debug(f"Restarting finished for service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
self.restart_countdown -= 1

115
src/primaite/simulator/system/software.py
View File

@@ -1,9 +1,10 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, List, Set
from typing import Any, Dict
from primaite.simulator.core import SimComponent
from primaite.simulator.core import Action, ActionManager, SimComponent
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.sys_log import SysLog
class SoftwareType(Enum):
@@ -62,11 +63,11 @@ class Software(SimComponent):
name: str
"The name of the software."
health_state_actual: SoftwareHealthState
health_state_actual: SoftwareHealthState = SoftwareHealthState.GOOD
"The actual health state of the software."
health_state_visible: SoftwareHealthState
health_state_visible: SoftwareHealthState = SoftwareHealthState.GOOD
"The health state of the software visible to the red agent."
criticality: SoftwareCriticality
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
"The criticality level of the software."
patching_count: int = 0
"The count of patches applied to the software, defaults to 0."
@@ -74,30 +75,44 @@ class Software(SimComponent):
"The count of times the software has been scanned, defaults to 0."
revealed_to_red: bool = False
"Indicates if the software has been revealed to red agent, defaults is False."
software_manager: Any = None
"An instance of Software Manager that is used by the parent node."
sys_log: SysLog = None
"An instance of SysLog that is used by the parent node."
def _init_action_manager(self) -> ActionManager:
am = super()._init_action_manager()
am.add_action(
"compromise",
Action(
func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED),
),
)
am.add_action("scan", Action(func=lambda request, context: self.scan()))
return am
@abstractmethod
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
Produce a dictionary describing the current state of this object.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: A dictionary containing key-value pairs representing the current state of the software.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
def apply_action(self, action: List[str]) -> None:
"""
Applies a list of actions to the software.
The specifics of how these actions are applied should be implemented in subclasses.
:param action: A list of actions to apply.
:type action: List[str]
"""
pass
state = super().describe_state()
state.update(
{
"health_state": self.health_state_actual.name,
"health_state_red_view": self.health_state_visible.name,
"criticality": self.criticality.name,
"patching_count": self.patching_count,
"scanning_count": self.scanning_count,
"revealed_to_red": self.revealed_to_red,
}
)
return state
def reset_component_for_episode(self, episode: int):
"""
@@ -109,15 +124,42 @@ class Software(SimComponent):
"""
pass
@staticmethod
def get_install():
def set_health_state(self, health_state: SoftwareHealthState) -> None:
"""
This method ensures the software has to have a way to install it.
Assign a new health state to this software.
This can be used by the software manager to install the software.
Note: this should only be possible when the software is currently running, but the software base class has no
operating state, only subclasses do. So subclasses will need to implement this check. TODO: check if this should
be changed so that the base Software class has a running attr.
:param health_state: New health state to assign to the software
:type health_state: SoftwareHealthState
"""
self.health_state_actual = health_state
def install(self) -> None:
"""
Perform first-time setup of this service on a node.
This is an abstract class that should be overwritten by specific applications or services. It must be called
after the service is already associate with a node. For example, a service may need to authenticate with a
server during installation, or create files in the node's filesystem.
"""
pass
def uninstall(self) -> None:
"""Uninstall this service from a node.
This is an abstract class that should be overwritten by applications or services. It must be called after the
`install` method has already been run on that node. It should undo any installation steps, for example by
deleting files, or contacting a server.
"""
pass
def scan(self) -> None:
"""Update the observed health status to match the actual health status."""
self.health_state_visible = self.health_state_actual
class IOSoftware(Software):
"""
@@ -137,21 +179,30 @@ class IOSoftware(Software):
"Indicates if the software uses TCP protocol for communication. Default is True."
udp: bool = True
"Indicates if the software uses UDP protocol for communication. Default is True."
ports: Set[Port]
"The set of ports to which the software is connected."
port: Port
"The port to which the software is connected."
@abstractmethod
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
Produce a dictionary describing the current state of this object.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
:return: A dictionary containing key-value pairs representing the current state of the software.
:return: Current state of this object and child objects.
:rtype: Dict
"""
pass
state = super().describe_state()
state.update(
{
"installing_count": self.installing_count,
"max_sessions": self.max_sessions,
"tcp": self.tcp,
"udp": self.udp,
"ports": [port.name for port in self.ports], # TODO: not sure if this should be port.name or port.value
}
)
return state
def send(self, payload: Any, session_id: str, **kwargs) -> bool:
"""