From 81fd43035d84273a121b8d9601f6b8ab8b8432ea Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 1 Mar 2024 22:51:01 +0000 Subject: [PATCH] #2358 - the node-specific properties in Network class now simply use node.__class__.__name__ to check their type for filtering by type. Tests updated to use the new property function names --- src/primaite/simulator/network/container.py | 39 +++++++++++-------- tests/conftest.py | 2 +- .../nodes/test_node_config.py | 13 ++++--- ...software_installation_and_configuration.py | 6 +-- .../_simulator/_network/test_container.py | 8 ++-- 5 files changed, 37 insertions(+), 31 deletions(-) diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index b5a16430..6c2f38c5 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -8,10 +8,6 @@ from prettytable import MARKDOWN, PrettyTable from primaite import getLogger from primaite.simulator.core import RequestManager, RequestType, SimComponent from primaite.simulator.network.hardware.base import Link, Node, WiredNetworkInterface -from primaite.simulator.network.hardware.nodes.host.computer import Computer -from primaite.simulator.network.hardware.nodes.host.server import Server -from primaite.simulator.network.hardware.nodes.network.router import Router -from primaite.simulator.network.hardware.nodes.network.switch import Switch from primaite.simulator.system.applications.application import Application from primaite.simulator.system.services.service import Service @@ -85,24 +81,29 @@ class Network(SimComponent): self.links[link_id].apply_timestep(timestep=timestep) @property - def routers(self) -> List[Router]: + def router_nodes(self) -> List[Node]: """The Routers in the Network.""" - return [node for node in self.nodes.values() if isinstance(node, Router)] + return [node for node in self.nodes.values() if node.__class__.__name__ == "Router"] @property - def switches(self) -> List[Switch]: + def switch_nodes(self) -> List[Node]: """The Switches in the Network.""" - return [node for node in self.nodes.values() if isinstance(node, Switch)] + return [node for node in self.nodes.values() if node.__class__.__name__ == "Switch"] @property - def computers(self) -> List[Computer]: + def computer_nodes(self) -> List[Node]: """The Computers in the Network.""" - return [node for node in self.nodes.values() if isinstance(node, Computer) and not isinstance(node, Server)] + return [node for node in self.nodes.values() if node.__class__.__name__ == "Computer"] @property - def servers(self) -> List[Server]: + def server_nodes(self) -> List[Node]: """The Servers in the Network.""" - return [node for node in self.nodes.values() if isinstance(node, Server)] + return [node for node in self.nodes.values() if node.__class__.__name__ == "Server"] + + @property + def firewall_nodes(self) -> List[Node]: + """The Firewalls in the Network.""" + return [node for node in self.nodes.values() if node.__class__.__name__ == "Firewall"] def show(self, nodes: bool = True, ip_addresses: bool = True, links: bool = True, markdown: bool = False): """ @@ -117,10 +118,11 @@ class Network(SimComponent): :param markdown: Use Markdown style in table output. Defaults to False. """ nodes_type_map = { - "Router": self.routers, - "Switch": self.switches, - "Server": self.servers, - "Computer": self.computers, + "Router": self.router_nodes, + "Firewall": self.firewall_nodes, + "Switch": self.switch_nodes, + "Server": self.server_nodes, + "Computer": self.computer_nodes, } if nodes: table = PrettyTable(["Node", "Type", "Operating State"]) @@ -143,7 +145,10 @@ class Network(SimComponent): for node in nodes: for i, port in node.network_interface.items(): if hasattr(port, "ip_address"): - table.add_row([node.hostname, i, port.ip_address, port.subnet_mask, node.default_gateway]) + port_str = port.port_name if port.port_name else port.port_num + table.add_row( + [node.hostname, port_str, port.ip_address, port.subnet_mask, node.default_gateway] + ) print(table) if links: diff --git a/tests/conftest.py b/tests/conftest.py index dbfff2f3..a8c3bdd1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -422,7 +422,7 @@ def install_stuff_to_sim(sim: Simulation): assert len(sim.network.nodes) == 6 assert len(sim.network.links) == 5 # 5.1: Assert the router is correctly configured - r = sim.network.routers[0] + r = sim.network.router_nodes[0] for i, acl_rule in enumerate(r.acl.acl): if i == 1: assert acl_rule.src_port == acl_rule.dst_port == Port.DNS diff --git a/tests/integration_tests/configuration_file_parsing/nodes/test_node_config.py b/tests/integration_tests/configuration_file_parsing/nodes/test_node_config.py index f23e7612..8797bf2e 100644 --- a/tests/integration_tests/configuration_file_parsing/nodes/test_node_config.py +++ b/tests/integration_tests/configuration_file_parsing/nodes/test_node_config.py @@ -11,9 +11,9 @@ def test_example_config(): network: Network = game.simulation.network assert len(network.nodes) == 10 # 10 nodes in example network - assert len(network.routers) == 1 # 1 router in network - assert len(network.switches) == 2 # 2 switches in network - assert len(network.servers) == 5 # 5 servers in network + assert len(network.router_nodes) == 1 # 1 router in network + assert len(network.switch_nodes) == 2 # 2 switches in network + assert len(network.server_nodes) == 5 # 5 servers in network def test_dmz_config(): @@ -23,9 +23,10 @@ def test_dmz_config(): network: Network = game.simulation.network assert len(network.nodes) == 9 # 9 nodes in network - assert len(network.routers) == 2 # 2 routers in network - assert len(network.switches) == 3 # 3 switches in network - assert len(network.servers) == 2 # 2 servers in network + assert len(network.router_nodes) == 1 # 1 router in network + assert len(network.firewall_nodes) == 1 # 1 firewall in network + assert len(network.switch_nodes) == 3 # 3 switches in network + assert len(network.server_nodes) == 2 # 2 servers in network def test_basic_config(): diff --git a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py index f3dc51bd..7da66547 100644 --- a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py +++ b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py @@ -60,9 +60,9 @@ def test_example_config(): network: Network = game.simulation.network assert len(network.nodes) == 10 # 10 nodes in example network - assert len(network.routers) == 1 # 1 router in network - assert len(network.switches) == 2 # 2 switches in network - assert len(network.servers) == 5 # 5 servers in network + assert len(network.router_nodes) == 1 # 1 router in network + assert len(network.switch_nodes) == 2 # 2 switches in network + assert len(network.server_nodes) == 5 # 5 servers in network def test_node_software_install(): diff --git a/tests/unit_tests/_primaite/_simulator/_network/test_container.py b/tests/unit_tests/_primaite/_simulator/_network/test_container.py index 2cfc3f11..f0e386b8 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/test_container.py +++ b/tests/unit_tests/_primaite/_simulator/_network/test_container.py @@ -26,10 +26,10 @@ def filter_keys_nested_item(data, keys): @pytest.fixture(scope="function") def network(example_network) -> Network: - assert len(example_network.routers) is 1 - assert len(example_network.switches) is 2 - assert len(example_network.computers) is 2 - assert len(example_network.servers) is 2 + assert len(example_network.router_nodes) is 1 + assert len(example_network.switch_nodes) is 2 + assert len(example_network.computer_nodes) is 2 + assert len(example_network.server_nodes) is 2 example_network.show()