#1706 - Refactored a bunch of if statements in base.py to improve readability

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

@@ -135,22 +135,26 @@ class NIC(SimComponent):
 
     def enable(self):
         """Attempt to enable the NIC."""
-        if not self.enabled:
+        if self.enabled:
-            if self.connected_node:
+            return
-                if self.connected_node.operating_state == NodeOperatingState.ON:
+        if not self.connected_node:
+            _LOGGER.error(f"NIC {self} cannot be enabled as it is not connected to a Node")
+            return
+        if self.connected_node.operating_state != NodeOperatingState.ON:
+            self.connected_node.sys_log.error(f"NIC {self} cannot be enabled as the endpoint is not turned on")
+            return
+
         self.enabled = True
         self.connected_node.sys_log.info(f"NIC {self} enabled")
         self.pcap = PacketCapture(hostname=self.connected_node.hostname, ip_address=self.ip_address)
         if self.connected_link:
             self.connected_link.endpoint_up()
-                else:
-                    self.connected_node.sys_log.error(f"NIC {self} cannot be enabled as the endpoint is not turned on")
-            else:
-                _LOGGER.error(f"NIC {self} cannot be enabled as it is not connected to a Node")
 
     def disable(self):
         """Disable the NIC."""
-        if self.enabled:
+        if not self.enabled:
+            return
+
         self.enabled = False
         if self.connected_node:
             self.connected_node.sys_log.info(f"NIC {self} disabled")
@@ -166,15 +170,17 @@ class NIC(SimComponent):
         :param link: The link to which the NIC is connected.
         :type link: :class:`~primaite.simulator.network.transmission.physical_layer.Link`
         """
-        if not self.connected_link:
+        if self.connected_link:
-            if self.connected_link != link:
+            _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it already has a connection")
+            return
+
+        if self.connected_link == link:
+            _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it is already connected")
+            return
+
         # TODO: Inform the Node that a link has been connected
         self.connected_link = link
         _LOGGER.info(f"NIC {self} connected to Link {link}")
-            else:
-                _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it is already connected")
-        else:
-            _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it already has a connection")
 
     def disconnect_link(self):
         """Disconnect the NIC from the connected Link."""
@@ -214,7 +220,6 @@ class NIC(SimComponent):
             self.pcap.capture(frame)
             self.connected_link.transmit_frame(sender_nic=self, frame=frame)
             return True
-        else:
         # Cannot send Frame as the NIC is not enabled
         return False
 
@@ -233,7 +238,6 @@ class NIC(SimComponent):
             self.pcap.capture(frame)
             self.connected_node.receive_frame(frame=frame, from_nic=self)
             return True
-        else:
         return False
 
     def describe_state(self) -> Dict:
@@ -290,26 +294,27 @@ class SwitchPort(SimComponent):
 
     def enable(self):
         """Attempt to enable the SwitchPort."""
-        if not self.enabled:
+        if self.enabled:
-            if self.connected_node:
+            return
-                if self.connected_node.operating_state == NodeOperatingState.ON:
+
+        if not self.connected_node:
+            _LOGGER.error(f"SwitchPort {self} cannot be enabled as it is not connected to a Node")
+            return
+
+        if self.connected_node.operating_state != NodeOperatingState.ON:
+            self.connected_node.sys_log.info(f"SwitchPort {self} cannot be enabled as the endpoint is not turned on")
+            return
+
         self.enabled = True
         self.connected_node.sys_log.info(f"SwitchPort {self} enabled")
         self.pcap = PacketCapture(hostname=self.connected_node.hostname)
         if self.connected_link:
             self.connected_link.endpoint_up()
-                else:
-                    self.connected_node.sys_log.info(
-                        f"SwitchPort {self} cannot be enabled as the endpoint is not turned on"
-                    )
-            else:
-                msg = f"SwitchPort {self} cannot be enabled as it is not connected to a Node"
-                _LOGGER.error(msg)
-                raise NetworkError(msg)
 
     def disable(self):
         """Disable the SwitchPort."""
-        if self.enabled:
+        if not self.enabled:
+            return
         self.enabled = False
         if self.connected_node:
             self.connected_node.sys_log.info(f"SwitchPort {self} disabled")
@@ -324,16 +329,18 @@ class SwitchPort(SimComponent):
 
         :param link: The link to which the SwitchPort is connected.
         """
-        if not self.connected_link:
+        if self.connected_link:
-            if self.connected_link != link:
+            _LOGGER.error(f"Cannot connect link to SwitchPort {self.mac_address} as it already has a connection")
+            return
+
+        if self.connected_link == link:
+            _LOGGER.error(f"Cannot connect Link to SwitchPort {self.mac_address} as it is already connected")
+            return
+
         # TODO: Inform the Switch that a link has been connected
         self.connected_link = link
         _LOGGER.info(f"SwitchPort {self} connected to Link {link}")
         self.enable()
-            else:
-                _LOGGER.error(f"Cannot connect Link to SwitchPort {self.mac_address} as it is already connected")
-        else:
-            _LOGGER.error(f"Cannot connect link to SwitchPort {self.mac_address} as it already has a connection")
 
     def disconnect_link(self):
         """Disconnect the SwitchPort from the connected Link."""
@@ -353,7 +360,6 @@ class SwitchPort(SimComponent):
             self.pcap.capture(frame)
             self.connected_link.transmit_frame(sender_nic=self, frame=frame)
             return True
-        else:
         # Cannot send Frame as the SwitchPort is not enabled
         return False
 
@@ -370,7 +376,6 @@ class SwitchPort(SimComponent):
             self.pcap.capture(frame)
             self.connected_node.forward_frame(frame=frame, incoming_port=self)
             return True
-        else:
         return False
 
     def describe_state(self) -> Dict:
@@ -468,28 +473,25 @@ class Link(SimComponent):
         :param frame: The network frame to be sent.
         :return: True if the Frame can be sent, otherwise False.
         """
-        if self._can_transmit(frame):
+        can_transmit = self._can_transmit(frame)
+        if not can_transmit:
+            _LOGGER.info(f"Cannot transmit frame as {self} is at capacity")
+            return False
+
         receiver = self.endpoint_a
         if receiver == sender_nic:
             receiver = self.endpoint_b
         frame_size = frame.size_Mbits
-            sent = receiver.receive_frame(frame)
+
-            if sent:
+        if receiver.receive_frame(frame):
             # Frame transmitted successfully
             # Load the frame size on the link
             self.current_load += frame_size
-                (
             _LOGGER.info(
                 f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits "
                 f"({self.current_load_percent})"
             )
-                )
             return True
-            # Received NIC disabled, reply
-
-            return False
-        else:
-            _LOGGER.info(f"Cannot transmit frame as {self} is at capacity")
         return False
 
     def reset_component_for_episode(self, episode: int):
@@ -624,15 +626,29 @@ class ARPCache:
         :param from_nic: The NIC that received the ARP packet.
         :param arp_packet: The ARP packet to be processed.
         """
-        if arp_packet.request:
+        # ARP Reply
+        if not arp_packet.request:
+            self.sys_log.info(
+                f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}"
+            )
+            self._add_arp_cache_entry(
+                ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic
+            )
+            return
+
+        # ARP Request
         self.sys_log.info(
             f"Received ARP request for {arp_packet.target_ip} from "
             f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} "
         )
-            if arp_packet.target_ip == from_nic.ip_address:
+
-                self._add_arp_cache_entry(
+        # Unmatched ARP Request
-                    ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic
+        if arp_packet.target_ip != from_nic.ip_address:
-                )
+            self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip}")
+            return
+
+        # Matched ARP request
+        self._add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic)
         arp_packet = arp_packet.generate_reply(from_nic.mac_address)
         self.sys_log.info(
             f"Sending ARP reply from {arp_packet.sender_mac_addr}/{arp_packet.sender_ip} "
@@ -652,15 +668,6 @@ class ARPCache:
         )
         frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet)
         from_nic.send_frame(frame)
-            else:
-                self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip}")
-        else:
-            self.sys_log.info(
-                f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}"
-            )
-            self._add_arp_cache_entry(
-                ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic
-            )
 
 
 class ICMP:
@@ -721,7 +728,15 @@ class ICMP:
             was not found in the ARP cache.
         """
         nic = self.arp.get_arp_cache_nic(target_ip_address)
-        if nic:
+        # TODO: Eventually this ARP request needs to be done elsewhere. It's not the resonsibility of the
+        # ping function to handle ARP lookups
+        # No existing ARP entry
+        if not nic:
+            self.sys_log.info(f"No entry in ARP cache for {target_ip_address}")
+            self.arp.send_arp_request(target_ip_address)
+            return 0, None
+
+        # ARP entry exists
         sequence += 1
         target_mac_address = self.arp.get_arp_cache_mac_address(target_ip_address)
         src_nic = self.arp.get_arp_cache_nic(target_ip_address)
@@ -740,10 +755,6 @@ class ICMP:
         self.sys_log.info(f"Sending echo request to {target_ip_address}")
         nic.send_frame(frame)
         return sequence, icmp_packet.identifier
-        else:
-            self.sys_log.info(f"No entry in ARP cache for {target_ip_address}")
-            self.arp.send_arp_request(target_ip_address)
-            return 0, None
 
 
 class NodeOperatingState(Enum):