Merge branch 'dev' into feature/1800_model-a-basic-router-subclass-of-node

src/primaite/simulator/network/container.py

@@ -1,10 +1,13 @@
-from typing import Dict
+from typing import Any, Dict, Union
 
+from primaite import getLogger
 from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent
-from primaite.simulator.network.hardware.base import Link, Node
+from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort
+
+_LOGGER = getLogger(__name__)
 
 
-class NetworkContainer(SimComponent):
+class Network(SimComponent):
     """Top level container object representing the physical network."""
 
     nodes: Dict[str, Node] = {}
@@ -40,3 +43,74 @@ class NetworkContainer(SimComponent):
             }
         )
         return state
+
+    def add_node(self, node: Node) -> None:
+        """
+        Add an existing node to the network.
+
+        :param node: Node instance that the network should keep track of.
+        :type node: Node
+        """
+        if node in self:
+            _LOGGER.warning(f"Can't add node {node.uuid}. It is already in the network.")
+            return
+        self.nodes[node.uuid] = node
+        node.parent = self
+        _LOGGER.info(f"Added node {node.uuid} to Network {self.uuid}")
+
+    def remove_node(self, node: Node) -> None:
+        """
+        Remove a node from the network.
+
+        :param node: Node instance that is currently part of the network that should be removed.
+        :type node: Node
+        """
+        if node not in self:
+            _LOGGER.warning(f"Can't remove node {node.uuid}. It's not in the network.")
+            return
+        self.nodes.pop(node.uuid)
+        node.parent = None
+        _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}")
+
+    def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None:
+        """Connect two nodes on the network by creating a link between an NIC/SwitchPort of each one.
+
+        :param endpoint_a: The endpoint to which to connect the link on the first node
+        :type endpoint_a: Union[NIC, SwitchPort]
+        :param endpoint_b: The endpoint to which to connct the link on the second node
+        :type endpoint_b: Union[NIC, SwitchPort]
+        :raises RuntimeError: _description_
+        """
+        node_a = endpoint_a.parent
+        node_b = endpoint_b.parent
+        if node_a not in self:
+            self.add_node(node_a)
+        if node_b not in self:
+            self.add_node(node_b)
+        if node_a is node_b:
+            _LOGGER.warn(f"Cannot link endpoint {endpoint_a} to {endpoint_b} because they belong to the same node.")
+            return
+
+        link = Link(endpoint_a=endpoint_a, endpoint_b=endpoint_b, **kwargs)
+        self.links[link.uuid] = link
+        link.parent = self
+        _LOGGER.info(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}")
+
+    def remove_link(self, link: Link) -> None:
+        """Disconnect a link from the network.
+
+        :param link: The link to be removed
+        :type link: Link
+        """
+        link.endpoint_a.disconnect_link()
+        link.endpoint_b.disconnect_link()
+        self.links.pop(link.uuid)
+        link.parent = None
+        _LOGGER.info(f"Removed link {link.uuid} from network {self.uuid}.")
+
+    def __contains__(self, item: Any) -> bool:
+        if isinstance(item, Node):
+            return item.uuid in self.nodes
+        elif isinstance(item, Link):
+            return item.uuid in self.links
+        return False

src/primaite/simulator/network/hardware/base.py

@@ -970,6 +970,7 @@ class Node(SimComponent):
         if nic.uuid not in self.nics:
             self.nics[nic.uuid] = nic
             nic.connected_node = self
+            nic.parent = self
             self.sys_log.info(f"Connected NIC {nic}")
             if self.operating_state == NodeOperatingState.ON:
                 nic.enable()
@@ -990,6 +991,7 @@ class Node(SimComponent):
             nic = self.nics.get(nic)
         if nic or nic.uuid in self.nics:
             self.nics.pop(nic.uuid)
+            nic.parent = None
             nic.disable()
             self.sys_log.info(f"Disconnected NIC {nic}")
         else:
@@ -1073,6 +1075,7 @@ class Switch(Node):
             self.switch_ports = {i: SwitchPort() for i in range(1, self.num_ports + 1)}
         for port_num, port in self.switch_ports.items():
             port.connected_node = self
+            port.parent = self
             port.port_num = port_num
 
     def show(self):