From c1a5a26ffca6beb906bad50c0855d6a6dbc9252c Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 4 Feb 2025 10:21:56 +0000 Subject: [PATCH 1/5] #2887 - Actioning review comments --- src/primaite/simulator/core.py | 2 +- .../simulator/network/hardware/base.py | 8 ++++---- .../network/hardware/nodes/host/computer.py | 4 ++-- .../network/hardware/nodes/host/host_node.py | 4 ++-- .../network/hardware/nodes/host/server.py | 8 ++++---- .../network/hardware/nodes/network/firewall.py | 4 ++-- .../network/hardware/nodes/network/router.py | 18 +++++++++--------- .../network/hardware/nodes/network/switch.py | 4 ++-- .../hardware/nodes/network/wireless_router.py | 4 ++-- .../system/services/dns/dns_client.py | 4 ++-- .../system/services/ftp/ftp_server.py | 5 ++--- .../system/services/ntp/ntp_client.py | 4 +--- 12 files changed, 33 insertions(+), 36 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index dc4ae73b..567a0493 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -3,7 +3,7 @@ """Core of the PrimAITE Simulator.""" import warnings from abc import abstractmethod -from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union +from typing import Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union from uuid import uuid4 from prettytable import PrettyTable diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 0564e1f3..36623a6f 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1525,16 +1525,13 @@ class Node(SimComponent, ABC): _identifier: ClassVar[str] = "unknown" """Identifier for this particular class, used for printing and logging. Each subclass redefines this.""" - config: Node.ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema()) - """Configuration items within Node""" - class ConfigSchema(BaseModel, ABC): """Configuration Schema for Node based classes.""" model_config = ConfigDict(arbitrary_types_allowed=True) """Configure pydantic to allow arbitrary types, let the instance have attributes not present in the model.""" - hostname: str = "default" + hostname: str "The node hostname on the network." revealed_to_red: bool = False @@ -1572,6 +1569,9 @@ class Node(SimComponent, ABC): operating_state: Any = None + config: ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema()) + """Configuration items within Node""" + @property def dns_server(self) -> Optional[IPv4Address]: """Convenience method to access the dns_server IP.""" diff --git a/src/primaite/simulator/network/hardware/nodes/host/computer.py b/src/primaite/simulator/network/hardware/nodes/host/computer.py index 85857a44..a0450443 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/host/computer.py @@ -37,11 +37,11 @@ class Computer(HostNode, identifier="computer"): SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient} - config: "Computer.ConfigSchema" = Field(default_factory=lambda: Computer.ConfigSchema()) - class ConfigSchema(HostNode.ConfigSchema): """Configuration Schema for Computer class.""" hostname: str = "Computer" + config: ConfigSchema = Field(default_factory=lambda: Computer.ConfigSchema()) + pass diff --git a/src/primaite/simulator/network/hardware/nodes/host/host_node.py b/src/primaite/simulator/network/hardware/nodes/host/host_node.py index 424e39f1..1aa482db 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py +++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py @@ -330,8 +330,6 @@ class HostNode(Node, identifier="HostNode"): network_interface: Dict[int, NIC] = {} "The NICs on the node by port id." - config: HostNode.ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema()) - class ConfigSchema(Node.ConfigSchema): """Configuration Schema for HostNode class.""" @@ -339,6 +337,8 @@ class HostNode(Node, identifier="HostNode"): subnet_mask: IPV4Address = "255.255.255.0" ip_address: IPV4Address + config: ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema()) + def __init__(self, **kwargs): super().__init__(**kwargs) self.connect_nic(NIC(ip_address=kwargs["config"].ip_address, subnet_mask=kwargs["config"].subnet_mask)) diff --git a/src/primaite/simulator/network/hardware/nodes/host/server.py b/src/primaite/simulator/network/hardware/nodes/host/server.py index bdf4e8c2..3a9fc2f9 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/server.py +++ b/src/primaite/simulator/network/hardware/nodes/host/server.py @@ -33,22 +33,22 @@ class Server(HostNode, identifier="server"): * Web Browser """ - config: "Server.ConfigSchema" = Field(default_factory=lambda: Server.ConfigSchema()) - class ConfigSchema(HostNode.ConfigSchema): """Configuration Schema for Server class.""" hostname: str = "server" + config: ConfigSchema = Field(default_factory=lambda: Server.ConfigSchema()) + class Printer(HostNode, identifier="printer"): """Printer? I don't even know her!.""" # TODO: Implement printer-specific behaviour - config: "Printer.ConfigSchema" = Field(default_factory=lambda: Printer.ConfigSchema()) - class ConfigSchema(HostNode.ConfigSchema): """Configuration Schema for Printer class.""" hostname: str = "printer" + + config: ConfigSchema = Field(default_factory=lambda: Printer.ConfigSchema()) diff --git a/src/primaite/simulator/network/hardware/nodes/network/firewall.py b/src/primaite/simulator/network/hardware/nodes/network/firewall.py index 99dd48c4..c4ddea58 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/firewall.py +++ b/src/primaite/simulator/network/hardware/nodes/network/firewall.py @@ -100,14 +100,14 @@ class Firewall(Router, identifier="firewall"): _identifier: str = "firewall" - config: "Firewall.ConfigSchema" = Field(default_factory=lambda: Firewall.ConfigSchema()) - class ConfigSchema(Router.ConfigSchema): """Configuration Schema for Firewall 'Nodes' within PrimAITE.""" hostname: str = "firewall" num_ports: int = 0 + config: ConfigSchema = Field(default_factory=lambda: Firewall.ConfigSchema()) + def __init__(self, **kwargs): if not kwargs.get("sys_log"): kwargs["sys_log"] = SysLog(kwargs["config"].hostname) diff --git a/src/primaite/simulator/network/hardware/nodes/network/router.py b/src/primaite/simulator/network/hardware/nodes/network/router.py index dd32fa31..7138cf4f 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/router.py @@ -7,7 +7,7 @@ from ipaddress import IPv4Address, IPv4Network from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union from prettytable import MARKDOWN, PrettyTable -from pydantic import validate_call +from pydantic import Field, validate_call from primaite.interface.request import RequestResponse from primaite.simulator.core import RequestManager, RequestType, SimComponent @@ -1201,6 +1201,14 @@ class Router(NetworkNode, identifier="router"): RouteTable, RouterARP, and RouterICMP services. """ + class ConfigSchema(NetworkNode.ConfigSchema): + """Configuration Schema for Routers.""" + + hostname: str = "router" + num_ports: int = 5 + + config: ConfigSchema = Field(default_factory=lambda: Router.ConfigSchema()) + SYSTEM_SOFTWARE: ClassVar[Dict] = { "UserSessionManager": UserSessionManager, "UserManager": UserManager, @@ -1214,14 +1222,6 @@ class Router(NetworkNode, identifier="router"): acl: AccessControlList route_table: RouteTable - config: "Router.ConfigSchema" - - class ConfigSchema(NetworkNode.ConfigSchema): - """Configuration Schema for Routers.""" - - hostname: str = "router" - num_ports: int = 5 - def __init__(self, **kwargs): if not kwargs.get("sys_log"): kwargs["sys_log"] = SysLog(kwargs["config"].hostname) diff --git a/src/primaite/simulator/network/hardware/nodes/network/switch.py b/src/primaite/simulator/network/hardware/nodes/network/switch.py index 3cb335f7..54e1d7ef 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/network/switch.py @@ -98,8 +98,6 @@ class Switch(NetworkNode, identifier="switch"): mac_address_table: Dict[str, SwitchPort] = {} "A MAC address table mapping destination MAC addresses to corresponding SwitchPorts." - config: "Switch.ConfigSchema" = Field(default_factory=lambda: Switch.ConfigSchema()) - class ConfigSchema(NetworkNode.ConfigSchema): """Configuration Schema for Switch nodes within PrimAITE.""" @@ -107,6 +105,8 @@ class Switch(NetworkNode, identifier="switch"): num_ports: int = 24 "The number of ports on the switch. Default is 24." + config: ConfigSchema = Field(default_factory=lambda: Switch.ConfigSchema()) + def __init__(self, **kwargs): super().__init__(**kwargs) for i in range(1, kwargs["config"].num_ports + 1): diff --git a/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py b/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py index 348c2aaa..2beb03d6 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py @@ -123,8 +123,6 @@ class WirelessRouter(Router, identifier="wireless_router"): network_interfaces: Dict[str, Union[RouterInterface, WirelessAccessPoint]] = {} network_interface: Dict[int, Union[RouterInterface, WirelessAccessPoint]] = {} - config: "WirelessRouter.ConfigSchema" = Field(default_factory=lambda: WirelessRouter.ConfigSchema()) - class ConfigSchema(Router.ConfigSchema): """Configuration Schema for WirelessRouter nodes within PrimAITE.""" @@ -132,6 +130,8 @@ class WirelessRouter(Router, identifier="wireless_router"): airspace: AirSpace num_ports: int = 0 + config: ConfigSchema = Field(default_factory=lambda: WirelessRouter.ConfigSchema()) + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 4a1fd292..4f926a25 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/src/primaite/simulator/system/services/dns/dns_client.py @@ -25,12 +25,12 @@ class DNSClient(Service, identifier="DNSClient"): """ConfigSchema for DNSClient.""" type: str = "DNSClient" - dns_server: Optional[IPV4Address] = None dns_server: Optional[IPv4Address] = None "The DNS Server the client sends requests to." - 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." diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index ebb93a7b..147d2dbb 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -20,17 +20,16 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"): RFC 959: https://datatracker.ietf.org/doc/html/rfc959 """ - config: "FTPServer.ConfigSchema" = Field(default_factory=lambda: FTPServer.ConfigSchema()) - class ConfigSchema(FTPServiceABC.ConfigSchema): """ConfigSchema for FTPServer.""" type: str = "FTPServer" - server_password: Optional[str] = None server_password: Optional[str] = None """Password needed to connect to FTP server. Default is None.""" + config: ConfigSchema = Field(default_factory=lambda: FTPServer.ConfigSchema()) + def __init__(self, **kwargs): kwargs["name"] = "FTPServer" kwargs["port"] = PORT_LOOKUP["FTP"] diff --git a/src/primaite/simulator/system/services/ntp/ntp_client.py b/src/primaite/simulator/system/services/ntp/ntp_client.py index 72e8f6c0..7c3efd25 100644 --- a/src/primaite/simulator/system/services/ntp/ntp_client.py +++ b/src/primaite/simulator/system/services/ntp/ntp_client.py @@ -9,7 +9,6 @@ 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__) @@ -22,12 +21,11 @@ class NTPClient(Service, identifier="NTPClient"): """ConfigSchema for NTPClient.""" type: str = "NTPClient" - ntp_server_ip: Optional[IPV4Address] = None ntp_server_ip: Optional[IPv4Address] = None "The NTP server the client sends requests to." - config: "NTPClient.ConfigSchema" = Field(default_factory=lambda: NTPClient.ConfigSchema()) + config: ConfigSchema = Field(default_factory=lambda: NTPClient.ConfigSchema()) time: Optional[datetime] = None From 961136fb4213d2274800760ff220d79bf5f59382 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 4 Feb 2025 10:41:51 +0000 Subject: [PATCH 2/5] #2887 - Updates to extensible_nodes.rst --- docs/source/how_to_guides/extensible_nodes.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/how_to_guides/extensible_nodes.rst b/docs/source/how_to_guides/extensible_nodes.rst index 78ee550e..6651b618 100644 --- a/docs/source/how_to_guides/extensible_nodes.rst +++ b/docs/source/how_to_guides/extensible_nodes.rst @@ -53,4 +53,4 @@ class Router(NetworkNode, identifier="router"): Changes to YAML file. ===================== -Nodes defined within configuration YAML files for use with PrimAITE 3.X should still be compatible following these changes. +While effort has been made to ensure that nodes defined within configuration YAML files for use with PrimAITE 3.X remain compatible with PrimAITE v4+, it is encouraged to review for minor changes needed. From 05946431ca5eac277e6269a8d2c172cd0b084cf8 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 4 Feb 2025 11:19:13 +0000 Subject: [PATCH 3/5] #2887 - Correct type in documentation --- docs/source/how_to_guides/extensible_nodes.rst | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/source/how_to_guides/extensible_nodes.rst b/docs/source/how_to_guides/extensible_nodes.rst index 6651b618..043d0f06 100644 --- a/docs/source/how_to_guides/extensible_nodes.rst +++ b/docs/source/how_to_guides/extensible_nodes.rst @@ -18,8 +18,7 @@ Node classes all inherit from the base Node Class, though new classes should inh The use of an `__init__` method is not necessary, as configurable variables for the class should be specified within the `config` of the class, and passed at run time via your YAML configuration using the `from_config` method. - -An example of how additional Node classes is below, taken from `router.py` withing PrimAITE. +An example of how additional Node classes is below, taken from `router.py` within PrimAITE. .. code-block:: Python From 99e38fbbc2277c981abd904be44ea9311843cfc9 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 4 Feb 2025 14:25:26 +0000 Subject: [PATCH 4/5] #2887 - Removal of un-necessary code and cleanup following review comments --- src/primaite/simulator/network/creation.py | 3 +-- src/primaite/simulator/network/hardware/base.py | 8 -------- .../simulator/network/hardware/nodes/network/switch.py | 2 +- 3 files changed, 2 insertions(+), 11 deletions(-) diff --git a/src/primaite/simulator/network/creation.py b/src/primaite/simulator/network/creation.py index 5e0d0ce8..009ac861 100644 --- a/src/primaite/simulator/network/creation.py +++ b/src/primaite/simulator/network/creation.py @@ -167,7 +167,6 @@ class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"): # Optionally include a router in the LAN if config.include_router: default_gateway = IPv4Address(f"192.168.{config.subnet_base}.1") - # router = Router(hostname=f"router_{config.lan_name}", start_up_duration=0) router = Router.from_config( config={"hostname": f"router_{config.lan_name}", "type": "router", "start_up_duration": 0} ) @@ -230,7 +229,7 @@ class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"): "type": "computer", "hostname": f"pc_{i}_{config.lan_name}", "ip_address": f"192.168.{config.subnet_base}.{i+config.pcs_ip_block_start-1}", - "default_gateway": "192.168.10.1", + "default_gateway": default_gateway, "start_up_duration": 0, } pc = Computer.from_config(config=pc_cfg) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 36623a6f..29d46164 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -2242,10 +2242,6 @@ class Node(SimComponent, ABC): for app_id in self.applications: self.applications[app_id].close() - # Turn off all processes in the node - # for process_id in self.processes: - # self.processes[process_id] - def _start_up_actions(self): """Actions to perform when the node is starting up.""" # Turn on all the services in the node @@ -2258,10 +2254,6 @@ class Node(SimComponent, ABC): print(f"Starting application:{self.applications[app_id].config.type}") self.applications[app_id].run() - # Turn off all processes in the node - # for process_id in self.processes: - # self.processes[process_id] - def _install_system_software(self) -> None: """Preinstall required software.""" for _, software_class in self.SYSTEM_SOFTWARE.items(): diff --git a/src/primaite/simulator/network/hardware/nodes/network/switch.py b/src/primaite/simulator/network/hardware/nodes/network/switch.py index 54e1d7ef..8a9fdb24 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/network/switch.py @@ -109,7 +109,7 @@ class Switch(NetworkNode, identifier="switch"): def __init__(self, **kwargs): super().__init__(**kwargs) - for i in range(1, kwargs["config"].num_ports + 1): + for i in range(1, self.config.num_ports + 1): self.connect_nic(SwitchPort()) def _install_system_software(self): From 51bb3f5b07deede2506bb242eb8586e8307aedf4 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 4 Feb 2025 14:26:55 +0000 Subject: [PATCH 5/5] #2887 - Removal of un-necessary print statement that was used for debugging --- src/primaite/simulator/network/hardware/base.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 29d46164..6543d793 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -2250,8 +2250,6 @@ class Node(SimComponent, ABC): # Turn on all the applications in the node for app_id in self.applications: - print(app_id) - print(f"Starting application:{self.applications[app_id].config.type}") self.applications[app_id].run() def _install_system_software(self) -> None: