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

Merge branch '4.0.0a1-dev' into feature/3075_Migrate_notebooks_to_MilPac_(Core_changes)

Browse Source
This commit is contained in:
Nick Todd
2025-02-06 09:17:34 +00:00
parent 7b2a9c1d2a 0d1edf0362
commit 0d8672cdba
210 changed files with 3022 additions and 2983 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

@@ -172,7 +172,7 @@
"\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application, identifier=\"MSPaint\"):\n",
"class MSPaint(Application, discriminator=\"MSPaint\"):\n",
" def describe_state(self):\n",
" return super().describe_state()"
]

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

@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
..code::python
class WebBrowser(Application, identifier="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:

2
src/primaite/simulator/network/container.py
View File

@@ -180,7 +180,7 @@ class Network(SimComponent):
table.align = "l"
table.title = "Nodes"
for node in self.nodes.values():
table.add_row((node.config.hostname, type(node)._identifier, node.operating_state.name))
table.add_row((node.config.hostname, type(node)._discriminator, node.operating_state.name))
print(table)
if ip_addresses:

24
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, identifier="your_name"):
class YourNodeAdder(NetworkNodeAdder, discriminator="your-name"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
property_1 : str
property_2 : int
@@ -40,8 +40,8 @@ class NetworkNodeAdder(BaseModel):
"""
Base schema for node adders.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The identifier is used
by the from_config method to select the correct node adder at runtime.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The discriminator is
used by the from_config method to select the correct node adder at runtime.
"""
model_config = ConfigDict(extra="forbid")
@@ -50,20 +50,20 @@ class NetworkNodeAdder(BaseModel):
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str], **kwargs: Any) -> None:
"""
Register a network node adder class.
:param identifier: Unique name for the node adder to use for matching against primaite config entries.
:type identifier: str
:param discriminator: Unique name for the node adder to use for matching against primaite config entries.
:type discriminator: str
:raises ValueError: When attempting to register a name that is already reserved.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate node adder {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate node adder {discriminator}")
cls._registry[discriminator] = cls
@classmethod
@abstractmethod
@@ -99,13 +99,13 @@ class NetworkNodeAdder(BaseModel):
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
class OfficeLANAdder(NetworkNodeAdder, identifier="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

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

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump()
class UserManager(Service, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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."""
@@ -1522,15 +1522,18 @@ class Node(SimComponent, ABC):
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
_identifier: ClassVar[str] = "unknown"
"""Identifier for this particular class, used for printing and logging. Each subclass redefines this."""
# TODO: this should not be set for abstract classes.
_discriminator: ClassVar[str]
"""discriminator for this particular class, used for printing and logging. Each subclass redefines this."""
class ConfigSchema(BaseModel, ABC):
"""Configuration Schema for Node based classes."""
model_config = ConfigDict(arbitrary_types_allowed=True)
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
"""Configure pydantic to allow arbitrary types, let the instance have attributes not present in the model."""
type: str
hostname: str
"The node hostname on the network."
@@ -1569,6 +1572,8 @@ class Node(SimComponent, ABC):
operating_state: Any = None
users: Any = None # Temporary to appease "extra=forbid"
config: ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema())
"""Configuration items within Node"""
@@ -1586,22 +1591,22 @@ class Node(SimComponent, ABC):
obj = cls(config=cls.ConfigSchema(**config))
return obj
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a node type.
:param identifier: Uniquely specifies an node class by name. Used for finding items by config.
:type identifier: str
:param discriminator: Uniquely specifies an node class by name. Used for finding items by config.
:type discriminator: str
:raises ValueError: When attempting to register an node with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
identifier = identifier.lower()
if identifier in cls._registry:
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
cls._identifier = identifier
discriminator = discriminator.lower()
if discriminator in cls._registry:
raise ValueError(f"Tried to define new node {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
cls._discriminator = discriminator
def __init__(self, **kwargs):
"""
@@ -1635,17 +1640,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]:
"""

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

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import ClassVar, Dict
from typing import ClassVar, Dict, Literal
from pydantic import Field
@@ -7,7 +7,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(HostNode, identifier="computer"):
class Computer(HostNode, discriminator="computer"):
"""
A basic Computer class.
@@ -35,11 +35,12 @@ class Computer(HostNode, identifier="computer"):
* Web Browser
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient}
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "ftp-client": FTPClient}
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Computer class."""
type: Literal["computer"] = "computer"
hostname: str = "Computer"
config: ConfigSchema = Field(default_factory=lambda: Computer.ConfigSchema())

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

@@ -2,7 +2,7 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Optional
from typing import Any, ClassVar, Dict, Literal, Optional
from pydantic import Field
@@ -267,7 +267,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, identifier="HostNode"):
class HostNode(Node, discriminator="host-node"):
"""
Represents a host node in the network.
@@ -314,14 +314,14 @@ class HostNode(Node, identifier="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."""
@@ -333,9 +333,14 @@ class HostNode(Node, identifier="HostNode"):
class ConfigSchema(Node.ConfigSchema):
"""Configuration Schema for HostNode class."""
type: Literal["host-node"]
hostname: str = "HostNode"
subnet_mask: IPV4Address = "255.255.255.0"
ip_address: IPV4Address
services: Any = None # temporarily unset to appease extra="forbid"
applications: Any = None # temporarily unset to appease extra="forbid"
folders: Any = None # temporarily unset to appease extra="forbid"
network_interfaces: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema())
@@ -351,7 +356,7 @@ class HostNode(Node, identifier="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]:
@@ -361,7 +366,7 @@ class HostNode(Node, identifier="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):
"""
@@ -393,8 +398,8 @@ class HostNode(Node, identifier="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"

8
src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -1,11 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Literal
from pydantic import Field
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
class Server(HostNode, identifier="server"):
class Server(HostNode, discriminator="server"):
"""
A basic Server class.
@@ -36,12 +38,13 @@ class Server(HostNode, identifier="server"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Server class."""
type: Literal["server"] = "server"
hostname: str = "server"
config: ConfigSchema = Field(default_factory=lambda: Server.ConfigSchema())
class Printer(HostNode, identifier="printer"):
class Printer(HostNode, discriminator="printer"):
"""Printer? I don't even know her!."""
# TODO: Implement printer-specific behaviour
@@ -49,6 +52,7 @@ class Printer(HostNode, identifier="printer"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Printer class."""
type: Literal["printer"] = "printer"
hostname: str = "printer"
config: ConfigSchema = Field(default_factory=lambda: Printer.ConfigSchema())

5
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Dict, Final, Union
from typing import Dict, Final, Literal, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -26,7 +26,7 @@ DMZ_PORT_ID: Final[int] = 3
"""The Firewall port ID of the DMZ port."""
class Firewall(Router, identifier="firewall"):
class Firewall(Router, discriminator="firewall"):
"""
A Firewall class that extends the functionality of a Router.
@@ -103,6 +103,7 @@ class Firewall(Router, identifier="firewall"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for Firewall 'Nodes' within PrimAITE."""
type: Literal["firewall"] = "firewall"
hostname: str = "firewall"
num_ports: int = 0

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

@@ -1,13 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import Optional
from typing import Any, Optional
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node, identifier="NetworkNode"):
class NetworkNode(Node, discriminator="network-node"):
"""
Represents an abstract base class for a network node that can receive and process network frames.
@@ -16,6 +16,11 @@ class NetworkNode(Node, identifier="NetworkNode"):
provide functionality for receiving and processing frames received on their network interfaces.
"""
class ConfigSchema(Node.ConfigSchema):
"""Config schema for Node baseclass."""
num_ports: Any = None # temporarily unset to appease extra="forbid"
@abstractmethod
def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface):
"""
@@ -40,4 +45,4 @@ class NetworkNode(Node, identifier="NetworkNode"):
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")

29
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -4,7 +4,7 @@ from __future__ import annotations
import secrets
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
class Router(NetworkNode, identifier="router"):
class Router(NetworkNode, discriminator="router"):
"""
Represents a network router, managing routing and forwarding of IP packets across network interfaces.
@@ -1204,8 +1204,13 @@ class Router(NetworkNode, identifier="router"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Routers."""
type: Literal["router"] = "router"
hostname: str = "router"
num_ports: int = 5
acl: Any = None # temporarily unset to appease extra="forbid"
routes: Any = None # temporarily unset to appease extra="forbid"
ports: Any = None # temporarily unset to appease extra="forbid"
default_route: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: Router.ConfigSchema())
@@ -1625,16 +1630,20 @@ class Router(NetworkNode, identifier="router"):
:return: Configured router.
:rtype: Router
"""
ports = config.pop("ports", None)
acl = config.pop("acl", None)
routes = config.pop("routes", None)
default_route = config.pop("default_route", None)
router = Router(config=Router.ConfigSchema(**config))
if "ports" in config:
for port_num, port_cfg in config["ports"].items():
if ports:
for port_num, port_cfg in ports.items():
router.configure_port(
port=port_num,
ip_address=port_cfg["ip_address"],
subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")),
)
if "acl" in config:
for r_num, r_cfg in config["acl"].items():
if acl:
for r_num, r_cfg in acl.items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
@@ -1646,16 +1655,16 @@ class Router(NetworkNode, identifier="router"):
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
if "routes" in config:
for route in config.get("routes"):
if routes:
for route in routes:
router.route_table.add_route(
address=IPv4Address(route.get("address")),
subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")),
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
metric=float(route.get("metric", 0)),
)
if "default_route" in config:
next_hop_ip_address = config["default_route"].get("next_hop_ip_address", None)
if default_route:
next_hop_ip_address = default_route.get("next_hop_ip_address", None)
if next_hop_ip_address:
router.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address)
router.operating_state = (

5
src/primaite/simulator/network/hardware/nodes/network/switch.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Optional
from typing import Dict, Literal, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
@@ -88,7 +88,7 @@ class SwitchPort(WiredNetworkInterface):
return False
class Switch(NetworkNode, identifier="switch"):
class Switch(NetworkNode, discriminator="switch"):
"""A class representing a Layer 2 network switch."""
network_interfaces: Dict[str, SwitchPort] = {}
@@ -101,6 +101,7 @@ class Switch(NetworkNode, identifier="switch"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Switch nodes within PrimAITE."""
type: Literal["switch"] = "switch"
hostname: str = "Switch"
num_ports: int = 24
"The number of ports on the switch. Default is 24."

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

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Literal, Optional, Union
from pydantic import Field, validate_call
@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
)
class WirelessRouter(Router, identifier="wireless_router"):
class WirelessRouter(Router, discriminator="wireless-router"):
"""
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.
@@ -126,10 +126,13 @@ class WirelessRouter(Router, identifier="wireless_router"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for WirelessRouter nodes within PrimAITE."""
type: Literal["wireless-router"] = "wireless-router"
hostname: str = "WirelessRouter"
airspace: AirSpace
num_ports: int = 0
router_interface: Any = None # temporarily unset to appease extra="forbid"
wireless_access_point: Any = None # temporarily unset to appease extra="forbid"
airspace: AirSpace
config: ConfigSchema = Field(default_factory=lambda: WirelessRouter.ConfigSchema())
def __init__(self, **kwargs):
@@ -137,7 +140,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
self.connect_nic(
WirelessAccessPoint(
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=kwargs["config"].airspace
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=self.airspace
)
)
@@ -236,7 +239,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
)
@classmethod
def from_config(cls, config: Dict, **kwargs) -> "WirelessRouter":
def from_config(cls, config: Dict, airspace: AirSpace) -> "WirelessRouter":
"""Generate the wireless router from config.
Schema:
@@ -263,7 +266,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
:return: WirelessRouter instance.
:rtype: WirelessRouter
"""
router = cls(config=cls.ConfigSchema(**config))
router = cls(config=cls.ConfigSchema(**config), airspace=airspace)
router.operating_state = (
NodeOperatingState.ON if not (p := config.get("operating_state")) else NodeOperatingState[p.upper()]
)

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

@@ -170,13 +170,13 @@ def arcd_uc2_network() -> Network:
client_1.power_on()
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
db_client_1: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
db_client_1: DatabaseClient = client_1.software_manager.software.get("database-client")
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")
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("data-manipulation-bot")
db_manipulation_bot.configure(
server_ip_address=IPv4Address("192.168.1.14"),
payload="DELETE",
@@ -199,10 +199,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],
@@ -244,7 +244,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"))
@@ -264,7 +264,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()
@@ -273,7 +273,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:

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

@@ -53,20 +53,20 @@ class Application(IOSoftware, ABC):
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: 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: Optional[str]
:param discriminator: Uniquely specifies an application class by name. Used for finding items by config.
:type discriminator: Optional[str]
:raises ValueError: When attempting to register an application with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator 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
if discriminator in cls._registry:
raise ValueError(f"Tried to define new application {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "Application":

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

@@ -37,7 +37,7 @@ class DatabaseClientConnection(BaseModel):
@property
def client(self) -> Optional[DatabaseClient]:
"""The DatabaseClient that holds this connection."""
return self.parent_node.software_manager.software.get("DatabaseClient")
return self.parent_node.software_manager.software.get("database-client")
def query(self, sql: str) -> bool:
"""
@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
return str(self)
class DatabaseClient(Application, identifier="DatabaseClient"):
class DatabaseClient(Application, discriminator="database-client"):
"""
A DatabaseClient application.
@@ -72,7 +72,7 @@ class DatabaseClient(Application, identifier="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, identifier="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)

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

@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
return state
class NMAP(Application, identifier="NMAP"):
class NMAP(Application, discriminator="nmap"):
"""
A class representing the NMAP application for network scanning.
@@ -55,7 +55,7 @@ class NMAP(Application, identifier="NMAP"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for NMAP."""
type: str = "NMAP"
type: str = "nmap"
config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema())
@@ -70,7 +70,7 @@ class NMAP(Application, identifier="NMAP"):
}
def __init__(self, **kwargs):
kwargs["name"] = "NMAP"
kwargs["name"] = "nmap"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)

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, identifier="C2Beacon"):
class C2Beacon(AbstractC2, discriminator="c2-beacon"):
"""
C2 Beacon Application.
@@ -39,7 +39,7 @@ class C2Beacon(AbstractC2, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="C2Server"):
class C2Server(AbstractC2, discriminator="c2-server"):
"""
C2 Server Application.
@@ -37,7 +37,7 @@ class C2Server(AbstractC2, identifier="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, identifier="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, identifier="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"
@@ -65,7 +65,7 @@ class DataManipulationBot(Application, identifier="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"]
@@ -94,7 +94,7 @@ class DataManipulationBot(Application, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
"""
Represents a web browser in the simulation environment.
@@ -33,7 +33,7 @@ class WebBrowser(Application, identifier="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, identifier="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, identifier="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]:
"""
@@ -244,7 +244,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:

6
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, identifier="ARP"):
class ARP(Service, discriminator="arp"):
"""
The ARP (Address Resolution Protocol) Service.
@@ -26,14 +26,14 @@ class ARP(Service, identifier="ARP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ARP."""
type: str = "ARP"
type: str = "arp"
config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema())
arp: Dict[IPV4Address, ARPEntry] = {}
def __init__(self, **kwargs):
kwargs["name"] = "ARP"
kwargs["name"] = "arp"
kwargs["port"] = PORT_LOOKUP["ARP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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, identifier="DatabaseService"):
class DatabaseService(Service, discriminator="database-service"):
"""
A class for simulating a generic SQL Server service.
@@ -29,12 +29,12 @@ class DatabaseService(Service, identifier="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
"""Password that needs to be provided by clients if they want to connect to the DatabaseService."""
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."""
@@ -46,7 +46,7 @@ class DatabaseService(Service, identifier="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)
@@ -70,7 +70,7 @@ class DatabaseService(Service, identifier="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)
@@ -94,7 +94,7 @@ class DatabaseService(Service, identifier="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(
@@ -128,7 +128,7 @@ class DatabaseService(Service, identifier="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(

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

@@ -18,24 +18,22 @@ if TYPE_CHECKING:
_LOGGER = getLogger(__name__)
class DNSClient(Service, identifier="DNSClient"):
class DNSClient(Service, discriminator="dns-client"):
"""Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
dns_server: Optional[IPv4Address] = None
type: str = "dns-client"
dns_server: Optional[IPV4Address] = None
"The DNS Server the client sends requests to."
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, identifier="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, identifier="FTPClient"):
class FTPClient(FTPServiceABC, discriminator="ftp-client"):
"""
A class for simulating an FTP client service.
@@ -26,15 +26,15 @@ class FTPClient(FTPServiceABC, identifier="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, identifier="FTPServer"):
class FTPServer(FTPServiceABC, discriminator="ftp-server"):
"""
A class for simulating an FTP server service.
@@ -23,15 +23,14 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"):
class ConfigSchema(FTPServiceABC.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
type: str = "ftp-server"
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
config: ConfigSchema = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
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, identifier="ICMP"):
class ICMP(Service, discriminator="icmp"):
"""
The Internet Control Message Protocol (ICMP) service.
@@ -27,14 +27,14 @@ class ICMP(Service, identifier="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,

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

@@ -9,20 +9,20 @@ from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPClient(Service, identifier="NTPClient"):
class NTPClient(Service, discriminator="ntp-client"):
"""Represents a NTP client as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPClient."""
type: str = "NTPClient"
ntp_server_ip: Optional[IPv4Address] = None
type: str = "ntp-client"
ntp_server_ip: Optional[IPV4Address] = None
"The NTP server the client sends requests to."
config: ConfigSchema = Field(default_factory=lambda: NTPClient.ConfigSchema())
@@ -30,7 +30,7 @@ class NTPClient(Service, identifier="NTPClient"):
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, identifier="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)

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

@@ -61,22 +61,22 @@ class Service(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a hostnode type.
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type identifier: str
:param discriminator: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type discriminator: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower()
if identifier in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
discriminator = discriminator.lower()
if discriminator in cls._registry:
raise ValueError(f"Tried to define new hostnode {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "Service":

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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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, identifier="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(