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

#3062 - First pass at unifying naming convention for discriminators (still errors) [skip ci]

Browse Source
This commit is contained in:
Marek Wolan
2025-02-03 16:24:03 +00:00
parent 055c853b0f
commit abccf4afc5
195 changed files with 2824 additions and 2802 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/core.py
View File

@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
..code::python
class WebBrowser(Application, discriminator="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager() # all requests generic to any Application get initialised
rm.add_request(...) # initialise any requests specific to the web browser

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:

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

@@ -396,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

@@ -22,7 +22,7 @@ class NetworkNodeAdder(BaseModel):
Here is a template that users can use to define custom node adders:
```
class YourNodeAdder(NetworkNodeAdder, discriminator="your_name"):
class YourNodeAdder(NetworkNodeAdder, discriminator="your-name"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
property_1 : str
property_2 : int
@@ -99,13 +99,13 @@ class NetworkNodeAdder(BaseModel):
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
class OfficeLANAdder(NetworkNodeAdder, discriminator="office_lan"):
class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
"""Creates an office LAN."""
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
"""Configuration schema for OfficeLANAdder."""
type: Literal["office_lan"] = "office_lan"
type: Literal["office-lan"] = "office-lan"
lan_name: str
"""Name of lan used for generating hostnames for new nodes."""
subnet_base: int

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

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump()
class UserManager(Service, discriminator="UserManager"):
class UserManager(Service, discriminator="user-manager"):
"""
Manages users within the PrimAITE system, handling creation, authentication, and administration.
@@ -836,7 +836,7 @@ class UserManager(Service, discriminator="UserManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserManager."""
type: str = "UserManager"
type: str = "user-manager"
config: "UserManager.ConfigSchema" = Field(default_factory=lambda: UserManager.ConfigSchema())
@@ -849,7 +849,7 @@ class UserManager(Service, discriminator="UserManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserManager"
kwargs["name"] = "user-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1037,7 +1037,7 @@ class UserManager(Service, discriminator="UserManager"):
@property
def _user_session_manager(self) -> "UserSessionManager":
return self.software_manager.software["UserSessionManager"] # noqa
return self.software_manager.software["user-session-manager"] # noqa
class UserSession(SimComponent):
@@ -1137,7 +1137,7 @@ class RemoteUserSession(UserSession):
return state
class UserSessionManager(Service, discriminator="UserSessionManager"):
class UserSessionManager(Service, discriminator="user-session-manager"):
"""
Manages user sessions on a Node, including local and remote sessions.
@@ -1147,7 +1147,7 @@ class UserSessionManager(Service, discriminator="UserSessionManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserSessionManager."""
type: str = "UserSessionManager"
type: str = "user-session-manager"
config: "UserSessionManager.ConfigSchema" = Field(default_factory=lambda: UserSessionManager.ConfigSchema())
@@ -1179,7 +1179,7 @@ class UserSessionManager(Service, discriminator="UserSessionManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserSessionManager"
kwargs["name"] = "user-session-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1289,7 +1289,7 @@ class UserSessionManager(Service, discriminator="UserSessionManager"):
:return: The UserManager instance.
"""
return self.software_manager.software["UserManager"] # noqa
return self.software_manager.software["user-manager"] # noqa
def pre_timestep(self, timestep: int) -> None:
"""Apply any pre-timestep logic that helps make sure we have the correct observations."""
@@ -1608,17 +1608,17 @@ class Node(SimComponent, ABC):
@property
def user_manager(self) -> Optional[UserManager]:
"""The Nodes User Manager."""
return self.software_manager.software.get("UserManager") # noqa
return self.software_manager.software.get("user-manager") # noqa
@property
def user_session_manager(self) -> Optional[UserSessionManager]:
"""The Nodes User Session Manager."""
return self.software_manager.software.get("UserSessionManager") # noqa
return self.software_manager.software.get("user-session-manager") # noqa
@property
def terminal(self) -> Optional[Terminal]:
"""The Nodes Terminal."""
return self.software_manager.software.get("Terminal")
"""The Node's Terminal."""
return self.software_manager.software.get("terminal")
def local_login(self, username: str, password: str) -> Optional[str]:
"""

2
src/primaite/simulator/network/hardware/nodes/host/computer.py
View File

@@ -33,6 +33,6 @@ class Computer(HostNode, discriminator="computer"):
* Web Browser
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient}
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "ftp-client": FTPClient}
pass

26
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -262,7 +262,7 @@ class NIC(IPWiredNetworkInterface):
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
class HostNode(Node, discriminator="HostNode"):
class HostNode(Node, discriminator="host-node"):
"""
Represents a host node in the network.
@@ -309,14 +309,14 @@ class HostNode(Node, discriminator="HostNode"):
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"HostARP": HostARP,
"ICMP": ICMP,
"DNSClient": DNSClient,
"NTPClient": NTPClient,
"WebBrowser": WebBrowser,
"NMAP": NMAP,
"UserSessionManager": UserSessionManager,
"UserManager": UserManager,
"Terminal": Terminal,
"icmp": ICMP,
"dns-client": DNSClient,
"ntp-client": NTPClient,
"web-browser": WebBrowser,
"nmap": NMAP,
"user-session-manager": UserSessionManager,
"user-manager": UserManager,
"terminal": Terminal,
}
"""List of system software that is automatically installed on nodes."""
@@ -337,7 +337,7 @@ class HostNode(Node, discriminator="HostNode"):
:return: NMAP application installed on the Node.
:rtype: Optional[NMAP]
"""
return self.software_manager.software.get("NMAP")
return self.software_manager.software.get("nmap")
@property
def arp(self) -> Optional[ARP]:
@@ -347,7 +347,7 @@ class HostNode(Node, discriminator="HostNode"):
:return: ARP Cache for given HostNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")
def default_gateway_hello(self):
"""
@@ -379,8 +379,8 @@ class HostNode(Node, discriminator="HostNode"):
dst_port = frame.udp.dst_port
can_accept_nmap = False
if self.software_manager.software.get("NMAP"):
if self.software_manager.software["NMAP"].operating_state == ApplicationOperatingState.RUNNING:
if self.software_manager.software.get("nmap"):
if self.software_manager.software["nmap"].operating_state == ApplicationOperatingState.RUNNING:
can_accept_nmap = True
accept_nmap = can_accept_nmap and frame.payload.__class__.__name__ == "PortScanPayload"

4
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -7,7 +7,7 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node, discriminator="NetworkNode"):
class NetworkNode(Node, discriminator="network-node"):
"""
Represents an abstract base class for a network node that can receive and process network frames.
@@ -40,4 +40,4 @@ class NetworkNode(Node, discriminator="NetworkNode"):
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")

2
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
)
class WirelessRouter(Router, discriminator="wireless_router"):
class WirelessRouter(Router, discriminator="wireless-router"):
"""
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.

12
src/primaite/simulator/network/networks.py
View File

@@ -160,7 +160,7 @@ def arcd_uc2_network() -> Network:
db_client_1: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
db_client_1.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_1.run()
web_browser_1 = client_1.software_manager.software.get("WebBrowser")
web_browser_1 = client_1.software_manager.software.get("web-browser")
web_browser_1.target_url = "http://arcd.com/users/"
client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
@@ -182,10 +182,10 @@ def arcd_uc2_network() -> Network:
)
client_2.power_on()
client_2.software_manager.install(DatabaseClient)
db_client_2 = client_2.software_manager.software.get("DatabaseClient")
db_client_2 = client_2.software_manager.software.get("database-client")
db_client_2.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_2.run()
web_browser_2 = client_2.software_manager.software.get("WebBrowser")
web_browser_2 = client_2.software_manager.software.get("web-browser")
web_browser_2.target_url = "http://arcd.com/users/"
network.connect(
endpoint_b=client_2.network_interface[1],
@@ -218,7 +218,7 @@ def arcd_uc2_network() -> Network:
network.connect(endpoint_b=database_server.network_interface[1], endpoint_a=switch_1.network_interface[3])
database_server.software_manager.install(DatabaseService)
database_service: DatabaseService = database_server.software_manager.software.get("DatabaseService") # noqa
database_service: DatabaseService = database_server.software_manager.software.get("database-service") # noqa
database_service.start()
database_service.configure_backup(backup_server=IPv4Address("192.168.1.16"))
@@ -234,7 +234,7 @@ def arcd_uc2_network() -> Network:
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
database_client: DatabaseClient = web_server.software_manager.software.get("database-client")
database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
network.connect(endpoint_b=web_server.network_interface[1], endpoint_a=switch_1.network_interface[2])
database_client.run()
@@ -243,7 +243,7 @@ def arcd_uc2_network() -> Network:
web_server.software_manager.install(WebServer)
# register the web_server to a domain
dns_server_service: DNSServer = domain_controller.software_manager.software.get("DNSServer") # noqa
dns_server_service: DNSServer = domain_controller.software_manager.software.get("dns-server") # noqa
dns_server_service.dns_register("arcd.com", web_server.network_interface[1].ip_address)
# Backup Server

4
src/primaite/simulator/sim_container.py
View File

@@ -38,8 +38,8 @@ class Simulation(SimComponent):
rm.add_request("network", RequestType(func=self.network._request_manager))
# pass through domain requests to the domain object
rm.add_request("domain", RequestType(func=self.domain._request_manager))
# if 'do_nothing' is requested, just return a success
rm.add_request("do_nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
# if 'do-nothing' is requested, just return a success
rm.add_request("do-nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
return rm
def describe_state(self) -> Dict:

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

@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
return str(self)
class DatabaseClient(Application, discriminator="DatabaseClient"):
class DatabaseClient(Application, discriminator="database-client"):
"""
A DatabaseClient application.
@@ -72,7 +72,7 @@ class DatabaseClient(Application, discriminator="DatabaseClient"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for DatabaseClient."""
type: str = "DatabaseClient"
type: str = "database-client"
db_server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
@@ -97,7 +97,7 @@ class DatabaseClient(Application, discriminator="DatabaseClient"):
"""Native Client Connection for using the client directly (similar to psql in a terminal)."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
kwargs["name"] = "database-client"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
return state
class NMAP(Application, discriminator="NMAP"):
class NMAP(Application, discriminator="nmap"):
"""
A class representing the NMAP application for network scanning.
@@ -55,7 +55,7 @@ class NMAP(Application, discriminator="NMAP"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for NMAP."""
type: str = "NMAP"
type: str = "nmap"
config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema())

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

@@ -162,11 +162,11 @@ class AbstractC2(Application):
:return: An FTPClient object is successful, else None
:rtype: union[FTPClient, None]
"""
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("FTPClient")
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("ftp-client")
if ftp_client is None:
self.sys_log.warning(f"{self.__class__.__name__}: No FTPClient. Attempting to install.")
self.software_manager.install(FTPClient)
ftp_client = self.software_manager.software.get("FTPClient")
ftp_client = self.software_manager.software.get("ftp-client")
# Force start if the service is stopped.
if ftp_client.operating_state == ServiceOperatingState.STOPPED:
@@ -189,11 +189,11 @@ class AbstractC2(Application):
:return: An FTPServer object is successful, else None
:rtype: Optional[FTPServer]
"""
ftp_server: Optional[FTPServer] = self.software_manager.software.get("FTPServer")
ftp_server: Optional[FTPServer] = self.software_manager.software.get("ftp-server")
if ftp_server is None:
self.sys_log.warning(f"{self.__class__.__name__}:No FTPServer installed. Attempting to install FTPServer.")
self.software_manager.install(FTPServer)
ftp_server = self.software_manager.software.get("FTPServer")
ftp_server = self.software_manager.software.get("ftp-server")
# Force start if the service is stopped.
if ftp_server.operating_state == ServiceOperatingState.STOPPED:

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

@@ -17,7 +17,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
class C2Beacon(AbstractC2, discriminator="C2Beacon"):
class C2Beacon(AbstractC2, discriminator="c2-beacon"):
"""
C2 Beacon Application.
@@ -39,7 +39,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Beacon."""
type: str = "C2Beacon"
type: str = "c2-beacon"
c2_server_ip_address: Optional[IPV4Address] = None
keep_alive_frequency: int = 5
masquerade_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"]
@@ -54,13 +54,13 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
"The currently in use terminal session."
def __init__(self, **kwargs):
kwargs["name"] = "C2Beacon"
kwargs["name"] = "c2-beacon"
super().__init__(**kwargs)
@property
def _host_terminal(self) -> Optional[Terminal]:
"""Return the Terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("Terminal")
"""Return the terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("terminal")
if host_terminal is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a terminal on its host.")
return host_terminal
@@ -68,7 +68,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
@property
def _host_ransomware_script(self) -> RansomwareScript:
"""Return the RansomwareScript that is installed on the same machine as the C2 Beacon."""
ransomware_script: RansomwareScript = self.software_manager.software.get("RansomwareScript")
ransomware_script: RansomwareScript = self.software_manager.software.get("ransomware-script")
if ransomware_script is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find installed ransomware on its host.")
return ransomware_script
@@ -300,7 +300,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = RansomwareOpts.model_validate(payload.payload)
@@ -324,7 +324,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ransomware_script is None:
@@ -351,7 +351,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns a tuple containing Request Response returned by the Terminal execute method.
:return: Returns a tuple containing Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ftp_server is None:
@@ -372,7 +372,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
)
# Using the terminal to start the FTP Client on the remote machine.
self.terminal_session.execute(command=["service", "start", "FTPClient"])
self.terminal_session.execute(command=["service", "start", "ftp-client"])
# Need to supply to the FTP Client the C2 Beacon's host IP.
host_network_interfaces = self.software_manager.node.network_interfaces
@@ -430,7 +430,7 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
# Using the terminal to send the target data back to the C2 Beacon.
exfil_response: RequestResponse = RequestResponse.from_bool(
self.terminal_session.execute(command=["service", "FTPClient", "send", ftp_opts])
self.terminal_session.execute(command=["service", "ftp-client", "send", ftp_opts])
)
# Validating that we successfully received the target data.
@@ -472,14 +472,14 @@ class C2Beacon(AbstractC2, discriminator="C2Beacon"):
def _command_terminal(self, payload: C2Packet) -> RequestResponse:
"""
C2 Command: Terminal.
C2 Command: terminal.
Creates a request that executes a terminal command.
This request is then sent to the terminal service in order to be executed.
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = TerminalOpts.model_validate(payload.payload)

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

@@ -16,7 +16,7 @@ from primaite.simulator.system.applications.red_applications.c2 import (
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
class C2Server(AbstractC2, discriminator="C2Server"):
class C2Server(AbstractC2, discriminator="c2-server"):
"""
C2 Server Application.
@@ -37,7 +37,7 @@ class C2Server(AbstractC2, discriminator="C2Server"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Server."""
type: str = "C2Server"
type: str = "c2-server"
config: ConfigSchema = Field(default_factory=lambda: C2Server.ConfigSchema())
@@ -125,7 +125,7 @@ class C2Server(AbstractC2, discriminator="C2Server"):
return rm
def __init__(self, **kwargs):
kwargs["name"] = "C2Server"
kwargs["name"] = "c2-server"
super().__init__(**kwargs)
self.run()

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

@@ -40,13 +40,13 @@ class DataManipulationAttackStage(IntEnum):
"Signifies that the attack has failed."
class DataManipulationBot(Application, discriminator="DataManipulationBot"):
class DataManipulationBot(Application, discriminator="data-manipulation-bot"):
"""A bot that simulates a script which performs a SQL injection attack."""
class ConfigSchema(Application.ConfigSchema):
"""Configuration schema for DataManipulationBot."""
type: str = "DataManipulationBot"
type: str = "data-manipulation-bot"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "DELETE"
@@ -64,7 +64,7 @@ class DataManipulationBot(Application, discriminator="DataManipulationBot"):
"Whether to repeat attacking once finished."
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulationBot"
kwargs["name"] = "data-manipulation-bot"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -92,7 +92,7 @@ class DataManipulationBot(Application, discriminator="DataManipulationBot"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the DataManipulationBot."""
db_client = self.software_manager.software.get("DatabaseClient")
db_client = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

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

@@ -32,13 +32,13 @@ class DoSAttackStage(IntEnum):
"Attack is completed."
class DoSBot(DatabaseClient, discriminator="DoSBot"):
class DoSBot(DatabaseClient, discriminator="dos-bot"):
"""A bot that simulates a Denial of Service attack."""
class ConfigSchema(DatabaseClient.ConfigSchema):
"""ConfigSchema for DoSBot."""
type: str = "DoSBot"
type: str = "dos-bot"
target_ip_address: Optional[IPV4Address] = None
target_port: Port = PORT_LOOKUP["POSTGRES_SERVER"]
payload: Optional[str] = None
@@ -72,7 +72,7 @@ class DoSBot(DatabaseClient, discriminator="DoSBot"):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.name = "dos-bot"
self.target_ip_address = self.config.target_ip_address
self.target_port = self.config.target_port
self.payload = self.config.payload

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

@@ -14,7 +14,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class RansomwareScript(Application, discriminator="RansomwareScript"):
class RansomwareScript(Application, discriminator="ransomware-script"):
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
:ivar payload: The attack stage query payload. (Default ENCRYPT)
@@ -23,7 +23,7 @@ class RansomwareScript(Application, discriminator="RansomwareScript"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for RansomwareScript."""
type: str = "RansomwareScript"
type: str = "ransomware-script"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "ENCRYPT"
@@ -38,7 +38,7 @@ class RansomwareScript(Application, discriminator="RansomwareScript"):
"Payload String for the payload stage"
def __init__(self, **kwargs):
kwargs["name"] = "RansomwareScript"
kwargs["name"] = "ransomware-script"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -63,7 +63,7 @@ class RansomwareScript(Application, discriminator="RansomwareScript"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the Ransomware Script."""
db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
db_client: DatabaseClient = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

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

@@ -23,7 +23,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebBrowser(Application, discriminator="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
"""
Represents a web browser in the simulation environment.
@@ -33,7 +33,7 @@ class WebBrowser(Application, discriminator="WebBrowser"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for WebBrowser."""
type: str = "WebBrowser"
type: str = "web-browser"
target_url: Optional[str] = None
config: "WebBrowser.ConfigSchema" = Field(default_factory=lambda: WebBrowser.ConfigSchema())
@@ -48,7 +48,7 @@ class WebBrowser(Application, discriminator="WebBrowser"):
"""Keep a log of visited websites and information about the visit, such as response code."""
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["name"] = "web-browser"
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
@@ -108,7 +108,7 @@ class WebBrowser(Application, discriminator="WebBrowser"):
return False
# get the IP address of the domain name via DNS
dns_client: DNSClient = self.software_manager.software.get("DNSClient")
dns_client: DNSClient = self.software_manager.software.get("dns-client")
domain_exists = dns_client.check_domain_exists(target_domain=parsed_url.hostname)
# if domain does not exist, the request fails

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

@@ -60,12 +60,12 @@ class SoftwareManager:
@property
def arp(self) -> "ARP":
"""Provides access to the ARP service instance, if installed."""
return self.software.get("ARP") # noqa
return self.software.get("arp") # noqa
@property
def icmp(self) -> "ICMP":
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
return self.software.get("icmp") # noqa
def get_open_ports(self) -> List[Port]:
"""
@@ -243,7 +243,7 @@ class SoftwareManager:
:param session: The transport session the payload originates from.
"""
if payload.__class__.__name__ == "PortScanPayload":
self.software.get("NMAP").receive(payload=payload, session_id=session_id)
self.software.get("nmap").receive(payload=payload, session_id=session_id)
return
main_receiver = self.port_protocol_mapping.get((port, protocol), None)
if main_receiver:

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

@@ -15,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class ARP(Service, discriminator="ARP"):
class ARP(Service, discriminator="arp"):
"""
The ARP (Address Resolution Protocol) Service.
@@ -26,7 +26,7 @@ class ARP(Service, discriminator="ARP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ARP."""
type: str = "ARP"
type: str = "arp"
config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema())

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

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DatabaseService(Service, discriminator="DatabaseService"):
class DatabaseService(Service, discriminator="database-service"):
"""
A class for simulating a generic SQL Server service.
@@ -29,11 +29,11 @@ class DatabaseService(Service, discriminator="DatabaseService"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DatabaseService."""
type: str = "DatabaseService"
type: str = "database-service"
backup_server_ip: Optional[IPv4Address] = None
db_password: Optional[str] = None
config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DatabaseService.ConfigSchema())
backup_server_ip: IPv4Address = None
"""IP address of the backup server."""
@@ -45,7 +45,7 @@ class DatabaseService(Service, discriminator="DatabaseService"):
"""File name of latest backup."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseService"
kwargs["name"] = "database-service"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@@ -69,7 +69,7 @@ class DatabaseService(Service, discriminator="DatabaseService"):
"""
super().install()
if not self.parent.software_manager.software.get("FTPClient"):
if not self.parent.software_manager.software.get("ftp-client"):
self.parent.sys_log.info(f"{self.name}: Installing FTPClient to enable database backups")
self.parent.software_manager.install(FTPClient)
@@ -93,7 +93,7 @@ class DatabaseService(Service, discriminator="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(
@@ -127,7 +127,7 @@ class DatabaseService(Service, discriminator="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(

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

@@ -18,21 +18,21 @@ if TYPE_CHECKING:
_LOGGER = getLogger(__name__)
class DNSClient(Service, discriminator="DNSClient"):
class DNSClient(Service, discriminator="dns-client"):
"""Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
type: str = "dns-client"
dns_server: Optional[IPV4Address] = None
config: "DNSClient.ConfigSchema" = Field(default_factory=lambda: DNSClient.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DNSClient.ConfigSchema())
dns_cache: Dict[str, IPv4Address] = {}
"A dict of known mappings between domain/URLs names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSClient"
kwargs["name"] = "dns-client"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

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

@@ -14,22 +14,22 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DNSServer(Service, discriminator="DNSServer"):
class DNSServer(Service, discriminator="dns-server"):
"""Represents a DNS Server as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSServer."""
type: str = "DNSServer"
type: str = "dns-server"
domain_mapping: dict = {}
config: "DNSServer.ConfigSchema" = Field(default_factory=lambda: DNSServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DNSServer.ConfigSchema())
dns_table: Dict[str, IPv4Address] = {}
"A dict of mappings between domain names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSServer"
kwargs["name"] = "dns-server"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

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

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPClient(FTPServiceABC, discriminator="FTPClient"):
class FTPClient(FTPServiceABC, discriminator="ftp-client"):
"""
A class for simulating an FTP client service.
@@ -26,15 +26,15 @@ class FTPClient(FTPServiceABC, discriminator="FTPClient"):
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"
type: str = "ftp-client"
config: ConfigSchema = Field(default_factory=lambda: FTPClient.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["name"] = "ftp-client"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -12,7 +12,7 @@ from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPServer(FTPServiceABC, discriminator="FTPServer"):
class FTPServer(FTPServiceABC, discriminator="ftp-server"):
"""
A class for simulating an FTP server service.
@@ -20,17 +20,18 @@ class FTPServer(FTPServiceABC, discriminator="FTPServer"):
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
"""
config: "FTPServer.ConfigSchema" = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
class ConfigSchema(FTPServiceABC.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
type: str = "ftp-server"
server_password: Optional[str] = None
config: ConfigSchema = Field(default_factory=lambda: FTPServer.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["name"] = "ftp-server"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class ICMP(Service, discriminator="ICMP"):
class ICMP(Service, discriminator="icmp"):
"""
The Internet Control Message Protocol (ICMP) service.
@@ -27,14 +27,14 @@ class ICMP(Service, discriminator="ICMP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ICMP."""
type: str = "ICMP"
type: str = "icmp"
config: "ICMP.ConfigSchema" = Field(default_factory=lambda: ICMP.ConfigSchema())
request_replies: Dict = {}
def __init__(self, **kwargs):
kwargs["name"] = "ICMP"
kwargs["name"] = "icmp"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["ICMP"]
super().__init__(**kwargs)

16
src/primaite/simulator/system/services/icmp/router_icmp.py
View File

@@ -1,13 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# class RouterICMP(ICMP):
# class RouterICMP(icmp):
# """
# A class to represent a router's Internet Control Message Protocol (ICMP) handler.
# A class to represent a router's Internet Control Message Protocol (icmp) handler.
#
# :param sys_log: System log for logging network events and errors.
# :type sys_log: SysLog
# :param arp_cache: The ARP cache for resolving MAC addresses.
# :param arp_cache: The arp cache for resolving MAC addresses.
# :type arp_cache: ARPCache
# :param router: The router to which this ICMP handler belongs.
# :param router: The router to which this icmp handler belongs.
# :type router: Router
# """
#
@@ -19,7 +19,7 @@
#
# def process_icmp(self, frame: Frame, from_network_interface: NIC, is_reattempt: bool = False):
# """
# Process incoming ICMP frames based on ICMP type.
# Process incoming icmp frames based on icmp type.
#
# :param frame: The incoming frame to process.
# :param from_network_interface: The network interface where the frame is coming from.
@@ -36,13 +36,13 @@
# self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
# target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
# src_nic = self.arp.get_arp_cache_network_interface(frame.ip.src_ip_address)
# tcp_header = TCPHeader(src_port=Port["ARP"], dst_port=Port["ARP"])
# tcp_header = TCPHeader(src_port=Port["arp"], dst_port=Port["arp"])
#
# # Network Layer
# ip_packet = IPPacket(
# src_ip_address=network_interface.ip_address,
# dst_ip_address=frame.ip.src_ip_address,
# protocol=IPProtocol["ICMP"],
# protocol=IPProtocol["icmp"],
# )
# # Data Link Layer
# ethernet_header = EthernetHeader(
@@ -54,7 +54,7 @@
# identifier=frame.icmp.identifier,
# sequence=frame.icmp.sequence + 1,
# )
# payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size
# payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard icmp 32 bytes size
# frame = Frame(
# ethernet=ethernet_header,
# ip=ip_packet,

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

@@ -15,23 +15,23 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPClient(Service, discriminator="NTPClient"):
class NTPClient(Service, discriminator="ntp-client"):
"""Represents a NTP client as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPClient."""
type: str = "NTPClient"
type: str = "ntp-client"
ntp_server_ip: Optional[IPV4Address] = None
config: "NTPClient.ConfigSchema" = Field(default_factory=lambda: NTPClient.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: NTPClient.ConfigSchema())
ntp_server: Optional[IPv4Address] = None
"The NTP server the client sends requests to."
time: Optional[datetime] = None
def __init__(self, **kwargs):
kwargs["name"] = "NTPClient"
kwargs["name"] = "ntp-client"
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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

@@ -13,18 +13,18 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPServer(Service, discriminator="NTPServer"):
class NTPServer(Service, discriminator="ntp-server"):
"""Represents a NTP server as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPServer."""
type: str = "NTPServer"
type: str = "ntp-server"
config: "NTPServer.ConfigSchema" = Field(default_factory=lambda: NTPServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: NTPServer.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["name"] = "ntp-server"
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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

@@ -129,13 +129,13 @@ class RemoteTerminalConnection(TerminalClientConnection):
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
class Terminal(Service, discriminator="Terminal"):
class Terminal(Service, discriminator="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"
type: str = "terminal"
config: "Terminal.ConfigSchema" = Field(default_factory=lambda: Terminal.ConfigSchema())
@@ -143,7 +143,7 @@ class Terminal(Service, discriminator="Terminal"):
"""Dictionary of connect requests made to remote nodes."""
def __init__(self, **kwargs):
kwargs["name"] = "Terminal"
kwargs["name"] = "terminal"
kwargs["port"] = PORT_LOOKUP["SSH"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@@ -186,7 +186,7 @@ class Terminal(Service, discriminator="Terminal"):
return RequestResponse(status="failure", data={})
rm.add_request(
"node_session_remote_login",
"node-session-remote-login",
request_type=RequestType(func=_remote_login),
)

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

@@ -21,15 +21,15 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebServer(Service, discriminator="WebServer"):
class WebServer(Service, discriminator="web-server"):
"""Class used to represent a Web Server Service in simulation."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for WebServer."""
type: str = "WebServer"
type: str = "web-server"
config: "WebServer.ConfigSchema" = Field(default_factory=lambda: WebServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: WebServer.ConfigSchema())
response_codes_this_timestep: List[HttpStatusCode] = []
@@ -57,7 +57,7 @@ class WebServer(Service, discriminator="WebServer"):
return super().pre_timestep(timestep)
def __init__(self, **kwargs):
kwargs["name"] = "WebServer"
kwargs["name"] = "web-server"
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
@@ -146,7 +146,7 @@ class WebServer(Service, discriminator="WebServer"):
def _establish_db_connection(self) -> None:
"""Establish a connection to db."""
db_client = self.software_manager.software.get("DatabaseClient")
db_client = self.software_manager.software.get("database-client")
self.db_connection: DatabaseClientConnection = db_client.get_new_connection()
def send(