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

#2887 - Resolve conflicts from merge

Browse Source
This commit is contained in:
Charlie Crane
2025-01-23 09:17:27 +00:00
parent 3957142afd d2d09abcb1
commit fbef78db51
174 changed files with 7047 additions and 8412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -264,7 +264,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

2
src/primaite/simulator/_package_data/network_simulator_demo.ipynb
View File

@@ -664,7 +664,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

4
src/primaite/simulator/file_system/file.py
View File

@@ -130,8 +130,8 @@ class File(FileSystemItemABC):
Return False if corruption is detected, otherwise True
"""
warnings.warn("NODE_FILE_CHECKHASH is currently not implemented.")
self.sys_log.warning("NODE_FILE_CHECKHASH is currently not implemented.")
warnings.warn("node_file_checkhash is currently not implemented.")
self.sys_log.warning("node_file_checkhash is currently not implemented.")
return False
if self.deleted:

10
src/primaite/simulator/file_system/file_system.py
View File

@@ -30,6 +30,11 @@ class FileSystem(SimComponent):
num_file_deletions: int = 0
"Number of file deletions in the current step."
_default_folder_scan_duration: Optional[int] = None
"Override default scan duration for folders"
_default_folder_restore_duration: Optional[int] = None
"Override default restore duration for folders"
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Ensure a default root folder
@@ -258,6 +263,11 @@ class FileSystem(SimComponent):
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
self.folders[folder.uuid] = folder
# set the folder scan and restore durations.
if self._default_folder_scan_duration is not None:
folder.scan_duration = self._default_folder_scan_duration
if self._default_folder_restore_duration is not None:
folder.restore_duration = self._default_folder_restore_duration
return folder
def delete_folder(self, folder_name: str) -> bool:

5
src/primaite/simulator/file_system/file_system_item_abc.py
View File

@@ -43,6 +43,9 @@ def convert_size(size_bytes: int) -> str:
class FileSystemItemHealthStatus(Enum):
"""Status of the FileSystemItem."""
NONE = 0
"""File system item health status is not known."""
GOOD = 1
"""File/Folder is OK."""
@@ -72,7 +75,7 @@ class FileSystemItemABC(SimComponent):
health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD
"Actual status of the current FileSystemItem"
visible_health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.GOOD
visible_health_status: FileSystemItemHealthStatus = FileSystemItemHealthStatus.NONE
"Visible status of the current FileSystemItem"
previous_hash: Optional[str] = None

21
src/primaite/simulator/file_system/folder.py
View File

@@ -46,7 +46,7 @@ class Folder(FileSystemItemABC):
:param sys_log: The SysLog instance to us to create system logs.
"""
super().__init__(**kwargs)
self._scanned_this_step: bool = False
self.sys_log.info(f"Created file /{self.name} (id: {self.uuid})")
def _init_request_manager(self) -> RequestManager:
@@ -83,6 +83,7 @@ class Folder(FileSystemItemABC):
state = super().describe_state()
state["files"] = {file.name: file.describe_state() for uuid, file in self.files.items()}
state["deleted_files"] = {file.name: file.describe_state() for uuid, file in self.deleted_files.items()}
state["scanned_this_step"] = self._scanned_this_step
return state
def show(self, markdown: bool = False):
@@ -135,7 +136,7 @@ class Folder(FileSystemItemABC):
def pre_timestep(self, timestep: int) -> None:
"""Apply pre-timestep logic."""
super().pre_timestep(timestep)
self._scanned_this_step = False
for file in self.files.values():
file.pre_timestep(timestep)
@@ -148,9 +149,17 @@ class Folder(FileSystemItemABC):
for file_id in self.files:
file = self.get_file_by_id(file_uuid=file_id)
file.scan()
if file.visible_health_status == FileSystemItemHealthStatus.CORRUPT:
self.health_status = FileSystemItemHealthStatus.CORRUPT
# set folder health to worst file's health by generating a list of file healths. If no files, use 0
self.health_status = FileSystemItemHealthStatus(
max(
[f.health_status.value for f in self.files.values()]
or [
0,
]
)
)
self.visible_health_status = self.health_status
self._scanned_this_step = True
def _reveal_to_red_timestep(self) -> None:
"""Apply reveal to red timestep."""
@@ -387,8 +396,8 @@ class Folder(FileSystemItemABC):
Return False if corruption is detected, otherwise True
"""
warnings.warn("NODE_FOLDER_CHECKHASH is currently not implemented.")
self.sys_log.error("NODE_FOLDER_CHECKHASH is currently not implemented.")
warnings.warn("node_folder_checkhash is currently not implemented.")
self.sys_log.error("node_folder_checkhash is currently not implemented.")
return False
if self.deleted:

6
src/primaite/simulator/network/creation.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC, abstractmethod
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Literal, Type
from typing import Any, ClassVar, Dict, Literal, Optional, Type
from pydantic import BaseModel, model_validator
@@ -49,7 +49,7 @@ class NetworkNodeAdder(BaseModel):
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None:
"""
Register a network node adder class.
@@ -58,6 +58,8 @@ class NetworkNodeAdder(BaseModel):
:raises ValueError: When attempting to register a name that is already reserved.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate node adder {identifier}")
cls._registry[identifier] = cls

27
src/primaite/simulator/network/hardware/base.py
View File

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump()
class UserManager(Service):
class UserManager(Service, identifier="UserManager"):
"""
Manages users within the PrimAITE system, handling creation, authentication, and administration.
@@ -833,11 +833,18 @@ class UserManager(Service):
:param disabled_admins: A dictionary of currently disabled admin users by their usernames
"""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserManager."""
type: str = "UserManager"
config: "UserManager.ConfigSchema" = Field(default_factory=lambda: UserManager.ConfigSchema())
users: Dict[str, User] = {}
def __init__(self, **kwargs):
"""
Initializes a UserManager instanc.
Initializes a UserManager instance.
:param username: The username for the default admin user
:param password: The password for the default admin user
@@ -1130,13 +1137,20 @@ class RemoteUserSession(UserSession):
return state
class UserSessionManager(Service):
class UserSessionManager(Service, identifier="UserSessionManager"):
"""
Manages user sessions on a Node, including local and remote sessions.
This class handles authentication, session management, and session timeouts for users interacting with the Node.
"""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserSessionManager."""
type: str = "UserSessionManager"
config: "UserSessionManager.ConfigSchema" = Field(default_factory=lambda: UserSessionManager.ConfigSchema())
local_session: Optional[UserSession] = None
"""The current local user session, if any."""
@@ -1554,7 +1568,6 @@ class Node(SimComponent, ABC):
red_scan_countdown: int = 0
"Time steps until reveal to red scan is complete."
@classmethod
def from_config(cls, config: Dict) -> "Node":
"""Create Node object from a given configuration dictionary."""
@@ -1564,7 +1577,7 @@ class Node(SimComponent, ABC):
obj = cls(config=cls.ConfigSchema(**config))
return obj
def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None:
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a node type.
@@ -1572,10 +1585,10 @@ class Node(SimComponent, ABC):
:type identifier: str
:raises ValueError: When attempting to register an node with a name that is already allocated.
"""
if identifier == "default":
super().__init_subclass__(**kwargs)
if identifier is None:
return
identifier = identifier.lower()
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls

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

@@ -1,10 +1,12 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from abc import abstractmethod
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, ClassVar, Dict, Optional, Set, Type
from pydantic import Field
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestPermissionValidator, RequestType
from primaite.simulator.system.software import IOSoftware, SoftwareHealthState
@@ -21,13 +23,20 @@ class ApplicationOperatingState(Enum):
"The application is being installed or updated."
class Application(IOSoftware):
class Application(IOSoftware, ABC):
"""
Represents an Application in the simulation environment.
Applications are user-facing programs that may perform input/output operations.
"""
class ConfigSchema(IOSoftware.ConfigSchema, ABC):
"""Config Schema for Application class."""
type: str
config: ConfigSchema = Field(default_factory=lambda: Application.ConfigSchema())
operating_state: ApplicationOperatingState = ApplicationOperatingState.CLOSED
"The current operating state of the Application."
execution_control_status: str = "manual"
@@ -44,21 +53,36 @@ class Application(IOSoftware):
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None:
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an application type.
:param identifier: Uniquely specifies an application class by name. Used for finding items by config.
:type identifier: str
:type identifier: Optional[str]
:raises ValueError: When attempting to register an application with a name that is already allocated.
"""
if identifier == "default":
return
super().__init_subclass__(**kwargs)
if identifier is None:
return
if identifier in cls._registry:
raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
@classmethod
def from_config(cls, config: Dict) -> "Application":
"""Create an application from a config dictionary.
:param config: dict of options for application components constructor
:type config: dict
:return: The application component.
:rtype: Application
"""
if config["type"] not in cls._registry:
raise ValueError(f"Invalid Application type {config['type']}")
application_class = cls._registry[config["type"]]
application_object = application_class(config=application_class.ConfigSchema(**config))
return application_object
def __init__(self, **kwargs):
super().__init__(**kwargs)

16
src/primaite/simulator/system/applications/database_client.py
View File

@@ -6,7 +6,7 @@ from typing import Any, Dict, Optional, Union
from uuid import uuid4
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel
from pydantic import BaseModel, Field
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
@@ -67,11 +67,19 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
Extends the Application class to provide functionality for connecting, querying, and disconnecting from a
Database Service. It mainly operates over TCP protocol.
:ivar server_ip_address: The IPv4 address of the Database Service server, defaults to None.
"""
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for DatabaseClient."""
type: str = "DatabaseClient"
db_server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
config: ConfigSchema = Field(default_factory=lambda: DatabaseClient.ConfigSchema())
server_ip_address: Optional[IPv4Address] = None
"""The IPv4 address of the Database Service server, defaults to None."""
server_password: Optional[str] = None
_query_success_tracker: Dict[str, bool] = {}
"""Keep track of connections that were established or verified during this step. Used for rewards."""
@@ -93,6 +101,8 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.server_ip_address = self.config.db_server_ip
self.server_password = self.config.server_password
def _init_request_manager(self) -> RequestManager:
"""

9
src/primaite/simulator/system/applications/nmap.py
View File

@@ -3,7 +3,7 @@ from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Final, List, Optional, Set, Tuple, Union
from prettytable import PrettyTable
from pydantic import validate_call
from pydantic import Field, validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
@@ -52,6 +52,13 @@ class NMAP(Application, identifier="NMAP"):
as ping scans to discover active hosts and port scans to detect open ports on those hosts.
"""
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for NMAP."""
type: str = "NMAP"
config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema())
_active_port_scans: Dict[str, PortScanPayload] = {}
_port_scan_responses: Dict[str, PortScanPayload] = {}

84
src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py
View File

@@ -2,9 +2,9 @@
from abc import abstractmethod
from enum import Enum
from ipaddress import IPv4Address
from typing import Dict, Optional, Union
from typing import Dict, Optional, Set, Union
from pydantic import BaseModel, Field, validate_call
from pydantic import Field, validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -45,10 +45,10 @@ class C2Payload(Enum):
"""C2 Input Command payload. Used by the C2 Server to send a command to the c2 beacon."""
OUTPUT = "output_command"
"""C2 Output Command. Used by the C2 Beacon to send the results of a Input command to the c2 server."""
"""C2 Output Command. Used by the C2 Beacon to send the results of an Input command to the c2 server."""
class AbstractC2(Application, identifier="AbstractC2"):
class AbstractC2(Application):
"""
An abstract command and control (c2) application.
@@ -60,9 +60,25 @@ class AbstractC2(Application, identifier="AbstractC2"):
Defaults to masquerading as HTTP (Port 80) via TCP.
Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite.
Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite.
"""
class ConfigSchema(Application.ConfigSchema):
"""Configuration for AbstractC2."""
keep_alive_frequency: int = Field(default=5, ge=1)
"""The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon."""
masquerade_protocol: IPProtocol = Field(default=PROTOCOL_LOOKUP["TCP"])
"""The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP."""
masquerade_port: Port = Field(default=PORT_LOOKUP["HTTP"])
"""The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP."""
listen_on_ports: Set[Port] = {PORT_LOOKUP["HTTP"], PORT_LOOKUP["FTP"], PORT_LOOKUP["DNS"]}
config: ConfigSchema = Field(default_factory=lambda: AbstractC2.ConfigSchema())
c2_connection_active: bool = False
"""Indicates if the c2 server and c2 beacon are currently connected."""
@@ -75,19 +91,6 @@ class AbstractC2(Application, identifier="AbstractC2"):
keep_alive_inactivity: int = 0
"""Indicates how many timesteps since the last time the c2 application received a keep alive."""
class _C2Opts(BaseModel):
"""A Pydantic Schema for the different C2 configuration options."""
keep_alive_frequency: int = Field(default=5, ge=1)
"""The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon."""
masquerade_protocol: IPProtocol = Field(default=PROTOCOL_LOOKUP["TCP"])
"""The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP."""
masquerade_port: Port = Field(default=PORT_LOOKUP["HTTP"])
"""The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP."""
c2_config: _C2Opts = _C2Opts()
"""
Holds the current configuration settings of the C2 Suite.
@@ -100,6 +103,12 @@ class AbstractC2(Application, identifier="AbstractC2"):
C2 beacon to reconfigure it's configuration settings.
"""
def __init__(self, **kwargs):
"""Initialise the C2 applications to by default listen for HTTP traffic."""
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
def _craft_packet(
self, c2_payload: C2Payload, c2_command: Optional[C2Command] = None, command_options: Optional[Dict] = {}
) -> C2Packet:
@@ -118,13 +127,13 @@ class AbstractC2(Application, identifier="AbstractC2"):
:type c2_command: C2Command.
:param command_options: The relevant C2 Beacon parameters.F
:type command_options: Dict
:return: Returns the construct C2Packet
:return: Returns the constructed C2Packet
:rtype: C2Packet
"""
constructed_packet = C2Packet(
masquerade_protocol=self.c2_config.masquerade_protocol,
masquerade_port=self.c2_config.masquerade_port,
keep_alive_frequency=self.c2_config.keep_alive_frequency,
masquerade_protocol=self.config.masquerade_protocol,
masquerade_port=self.config.masquerade_port,
keep_alive_frequency=self.config.keep_alive_frequency,
payload_type=c2_payload,
command=c2_command,
payload=command_options,
@@ -140,13 +149,6 @@ class AbstractC2(Application, identifier="AbstractC2"):
"""
return super().describe_state()
def __init__(self, **kwargs):
"""Initialise the C2 applications to by default listen for HTTP traffic."""
kwargs["listen_on_ports"] = {PORT_LOOKUP["HTTP"], PORT_LOOKUP["FTP"], PORT_LOOKUP["DNS"]}
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@property
def _host_ftp_client(self) -> Optional[FTPClient]:
"""Return the FTPClient that is installed C2 Application's host.
@@ -330,8 +332,8 @@ class AbstractC2(Application, identifier="AbstractC2"):
if self.send(
payload=keep_alive_packet,
dest_ip_address=self.c2_remote_connection,
dest_port=self.c2_config.masquerade_port,
ip_protocol=self.c2_config.masquerade_protocol,
dest_port=self.config.masquerade_port,
ip_protocol=self.config.masquerade_protocol,
session_id=session_id,
):
# Setting the keep_alive_sent guard condition to True. This is used to prevent packet storms.
@@ -340,8 +342,8 @@ class AbstractC2(Application, identifier="AbstractC2"):
self.sys_log.info(f"{self.name}: Keep Alive sent to {self.c2_remote_connection}")
self.sys_log.debug(
f"{self.name}: Keep Alive sent to {self.c2_remote_connection} "
f"Masquerade Port: {self.c2_config.masquerade_port} "
f"Masquerade Protocol: {self.c2_config.masquerade_protocol} "
f"Masquerade Port: {self.config.masquerade_port} "
f"Masquerade Protocol: {self.config.masquerade_protocol} "
)
return True
else:
@@ -376,15 +378,15 @@ class AbstractC2(Application, identifier="AbstractC2"):
# Updating the C2 Configuration attribute.
self.c2_config.masquerade_port = payload.masquerade_port
self.c2_config.masquerade_protocol = payload.masquerade_protocol
self.c2_config.keep_alive_frequency = payload.keep_alive_frequency
self.config.masquerade_port = payload.masquerade_port
self.config.masquerade_protocol = payload.masquerade_protocol
self.config.keep_alive_frequency = payload.keep_alive_frequency
self.sys_log.debug(
f"{self.name}: C2 Config Resolved Config from Keep Alive:"
f"Masquerade Port: {self.c2_config.masquerade_port}"
f"Masquerade Protocol: {self.c2_config.masquerade_protocol}"
f"Keep Alive Frequency: {self.c2_config.keep_alive_frequency}"
f"Masquerade Port: {self.config.masquerade_port}"
f"Masquerade Protocol: {self.config.masquerade_protocol}"
f"Keep Alive Frequency: {self.config.keep_alive_frequency}"
)
# This statement is intended to catch on the C2 Application that is listening for connection.
@@ -410,8 +412,8 @@ class AbstractC2(Application, identifier="AbstractC2"):
self.keep_alive_inactivity = 0
self.keep_alive_frequency = 5
self.c2_remote_connection = None
self.c2_config.masquerade_port = PORT_LOOKUP["HTTP"]
self.c2_config.masquerade_protocol = PROTOCOL_LOOKUP["TCP"]
self.config.masquerade_port = PORT_LOOKUP["HTTP"]
self.config.masquerade_protocol = PROTOCOL_LOOKUP["TCP"]
@abstractmethod
def _confirm_remote_connection(self, timestep: int) -> bool:

52
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -3,7 +3,7 @@ from ipaddress import IPv4Address
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from pydantic import Field, validate_call
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
@@ -12,8 +12,9 @@ from primaite.simulator.system.applications.red_applications.c2 import ExfilOpts
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.services.terminal.terminal import Terminal, TerminalClientConnection
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import PORT_LOOKUP
from primaite.utils.validation.ip_protocol import IPProtocol, PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
class C2Beacon(AbstractC2, identifier="C2Beacon"):
@@ -32,15 +33,30 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
2. Leveraging the terminal application to execute requests (dependent on the command given)
3. Sending the RequestResponse back to the C2 Server (Command output)
Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite.
Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite.
"""
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Beacon."""
type: str = "C2Beacon"
c2_server_ip_address: Optional[IPV4Address] = None
keep_alive_frequency: int = 5
masquerade_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"]
masquerade_port: Port = PORT_LOOKUP["HTTP"]
config: ConfigSchema = Field(default_factory=lambda: C2Beacon.ConfigSchema())
keep_alive_attempted: bool = False
"""Indicates if a keep alive has been attempted to be sent this timestep. Used to prevent packet storms."""
terminal_session: TerminalClientConnection = None
"The currently in use terminal session."
def __init__(self, **kwargs):
kwargs["name"] = "C2Beacon"
super().__init__(**kwargs)
@property
def _host_terminal(self) -> Optional[Terminal]:
"""Return the Terminal that is installed on the same machine as the C2 Beacon."""
@@ -119,10 +135,6 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
rm.add_request("configure", request_type=RequestType(func=_configure))
return rm
def __init__(self, **kwargs):
kwargs["name"] = "C2Beacon"
super().__init__(**kwargs)
# Configure is practically setter method for the ``c2.config`` attribute that also ties into the request manager.
@validate_call
def configure(
@@ -146,7 +158,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
masquerade_port | What port should the C2 traffic use? (TCP or UDP)
These configuration options are used to reassign the fields in the inherited inner class
``c2_config``.
``config``.
If a connection is already in progress then this method also sends a keep alive to the C2
Server in order for the C2 Server to sync with the new configuration settings.
@@ -162,9 +174,9 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:return: Returns True if the configuration was successful, False otherwise.
"""
self.c2_remote_connection = IPv4Address(c2_server_ip_address)
self.c2_config.keep_alive_frequency = keep_alive_frequency
self.c2_config.masquerade_port = masquerade_port
self.c2_config.masquerade_protocol = masquerade_protocol
self.config.keep_alive_frequency = keep_alive_frequency
self.config.masquerade_port = masquerade_port
self.config.masquerade_protocol = masquerade_protocol
self.sys_log.info(
f"{self.name}: Configured {self.name} with remote C2 server connection: {c2_server_ip_address=}."
)
@@ -263,14 +275,12 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
if self.send(
payload=output_packet,
dest_ip_address=self.c2_remote_connection,
dest_port=self.c2_config.masquerade_port,
ip_protocol=self.c2_config.masquerade_protocol,
dest_port=self.config.masquerade_port,
ip_protocol=self.config.masquerade_protocol,
session_id=session_id,
):
self.sys_log.info(f"{self.name}: Command output sent to {self.c2_remote_connection}")
self.sys_log.debug(
f"{self.name}: on {self.c2_config.masquerade_port} via {self.c2_config.masquerade_protocol}"
)
self.sys_log.debug(f"{self.name}: on {self.config.masquerade_port} via {self.config.masquerade_protocol}")
return True
else:
self.sys_log.warning(
@@ -562,7 +572,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:rtype bool:
"""
self.keep_alive_attempted = False # Resetting keep alive sent.
if self.keep_alive_inactivity == self.c2_config.keep_alive_frequency:
if self.keep_alive_inactivity == self.config.keep_alive_frequency:
self.sys_log.info(
f"{self.name}: Attempting to Send Keep Alive to {self.c2_remote_connection} at timestep {timestep}."
)
@@ -627,9 +637,9 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
self.c2_connection_active,
self.c2_remote_connection,
self.keep_alive_inactivity,
self.c2_config.keep_alive_frequency,
self.c2_config.masquerade_protocol,
self.c2_config.masquerade_port,
self.config.keep_alive_frequency,
self.config.masquerade_protocol,
self.config.masquerade_port,
]
)
print(table)

23
src/primaite/simulator/system/applications/red_applications/c2/c2_server.py
View File

@@ -2,7 +2,7 @@
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from pydantic import Field, validate_call
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
@@ -31,9 +31,16 @@ class C2Server(AbstractC2, identifier="C2Server"):
1. Sending commands to the C2 Beacon. (Command input)
2. Parsing terminal RequestResponses back to the Agent.
Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite.
Please refer to the Command-and-Control notebook for an in-depth example of the C2 Suite.
"""
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Server."""
type: str = "C2Server"
config: ConfigSchema = Field(default_factory=lambda: C2Server.ConfigSchema())
current_command_output: RequestResponse = None
"""The Request Response by the last command send. This attribute is updated by the method _handle_command_output."""
@@ -251,8 +258,8 @@ class C2Server(AbstractC2, identifier="C2Server"):
payload=command_packet,
dest_ip_address=self.c2_remote_connection,
session_id=self.c2_session.uuid,
dest_port=self.c2_config.masquerade_port,
ip_protocol=self.c2_config.masquerade_protocol,
dest_port=self.config.masquerade_port,
ip_protocol=self.config.masquerade_protocol,
):
self.sys_log.info(f"{self.name}: Successfully sent {given_command}.")
self.sys_log.info(f"{self.name}: Awaiting command response {given_command}.")
@@ -334,11 +341,11 @@ class C2Server(AbstractC2, identifier="C2Server"):
:return: Returns False if the C2 beacon is considered dead. Otherwise True.
:rtype bool:
"""
if self.keep_alive_inactivity > self.c2_config.keep_alive_frequency:
if self.keep_alive_inactivity > self.config.keep_alive_frequency:
self.sys_log.info(f"{self.name}: C2 Beacon connection considered dead due to inactivity.")
self.sys_log.debug(
f"{self.name}: Did not receive expected keep alive connection from {self.c2_remote_connection}"
f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.c2_config.keep_alive_frequency}"
f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.config.keep_alive_frequency}"
f"{self.name}: Last Keep Alive received at {(timestep - self.keep_alive_inactivity)}"
)
self._reset_c2_connection()
@@ -389,8 +396,8 @@ class C2Server(AbstractC2, identifier="C2Server"):
[
self.c2_connection_active,
self.c2_remote_connection,
self.c2_config.masquerade_protocol,
self.c2_config.masquerade_port,
self.config.masquerade_protocol,
self.config.masquerade_port,
]
)
print(table)

21
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -3,6 +3,8 @@ from enum import IntEnum
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestResponse
@@ -10,6 +12,7 @@ from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
@@ -40,6 +43,18 @@ class DataManipulationAttackStage(IntEnum):
class DataManipulationBot(Application, identifier="DataManipulationBot"):
"""A bot that simulates a script which performs a SQL injection attack."""
class ConfigSchema(Application.ConfigSchema):
"""Configuration schema for DataManipulationBot."""
type: str = "DataManipulationBot"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "DELETE"
port_scan_p_of_success: float = 0.1
data_manipulation_p_of_success: float = 0.1
config: "DataManipulationBot.ConfigSchema" = Field(default_factory=lambda: DataManipulationBot.ConfigSchema())
payload: Optional[str] = None
port_scan_p_of_success: float = 0.1
data_manipulation_p_of_success: float = 0.1
@@ -56,6 +71,12 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None
self.server_ip_address = self.config.server_ip
self.server_password = self.config.server_password
self.payload = self.config.payload
self.port_scan_p_of_success = self.config.port_scan_p_of_success
self.data_manipulation_p_of_success = self.config.data_manipulation_p_of_success
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

25
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -3,11 +3,14 @@ from enum import IntEnum
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
@@ -32,6 +35,20 @@ class DoSAttackStage(IntEnum):
class DoSBot(DatabaseClient, identifier="DoSBot"):
"""A bot that simulates a Denial of Service attack."""
class ConfigSchema(DatabaseClient.ConfigSchema):
"""ConfigSchema for DoSBot."""
type: str = "DoSBot"
target_ip_address: Optional[IPV4Address] = None
target_port: Port = PORT_LOOKUP["POSTGRES_SERVER"]
payload: Optional[str] = None
repeat: bool = False
port_scan_p_of_success: float = 0.1
dos_intensity: float = 1.0
max_sessions: int = 1000
config: "DoSBot.ConfigSchema" = Field(default_factory=lambda: DoSBot.ConfigSchema())
target_ip_address: Optional[IPv4Address] = None
"""IP address of the target service."""
@@ -56,7 +73,13 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.max_sessions = 1000 # override normal max sessions
self.target_ip_address = self.config.target_ip_address
self.target_port = self.config.target_port
self.payload = self.config.payload
self.repeat = self.config.repeat
self.port_scan_p_of_success = self.config.port_scan_p_of_success
self.dos_intensity = self.config.dos_intensity
self.max_sessions = self.config.max_sessions
def _init_request_manager(self) -> RequestManager:
"""

15
src/primaite/simulator/system/applications/red_applications/ransomware_script.py
View File

@@ -3,12 +3,14 @@ from ipaddress import IPv4Address
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
@@ -18,6 +20,16 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
:ivar payload: The attack stage query payload. (Default ENCRYPT)
"""
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for RansomwareScript."""
type: str = "RansomwareScript"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "ENCRYPT"
config: "RansomwareScript.ConfigSchema" = Field(default_factory=lambda: RansomwareScript.ConfigSchema())
server_ip_address: Optional[IPv4Address] = None
"""IP address of node which hosts the database."""
server_password: Optional[str] = None
@@ -32,6 +44,9 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None
self.server_ip_address = self.config.server_ip
self.server_password = self.config.server_password
self.payload = self.config.payload
def describe_state(self) -> Dict:
"""

12
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -4,7 +4,7 @@ from ipaddress import IPv4Address
from typing import Dict, List, Optional
from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, Field
from primaite import getLogger
from primaite.interface.request import RequestResponse
@@ -30,7 +30,13 @@ class WebBrowser(Application, identifier="WebBrowser"):
The application requests and loads web pages using its domain name and requesting IP addresses using DNS.
"""
target_url: Optional[str] = None
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for WebBrowser."""
type: str = "WebBrowser"
target_url: Optional[str] = None
config: "WebBrowser.ConfigSchema" = Field(default_factory=lambda: WebBrowser.ConfigSchema())
domain_name_ip_address: Optional[IPv4Address] = None
"The IP address of the domain name for the webpage."
@@ -86,7 +92,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
:param: url: The address of the web page the browser requests
:type: url: str
"""
url = url or self.target_url
url = url or self.config.target_url
if not self._can_perform_action():
return False

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

@@ -106,7 +106,7 @@ class SoftwareManager:
return True
return False
def install(self, software_class: Type[IOSoftware], **install_kwargs):
def install(self, software_class: Type[IOSoftware], software_config: Optional[IOSoftware.ConfigSchema] = None):
"""
Install an Application or Service.
@@ -115,13 +115,22 @@ class SoftwareManager:
if software_class in self._software_class_to_name_map:
self.sys_log.warning(f"Cannot install {software_class} as it is already installed")
return
software = software_class(
software_manager=self,
sys_log=self.sys_log,
file_system=self.file_system,
dns_server=self.dns_server,
**install_kwargs,
)
if software_config is None:
software = software_class(
software_manager=self,
sys_log=self.sys_log,
file_system=self.file_system,
dns_server=self.dns_server,
)
else:
software = software_class(
software_manager=self,
sys_log=self.sys_log,
file_system=self.file_system,
dns_server=self.dns_server,
config=software_config,
)
software.parent = self.node
if isinstance(software, Application):
self.node.applications[software.uuid] = software

10
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -5,6 +5,7 @@ from abc import abstractmethod
from typing import Any, Dict, Optional, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
@@ -14,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class ARP(Service):
class ARP(Service, identifier="ARP"):
"""
The ARP (Address Resolution Protocol) Service.
@@ -22,6 +23,13 @@ class ARP(Service):
sends ARP requests and replies, and processes incoming ARP packets.
"""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ARP."""
type: str = "ARP"
config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema())
arp: Dict[IPV4Address, ARPEntry] = {}
def __init__(self, **kwargs):

13
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -3,6 +3,8 @@ from ipaddress import IPv4Address
from typing import Any, Dict, List, Literal, Optional, Union
from uuid import uuid4
from pydantic import Field
from primaite import getLogger
from primaite.simulator.file_system.file_system import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
@@ -17,13 +19,21 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DatabaseService(Service):
class DatabaseService(Service, identifier="DatabaseService"):
"""
A class for simulating a generic SQL Server service.
This class inherits from the `Service` class and provides methods to simulate a SQL database.
"""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DatabaseService."""
type: str = "DatabaseService"
backup_server_ip: Optional[IPv4Address] = None
config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema())
password: Optional[str] = None
"""Password that needs to be provided by clients if they want to connect to the DatabaseService."""
@@ -42,6 +52,7 @@ class DatabaseService(Service):
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self._create_db_file()
self.backup_server_ip = self.config.backup_server_ip
def install(self):
"""

10
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -2,6 +2,8 @@
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest
from primaite.simulator.system.core.software_manager import SoftwareManager
@@ -12,9 +14,15 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DNSClient(Service):
class DNSClient(Service, identifier="DNSClient"):
"""Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
config: "DNSClient.ConfigSchema" = Field(default_factory=lambda: DNSClient.ConfigSchema())
dns_cache: Dict[str, IPv4Address] = {}
"A dict of known mappings between domain/URLs names and IPv4 addresses."
dns_server: Optional[IPv4Address] = None

11
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -3,6 +3,7 @@ from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.dns import DNSPacket
@@ -13,9 +14,17 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DNSServer(Service):
class DNSServer(Service, identifier="DNSServer"):
"""Represents a DNS Server as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSServer."""
type: str = "DNSServer"
domain_mapping: dict = {}
config: "DNSServer.ConfigSchema" = Field(default_factory=lambda: DNSServer.ConfigSchema())
dns_table: Dict[str, IPv4Address] = {}
"A dict of mappings between domain names and IPv4 addresses."

19
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -2,6 +2,8 @@
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
@@ -9,20 +11,28 @@ from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
from primaite.simulator.system.services.service import Service
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPClient(FTPServiceABC):
class FTPClient(FTPServiceABC, identifier="FTPClient"):
"""
A class for simulating an FTP client service.
This class inherits from the `Service` class and provides methods to emulate FTP
This class inherits from the `FTPServiceABC` class and provides methods to emulate FTP
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
"""
config: "FTPClient.ConfigSchema" = Field(default_factory=lambda: FTPClient.ConfigSchema())
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for FTPClient."""
type: str = "FTPClient"
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["port"] = PORT_LOOKUP["FTP"]
@@ -108,6 +118,7 @@ class FTPClient(FTPServiceABC):
session_id: Optional[str] = None,
is_reattempt: Optional[bool] = False,
) -> bool:
self._active = True
"""
Connects the client to a given FTP server.
@@ -164,6 +175,7 @@ class FTPClient(FTPServiceABC):
:param: is_reattempt: Set to True if attempt to disconnect from FTP Server has been attempted. Default False.
:type: is_reattempt: Optional[bool]
"""
self._active = True
# send a disconnect request payload to FTP server
payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.QUIT)
software_manager: SoftwareManager = self.software_manager
@@ -209,6 +221,7 @@ class FTPClient(FTPServiceABC):
:param: session_id: The id of the session
:type: session_id: Optional[str]
"""
self._active = True
# check if the file to transfer exists on the client
file_to_transfer: File = self.file_system.get_file(folder_name=src_folder_name, file_name=src_file_name)
if not file_to_transfer:
@@ -266,6 +279,7 @@ class FTPClient(FTPServiceABC):
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"].
:type: dest_port: Optional[int]
"""
self._active = True
# check if FTP is currently connected to IP
self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
@@ -317,6 +331,7 @@ class FTPClient(FTPServiceABC):
This helps prevent an FTP request loop - FTP client and servers can exist on
the same node.
"""
self._active = True
if not self._can_perform_action():
return False

14
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -1,26 +1,36 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Any, Optional
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
from primaite.simulator.system.services.service import Service
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPServer(FTPServiceABC):
class FTPServer(FTPServiceABC, identifier="FTPServer"):
"""
A class for simulating an FTP server service.
This class inherits from the `Service` class and provides methods to emulate FTP
This class inherits from the `FTPServiceABC` class and provides methods to emulate FTP
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
"""
config: "FTPServer.ConfigSchema" = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["port"] = PORT_LOOKUP["FTP"]

24
src/primaite/simulator/system/services/ftp/ftp_service.py
View File

@@ -3,9 +3,11 @@ from abc import ABC
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import StrictBool
from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.utils.validation.port import Port
@@ -16,9 +18,22 @@ class FTPServiceABC(Service, ABC):
Contains shared methods between both classes.
"""
_active: StrictBool = False
"""Flag that is True on timesteps where service transmits data and False when idle. Used for describe_state."""
def pre_timestep(self, timestep: int) -> None:
"""When a new timestep begins, clear the _active attribute."""
self._active = False
return super().pre_timestep(timestep)
def describe_state(self) -> Dict:
"""Returns a Dict of the FTPService state."""
return super().describe_state()
state = super().describe_state()
# override so that the service is shows as running only if actively transmitting data this timestep
if self.operating_state == ServiceOperatingState.RUNNING and not self._active:
state["operating_state"] = ServiceOperatingState.STOPPED.value
return state
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
"""
@@ -29,6 +44,7 @@ class FTPServiceABC(Service, ABC):
:param: session_id: session ID linked to the FTP Packet. Optional.
:type: session_id: Optional[str]
"""
self._active = True
if payload.ftp_command is not None:
self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.")
@@ -51,6 +67,7 @@ class FTPServiceABC(Service, ABC):
:param: payload: The FTP Packet that contains the file data
:type: FTPPacket
"""
self._active = True
try:
file_name = payload.ftp_command_args["dest_file_name"]
folder_name = payload.ftp_command_args["dest_folder_name"]
@@ -106,6 +123,7 @@ class FTPServiceABC(Service, ABC):
:param: is_response: is true if the data being sent is in response to a request. Default False.
:type: is_response: bool
"""
self._active = True
# send STOR request
payload: FTPPacket = FTPPacket(
ftp_command=FTPCommand.STOR,
@@ -135,6 +153,7 @@ class FTPServiceABC(Service, ABC):
:param: payload: The FTP Packet that contains the file data
:type: FTPPacket
"""
self._active = True
try:
# find the file
file_name = payload.ftp_command_args["src_file_name"]
@@ -181,6 +200,7 @@ class FTPServiceABC(Service, ABC):
:return: True if successful, False otherwise.
"""
self._active = True
self.sys_log.info(f"{self.name}: Sending FTP {payload.ftp_command.name} {payload.ftp_command_args}")
return super().send(

11
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -3,6 +3,8 @@ import secrets
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, Union
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
@@ -14,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class ICMP(Service):
class ICMP(Service, identifier="ICMP"):
"""
The Internet Control Message Protocol (ICMP) service.
@@ -22,6 +24,13 @@ class ICMP(Service):
network diagnostics, notably the ping command.
"""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ICMP."""
type: str = "ICMP"
config: "ICMP.ConfigSchema" = Field(default_factory=lambda: ICMP.ConfigSchema())
request_replies: Dict = {}
def __init__(self, **kwargs):

11
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -3,6 +3,8 @@ from datetime import datetime
from ipaddress import IPv4Address
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.service import Service, ServiceOperatingState
@@ -12,9 +14,16 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPClient(Service):
class NTPClient(Service, identifier="NTPClient"):
"""Represents a NTP client as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPClient."""
type: str = "NTPClient"
config: "NTPClient.ConfigSchema" = Field(default_factory=lambda: NTPClient.ConfigSchema())
ntp_server: Optional[IPv4Address] = None
"The NTP server the client sends requests to."
time: Optional[datetime] = None

11
src/primaite/simulator/system/services/ntp/ntp_server.py
View File

@@ -2,6 +2,8 @@
from datetime import datetime
from typing import Dict, Optional
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.service import Service
@@ -11,9 +13,16 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPServer(Service):
class NTPServer(Service, identifier="NTPServer"):
"""Represents a NTP server as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPServer."""
type: str = "NTPServer"
config: "NTPServer.ConfigSchema" = Field(default_factory=lambda: NTPServer.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["port"] = PORT_LOOKUP["NTP"]

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

@@ -1,10 +1,12 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from abc import abstractmethod
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, ClassVar, Dict, Optional, Type
from pydantic import Field
from primaite import getLogger
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestPermissionValidator, RequestType
@@ -37,6 +39,13 @@ class Service(IOSoftware):
Services are programs that run in the background and may perform input/output operations.
"""
class ConfigSchema(IOSoftware.ConfigSchema, ABC):
"""Config Schema for Service class."""
type: str
config: "Service.ConfigSchema" = Field(default_factory=lambda: Service.ConfigSchema())
operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED
"The current operating state of the Service."
@@ -52,7 +61,7 @@ class Service(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None:
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a hostnode type.
@@ -60,15 +69,30 @@ class Service(IOSoftware):
:type identifier: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
"""
if identifier == "default":
super().__init_subclass__(**kwargs)
if identifier is None:
return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower()
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
@classmethod
def from_config(cls, config: Dict) -> "Service":
"""Create a service from a config dictionary.
:param config: dict of options for service components constructor
:type config: dict
:return: The service component.
:rtype: Service
"""
if config["type"] not in cls._registry:
raise ValueError(f"Invalid service type {config['type']}")
service_class = cls._registry[config["type"]]
service_object = service_class(config=service_class.ConfigSchema(**config))
return service_object
def _can_perform_action(self) -> bool:
"""
Checks if the service can perform actions.
@@ -232,14 +256,14 @@ class Service(IOSoftware):
def disable(self) -> bool:
"""Disable the service."""
self.sys_log.info(f"Disabling Application {self.name}")
self.sys_log.info(f"Disabling Service {self.name}")
self.operating_state = ServiceOperatingState.DISABLED
return True
def enable(self) -> bool:
"""Enable the disabled service."""
if self.operating_state == ServiceOperatingState.DISABLED:
self.sys_log.info(f"Enabling Application {self.name}")
self.sys_log.info(f"Enabling Service {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
return True
return False

15
src/primaite/simulator/system/services/terminal/terminal.py
View File

@@ -7,7 +7,7 @@ from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional, Union
from uuid import uuid4
from pydantic import BaseModel
from pydantic import BaseModel, Field
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
@@ -129,9 +129,16 @@ class RemoteTerminalConnection(TerminalClientConnection):
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
class Terminal(Service):
class Terminal(Service, identifier="Terminal"):
"""Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for Terminal."""
type: str = "Terminal"
config: "Terminal.ConfigSchema" = Field(default_factory=lambda: Terminal.ConfigSchema())
_client_connection_requests: Dict[str, Optional[Union[str, TerminalClientConnection]]] = {}
"""Dictionary of connect requests made to remote nodes."""
@@ -179,7 +186,7 @@ class Terminal(Service):
return RequestResponse(status="failure", data={})
rm.add_request(
"ssh_to_remote",
"node_session_remote_login",
request_type=RequestType(func=_remote_login),
)
@@ -279,7 +286,6 @@ class Terminal(Service):
:param password: Password for login.
:return: boolean, True if successful, else False
"""
# TODO: Un-comment this when UserSessionManager is merged.
connection_uuid = self.parent.user_session_manager.local_login(username=username, password=password)
if connection_uuid:
self.sys_log.info(f"{self.name}: Login request authorised, connection uuid: {connection_uuid}")
@@ -406,7 +412,6 @@ class Terminal(Service):
if isinstance(payload, SSHPacket):
if payload.transport_message == SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST:
# validate & add connection
# TODO: uncomment this as part of 2781
username = payload.user_account.username
password = payload.user_account.password
connection_id = self.parent.user_session_manager.remote_login(

11
src/primaite/simulator/system/services/web_server/web_server.py
View File

@@ -3,6 +3,8 @@ from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse
from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.http import (
HttpRequestMethod,
@@ -19,9 +21,16 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebServer(Service):
class WebServer(Service, identifier="WebServer"):
"""Class used to represent a Web Server Service in simulation."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for WebServer."""
type: str = "WebServer"
config: "WebServer.ConfigSchema" = Field(default_factory=lambda: WebServer.ConfigSchema())
response_codes_this_timestep: List[HttpStatusCode] = []
def describe_state(self) -> Dict:

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

@@ -1,13 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import copy
from abc import abstractmethod
from abc import ABC, abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
from pydantic import BaseModel, ConfigDict, Field
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
@@ -70,7 +70,7 @@ class SoftwareCriticality(Enum):
"The highest level of criticality."
class Software(SimComponent):
class Software(SimComponent, ABC):
"""
A base class representing software in a simulator environment.
@@ -78,14 +78,22 @@ class Software(SimComponent):
It outlines the fundamental attributes and behaviors expected of any software in the simulation.
"""
class ConfigSchema(BaseModel, ABC):
"""Configurable options for all software."""
model_config = ConfigDict(extra="forbid")
starting_health_state: SoftwareHealthState = SoftwareHealthState.UNUSED
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
fixing_duration: int = 2
config: ConfigSchema = Field(default_factory=lambda: Software.ConfigSchema())
name: str
"The name of the software."
health_state_actual: SoftwareHealthState = SoftwareHealthState.UNUSED
"The actual health state of the software."
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."
fixing_count: int = 0
"The count of patches applied to the software, defaults to 0."
scanning_count: int = 0
@@ -100,11 +108,13 @@ class Software(SimComponent):
"The FileSystem of the Node the Software is installed on."
folder: Optional[Folder] = None
"The folder on the file system the Software uses."
fixing_duration: int = 2
"The number of ticks it takes to patch the software."
_fixing_countdown: Optional[int] = None
"Current number of ticks left to patch the software."
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.health_state_actual = self.config.starting_health_state # don't remove this
def _init_request_manager(self) -> RequestManager:
"""
Initialise the request manager.
@@ -152,7 +162,7 @@ class Software(SimComponent):
{
"health_state_actual": self.health_state_actual.value,
"health_state_visible": self.health_state_visible.value,
"criticality": self.criticality.value,
"criticality": self.config.criticality.value,
"fixing_count": self.fixing_count,
"scanning_count": self.scanning_count,
"revealed_to_red": self.revealed_to_red,
@@ -201,7 +211,7 @@ class Software(SimComponent):
def fix(self) -> bool:
"""Perform a fix on the software."""
if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD):
self._fixing_countdown = self.fixing_duration
self._fixing_countdown = self.config.fixing_duration
self.set_health_state(SoftwareHealthState.FIXING)
return True
return False
@@ -233,7 +243,7 @@ class Software(SimComponent):
super().pre_timestep(timestep)
class IOSoftware(Software):
class IOSoftware(Software, ABC):
"""
Represents software in a simulator environment that is capable of input/output operations.
@@ -243,6 +253,13 @@ class IOSoftware(Software):
required.
"""
class ConfigSchema(Software.ConfigSchema, ABC):
"""Configuration options for all IO Software."""
listen_on_ports: Set[Port] = Field(default_factory=set)
config: ConfigSchema = Field(default_factory=lambda: IOSoftware.ConfigSchema())
installing_count: int = 0
"The number of times the software has been installed. Default is 0."
max_sessions: int = 100
@@ -260,6 +277,10 @@ class IOSoftware(Software):
_connections: Dict[str, Dict] = {}
"Active connections."
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.listen_on_ports = self.config.listen_on_ports
@abstractmethod
def describe_state(self) -> Dict:
"""