diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 9684e9e8..f6b5b62c 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -18,7 +18,7 @@ from primaite.game.agent.scripted_agents.tap001 import TAP001 from primaite.game.science import graph_has_cycle, topological_sort from primaite.simulator import SIM_OUTPUT from primaite.simulator.network.creation import NetworkNodeAdder -from primaite.simulator.network.hardware.base import NetworkInterface, NodeOperatingState, UserManager +from primaite.simulator.network.hardware.base import NetworkInterface, Node, NodeOperatingState, UserManager from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.network.hardware.nodes.host.host_node import HostNode, NIC from primaite.simulator.network.hardware.nodes.host.server import Printer, Server @@ -280,21 +280,7 @@ class PrimaiteGame: new_node = None # Handle extended nodes - if n_type.lower() in HostNode._registry: - new_node = HostNode._registry[n_type]( - hostname=node_cfg["hostname"], - ip_address=node_cfg["ip_address"], - subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")), - default_gateway=node_cfg.get("default_gateway"), - dns_server=node_cfg.get("dns_server", None), - operating_state=NodeOperatingState.ON - if not (p := node_cfg.get("operating_state")) - else NodeOperatingState[p.upper()], - ) - elif n_type in NetworkNode._registry: - new_node = NetworkNode._registry[n_type](**node_cfg) - # Default PrimAITE nodes - elif n_type == "computer": + if n_type == "computer": new_node = Computer( hostname=node_cfg["hostname"], ip_address=node_cfg["ip_address"], @@ -339,6 +325,20 @@ class PrimaiteGame: if not (p := node_cfg.get("operating_state")) else NodeOperatingState[p.upper()], ) + elif n_type.lower() in Node._registry: + new_node = HostNode._registry[n_type]( + hostname=node_cfg["hostname"], + ip_address=node_cfg["ip_address"], + subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")), + default_gateway=node_cfg.get("default_gateway"), + dns_server=node_cfg.get("dns_server", None), + operating_state=NodeOperatingState.ON + if not (p := node_cfg.get("operating_state")) + else NodeOperatingState[p.upper()], + ) + elif n_type in NetworkNode._registry: + new_node = NetworkNode._registry[n_type](**node_cfg) + # Default PrimAITE nodes else: msg = f"invalid node type {n_type} in config" _LOGGER.error(msg) diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 6e019f32..1082e172 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -179,9 +179,8 @@ class Network(SimComponent): table.set_style(MARKDOWN) table.align = "l" table.title = "Nodes" - for node_type, nodes in nodes_type_map.items(): - for node in nodes: - table.add_row([node.hostname, node_type, node.operating_state.name]) + for node in self.nodes.values(): + table.add_row((node.hostname, type(node)._identifier, node.operating_state.name)) print(table) if ip_addresses: diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index b2bde574..51e200e7 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1542,6 +1542,9 @@ class Node(SimComponent): _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.""" + def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None: """ Register a node type. @@ -1557,6 +1560,7 @@ class Node(SimComponent): 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 def __init__(self, **kwargs): """ diff --git a/src/primaite/simulator/network/hardware/nodes/host/computer.py b/src/primaite/simulator/network/hardware/nodes/host/computer.py index 68c72554..4253d15c 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/host/computer.py @@ -5,7 +5,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): +class Computer(HostNode, identifier="computer"): """ A basic Computer class. 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 de119562..0c309136 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py +++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py @@ -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): +class HostNode(Node, identifier="HostNode"): """ Represents a host node in the network. diff --git a/src/primaite/simulator/network/hardware/nodes/host/server.py b/src/primaite/simulator/network/hardware/nodes/host/server.py index 379c9927..bf1ef39b 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/server.py +++ b/src/primaite/simulator/network/hardware/nodes/host/server.py @@ -2,7 +2,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode -class Server(HostNode): +class Server(HostNode, identifier="server"): """ A basic Server class. @@ -31,7 +31,7 @@ class Server(HostNode): """ -class Printer(HostNode): +class Printer(HostNode, identifier="printer"): """Printer? I don't even know her!.""" # TODO: Implement printer-specific behaviour diff --git a/src/primaite/simulator/network/hardware/nodes/network/firewall.py b/src/primaite/simulator/network/hardware/nodes/network/firewall.py index 47cfae57..84cf8530 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/firewall.py +++ b/src/primaite/simulator/network/hardware/nodes/network/firewall.py @@ -27,7 +27,7 @@ DMZ_PORT_ID: Final[int] = 3 """The Firewall port ID of the DMZ port.""" -class Firewall(Router): +class Firewall(Router, identifier="firewall"): """ A Firewall class that extends the functionality of a Router. diff --git a/src/primaite/simulator/network/hardware/nodes/network/network_node.py b/src/primaite/simulator/network/hardware/nodes/network/network_node.py index 5ff791cc..a5b8544f 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/network_node.py +++ b/src/primaite/simulator/network/hardware/nodes/network/network_node.py @@ -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): +class NetworkNode(Node, identifier="NetworkNode"): """ Represents an abstract base class for a network node that can receive and process network frames. diff --git a/src/primaite/simulator/network/hardware/nodes/network/router.py b/src/primaite/simulator/network/hardware/nodes/network/router.py index 1080dca8..e921faff 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/router.py @@ -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): +class Router(NetworkNode, identifier="router"): """ Represents a network router, managing routing and forwarding of IP packets across network interfaces. diff --git a/src/primaite/simulator/network/hardware/nodes/network/switch.py b/src/primaite/simulator/network/hardware/nodes/network/switch.py index 4324ac94..d29152a4 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/network/switch.py @@ -87,7 +87,7 @@ class SwitchPort(WiredNetworkInterface): return False -class Switch(NetworkNode): +class Switch(NetworkNode, identifier="switch"): """ A class representing a Layer 2 network switch. 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 00ac9ca4..aed314d2 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py @@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface): ) -class WirelessRouter(Router): +class WirelessRouter(Router, identifier="wireless_router"): """ A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.