Give node types their own identifiers and make the network show method use them
This commit is contained in:
@@ -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.game.science import graph_has_cycle, topological_sort
|
||||||
from primaite.simulator import SIM_OUTPUT
|
from primaite.simulator import SIM_OUTPUT
|
||||||
from primaite.simulator.network.creation import NetworkNodeAdder
|
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.computer import Computer
|
||||||
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode, NIC
|
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode, NIC
|
||||||
from primaite.simulator.network.hardware.nodes.host.server import Printer, Server
|
from primaite.simulator.network.hardware.nodes.host.server import Printer, Server
|
||||||
@@ -280,21 +280,7 @@ class PrimaiteGame:
|
|||||||
|
|
||||||
new_node = None
|
new_node = None
|
||||||
# Handle extended nodes
|
# Handle extended nodes
|
||||||
if n_type.lower() in HostNode._registry:
|
if n_type == "computer":
|
||||||
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":
|
|
||||||
new_node = Computer(
|
new_node = Computer(
|
||||||
hostname=node_cfg["hostname"],
|
hostname=node_cfg["hostname"],
|
||||||
ip_address=node_cfg["ip_address"],
|
ip_address=node_cfg["ip_address"],
|
||||||
@@ -339,6 +325,20 @@ class PrimaiteGame:
|
|||||||
if not (p := node_cfg.get("operating_state"))
|
if not (p := node_cfg.get("operating_state"))
|
||||||
else NodeOperatingState[p.upper()],
|
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:
|
else:
|
||||||
msg = f"invalid node type {n_type} in config"
|
msg = f"invalid node type {n_type} in config"
|
||||||
_LOGGER.error(msg)
|
_LOGGER.error(msg)
|
||||||
|
|||||||
@@ -179,9 +179,8 @@ class Network(SimComponent):
|
|||||||
table.set_style(MARKDOWN)
|
table.set_style(MARKDOWN)
|
||||||
table.align = "l"
|
table.align = "l"
|
||||||
table.title = "Nodes"
|
table.title = "Nodes"
|
||||||
for node_type, nodes in nodes_type_map.items():
|
for node in self.nodes.values():
|
||||||
for node in nodes:
|
table.add_row((node.hostname, type(node)._identifier, node.operating_state.name))
|
||||||
table.add_row([node.hostname, node_type, node.operating_state.name])
|
|
||||||
print(table)
|
print(table)
|
||||||
|
|
||||||
if ip_addresses:
|
if ip_addresses:
|
||||||
|
|||||||
@@ -1542,6 +1542,9 @@ class Node(SimComponent):
|
|||||||
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
|
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
|
||||||
"""Registry of application types. Automatically populated when subclasses are defined."""
|
"""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:
|
def __init_subclass__(cls, identifier: str = "default", **kwargs: Any) -> None:
|
||||||
"""
|
"""
|
||||||
Register a node type.
|
Register a node type.
|
||||||
@@ -1557,6 +1560,7 @@ class Node(SimComponent):
|
|||||||
if identifier in cls._registry:
|
if identifier in cls._registry:
|
||||||
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
|
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
|
||||||
cls._registry[identifier] = cls
|
cls._registry[identifier] = cls
|
||||||
|
cls._identifier = identifier
|
||||||
|
|
||||||
def __init__(self, **kwargs):
|
def __init__(self, **kwargs):
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -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
|
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
|
||||||
|
|
||||||
|
|
||||||
class Computer(HostNode):
|
class Computer(HostNode, identifier="computer"):
|
||||||
"""
|
"""
|
||||||
A basic Computer class.
|
A basic Computer class.
|
||||||
|
|
||||||
|
|||||||
@@ -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}"
|
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.
|
Represents a host node in the network.
|
||||||
|
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
|
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
|
||||||
|
|
||||||
|
|
||||||
class Server(HostNode):
|
class Server(HostNode, identifier="server"):
|
||||||
"""
|
"""
|
||||||
A basic Server class.
|
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!."""
|
"""Printer? I don't even know her!."""
|
||||||
|
|
||||||
# TODO: Implement printer-specific behaviour
|
# TODO: Implement printer-specific behaviour
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ DMZ_PORT_ID: Final[int] = 3
|
|||||||
"""The Firewall port ID of the DMZ port."""
|
"""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.
|
A Firewall class that extends the functionality of a Router.
|
||||||
|
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
|
|||||||
from primaite.simulator.system.services.arp.arp import ARP
|
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.
|
Represents an abstract base class for a network node that can receive and process network frames.
|
||||||
|
|
||||||
|
|||||||
@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
|
|||||||
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
|
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.
|
Represents a network router, managing routing and forwarding of IP packets across network interfaces.
|
||||||
|
|
||||||
|
|||||||
@@ -87,7 +87,7 @@ class SwitchPort(WiredNetworkInterface):
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
class Switch(NetworkNode):
|
class Switch(NetworkNode, identifier="switch"):
|
||||||
"""
|
"""
|
||||||
A class representing a Layer 2 network switch.
|
A class representing a Layer 2 network switch.
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user