primitive 0.2.68__py3-none-any.whl → 0.2.72__py3-none-any.whl

primitive/monitor/actions.py

@@ -94,7 +94,7 @@ class Monitor(BaseAction):
                         if hardware.get("isController", False):
                             self.primitive.hardware.get_and_set_switch_info()
                             self.primitive.network.push_switch_and_interfaces_info()
-                        self.primitive.network.push_own_network_interfaces()
+                        self.primitive.hardware.push_own_system_info()
                         self.primitive.hardware._sync_children(hardware=hardware)
 
                         sleep_amount = 5
@@ -191,11 +191,16 @@ class Monitor(BaseAction):
                         job_run_id = None
                     else:
                         hardware_id = hardware.get("id", None) if hardware else None
-                        execution_hardware_id = (
-                            job_run_data.get("executionHardware", {}).get("id", None)
-                            if job_run_data
-                            else None
-                        )
+                        execution_hardware_id = None
+                        if job_run_data:
+                            execution_hardware = job_run_data.get(
+                                "executionHardware", None
+                            )
+                            execution_hardware_id = (
+                                execution_hardware.get("id", None)
+                                if execution_hardware
+                                else None
+                            )
 
                         if (
                             hardware_id is not None

primitive/network/actions.py

@@ -1,11 +1,14 @@
 from pathlib import Path
-
+from subprocess import PIPE, Popen
+import json
 import requests
 import serial
 from loguru import logger
 from paramiko import SSHClient
 from typing import TypedDict
+import re
 
+from primitive.hardware.actions import does_executable_exist
 from primitive.messaging.provider import MESSAGE_TYPES
 from primitive.utils.actions import BaseAction
 
@@ -39,6 +42,11 @@ def mac_address_manufacturer_style_to_ieee(mac: str) -> str:
     return ":".join(mac[i : i + 2] for i in range(0, 12, 2))
 
 
+def natural_interface_key(s):
+    # extract numbers after "Ethernet" or subports
+    return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", s)]
+
+
 class SwitchConnectionInfo(TypedDict):
     vendor: str
     hostname: str
@@ -46,6 +54,12 @@ class SwitchConnectionInfo(TypedDict):
     password: str
 
 
+class MacAddressEntry(TypedDict):
+    ip_address: str | None
+    mac_address: str
+    vlan: str
+
+
 class Network(BaseAction):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
@@ -137,6 +151,8 @@ class Network(BaseAction):
                 return None
 
     def get_switch_info(self):
+        if self.switch_connection_info is None:
+            self.primitive.hardware.get_and_set_switch_info()
         if self.is_switch_api_enabled():
             switch_info = self.get_switch_info_via_api()
             if switch_info:
@@ -145,26 +161,46 @@ class Network(BaseAction):
         return None
 
     def get_interfaces_info(self):
+        if self.switch_connection_info is None:
+            self.primitive.hardware.get_and_set_switch_info()
         if self.is_switch_api_enabled():
             interfaces_info = self.get_interfaces_via_api()
             mac_address_info = self.get_mac_address_info_via_api()
+
             ip_arp_table_info = self.get_ip_arp_table_via_api()
+            controllers_neighbors = self.get_ip_arp_table_via_ip_command()
 
             if interfaces_info and mac_address_info and ip_arp_table_info:
                 for interface, mac_info in mac_address_info.items():
-                    if interface in interfaces_info.get("interfaces", {}):
-                        interfaces_info["interfaces"][interface]["mac_address"] = (
-                            mac_info.get("macAddress", "")
-                        )
-                    if interface in ip_arp_table_info:
-                        for ip_arp in ip_arp_table_info[interface]:
-                            if (
-                                interfaces_info["interfaces"][interface]["ip_address"]
-                                == ""
-                            ):
-                                interfaces_info["interfaces"][interface][
+                    if interface in interfaces_info:
+                        mac_addresses: dict[str, MacAddressEntry] = {}
+                        for entry in mac_info:
+                            mac_addresses[entry.get("macAddress", "")] = {
+                                "mac_address": entry.get("macAddress", ""),
+                                "ip_address": None,
+                                "vlan": entry.get("vlanId", ""),
+                            }
+
+                        for neighbor in controllers_neighbors:
+                            if neighbor.get("lladdr", "") in mac_addresses:
+                                mac_addresses[neighbor.get("lladdr", "")][
                                     "ip_address"
-                                ] = ip_arp.get("ip_address", "")
+                                ] = neighbor.get("dst", None)
+
+                        interfaces_info[interface]["mac_addresses"] = mac_addresses
+
+                        if interface in ip_arp_table_info:
+                            for ip_arp in ip_arp_table_info[interface]:
+                                for mac_address_entry in interfaces_info[interface][
+                                    "mac_addresses"
+                                ].values():
+                                    if (
+                                        ip_arp.get("mac_address", "")
+                                        in mac_address_entry["mac_address"]
+                                    ):
+                                        mac_address_entry["ip_address"] = ip_arp.get(
+                                            "ip_address", None
+                                        )
 
                 return interfaces_info
 
@@ -245,25 +281,27 @@ class Network(BaseAction):
                 # }
                 arista_interfaces_info = response.get("result", [])[0]
                 formatted_interfaces_info = {
-                    "interfaces": {
-                        k: {
-                            "interface_name": k,
-                            "interface_type": v.get("interfaceType", ""),
-                            "link_status": v.get("linkStatus", ""),
-                            "line_protocol_status": v.get("lineProtocolStatus", ""),
-                            "mac_address": "",
-                            "ip_address": "",
-                        }
-                        for k, v in dict(
-                            sorted(
-                                arista_interfaces_info.get(
-                                    "interfaceStatuses", {}
-                                ).items()
-                            )
-                        ).items()
-                    },
-                    # "raw_output": arista_interfaces_info,
+                    k: {
+                        "interface_name": k,
+                        "interface_type": v.get("interfaceType", ""),
+                        "link_status": v.get("linkStatus", ""),
+                        "line_protocol_status": v.get("lineProtocolStatus", ""),
+                        "mac_addresses": {},
+                    }
+                    for k, v in dict(
+                        sorted(
+                            arista_interfaces_info.get("interfaceStatuses", {}).items()
+                        )
+                    ).items()
                 }
+        if formatted_interfaces_info:
+            formatted_interfaces_info = {
+                k: formatted_interfaces_info[k]
+                for k in sorted(
+                    formatted_interfaces_info.keys(), key=natural_interface_key
+                )
+            }
+
         return formatted_interfaces_info
 
     def get_mac_address_info_via_api(self):
@@ -297,9 +335,14 @@ class Network(BaseAction):
                     .get("tableEntries", [])
                 )
                 table_entries.sort(key=lambda x: x["lastMove"])
+
                 for entry in table_entries:
                     if entry.get("interface") not in interface_to_mac_address_info:
-                        interface_to_mac_address_info[entry.get("interface")] = entry
+                        interface_to_mac_address_info[entry.get("interface")] = [entry]
+                    else:
+                        interface_to_mac_address_info[entry.get("interface")].append(
+                            entry
+                        )
 
         return interface_to_mac_address_info
 
@@ -329,6 +372,16 @@ class Network(BaseAction):
                     )
         return ip_to_mac_address_info
 
+    def get_ip_arp_table_via_ip_command(self):
+        if does_executable_exist("ip") is False:
+            return []
+
+        command = "ip --json neigh show"
+        ip_result = None
+        with Popen(command.split(" "), stdout=PIPE) as process:
+            ip_result = json.loads(process.stdout.read().decode("utf-8"))
+        return ip_result
+
     def serial_connect(self):
         self.ser = serial.Serial()
         self.ser.port = self.switch_tty_name
@@ -408,10 +461,11 @@ class Network(BaseAction):
 
         return False
 
-    def push_switch_and_interfaces_info(self):
+    def push_switch_and_interfaces_info(self, interfaces_info: dict | None = None):
+        logger.debug("Pushing switch and interfaces info")
         if self.primitive.messaging.ready and self.switch_connection_info is not None:
             switch_info = self.get_switch_info()
-            interfaces_info = self.get_interfaces_info()
+            interfaces_info = interfaces_info or self.get_interfaces_info()
 
             message = {"switch_info": {}, "interfaces_info": {}}
             if switch_info:
@@ -424,13 +478,4 @@ class Network(BaseAction):
                     message_type=MESSAGE_TYPES.SWITCH_AND_INTERFACES_INFO,
                     message=message,
                 )
-
-    def push_own_network_interfaces(self):
-        if self.primitive.messaging.ready:
-            message = {
-                "network_interfaces": self.primitive.hardware._get_network_interfaces()
-            }
-            self.primitive.messaging.send_message(
-                message_type=MESSAGE_TYPES.OWN_NETWORK_INTERFACES,
-                message=message,
-            )
+                logger.debug("Switch and interfaces info pushed")

primitive/network/commands.py

@@ -28,13 +28,24 @@ def switch(context):
     print_result(message=message, context=context)
 
 
-@cli.command("ports")
+@cli.command("interfaces")
 @click.pass_context
-def ports(context):
-    """Ports"""
+@click.option(
+    "--push",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Push current interface info.",
+)
+def interfaces(context, push: bool = False):
+    """Interfaces"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    ports_info = primitive.network.get_interfaces_info()
+    interfaces_info = primitive.network.get_interfaces_info()
+    if push:
+        primitive.network.push_switch_and_interfaces_info(
+            interfaces_info=interfaces_info
+        )
     if context.obj["JSON"]:
-        print_result(message=ports_info, context=context)
+        print_result(message=interfaces_info, context=context)
     else:
-        render_ports_table(ports_info.get("interfaces"))
+        render_ports_table(interfaces_info)

primitive/network/ssh.py

@@ -1,6 +1,8 @@
+from loguru import logger
 import paramiko
 import socket
 import time
+from paramiko import SSHClient
 
 
 def test_ssh_connection(hostname, username, password=None, key_filename=None, port=22):
@@ -25,7 +27,13 @@ def test_ssh_connection(hostname, username, password=None, key_filename=None, po
     try:
         if password:
             ssh_client.connect(
-                hostname=hostname, port=port, username=username, password=password
+                hostname=hostname,
+                port=port,
+                username=username,
+                password=password,
+                banner_timeout=60,
+                timeout=60,
+                auth_timeout=60,
             )
         elif key_filename:
             ssh_client.connect(
@@ -33,26 +41,29 @@ def test_ssh_connection(hostname, username, password=None, key_filename=None, po
                 port=port,
                 username=username,
                 key_filename=key_filename,
+                banner_timeout=60,
+                timeout=60,
+                auth_timeout=60,
             )
         else:
-            print(
+            logger.error(
                 "Error: Either password or key_filename must be provided for authentication."
             )
             return False
 
-        print(f"Successfully connected to {hostname} as {username}")
+        logger.info(f"Successfully connected to {hostname} as {username}")
         return True
     except paramiko.AuthenticationException:
-        print(f"Authentication failed for {username} on {hostname}")
+        logger.debug(f"Authentication failed for {username} on {hostname}")
         return False
-    except paramiko.SSHException as e:
-        print(f"SSH error connecting to {hostname}: {e}")
+    except paramiko.SSHException as exception:
+        logger.debug(f"SSH error connecting to {hostname}: {exception}")
         return False
-    except socket.error as e:
-        print(f"Socket error connecting to {hostname}: {e}")
+    except socket.error as exception:
+        logger.debug(f"Socket error connecting to {hostname}: {exception}")
         return False
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
+    except Exception as exception:
+        logger.debug(f"An unexpected error occurred: {exception}")
         return False
     finally:
         ssh_client.close()
@@ -85,10 +96,40 @@ def wait_for_ssh(
             hostname, username, password=password, key_filename=key_filename, port=port
         ):
             return True
-        print(f"Waiting for SSH to become available on {hostname}...")
+        logger.debug(f"Waiting for SSH to become available on {hostname}...")
         time.sleep(10)
 
-    print(
+    logger.warning(
         f"Timeout reached: Unable to connect to {hostname} via SSH within {timeout} seconds."
     )
     return False
+
+
+def run_command(
+    hostname,
+    username,
+    command: str,
+    password=None,
+    key_filename=None,
+    port=22,
+):
+    ssh_client = SSHClient()
+    ssh_client.load_system_host_keys()
+    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh_client.connect(
+        hostname=hostname,
+        port=port,
+        username=username,
+        password=password,
+        key_filename=key_filename,
+    )
+    stdin, stdout, stderr = ssh_client.exec_command(command)
+
+    stdout_string = stdout.read().decode("utf-8").rstrip("\n")
+    stderr_string = stderr.read().decode("utf-8").rstrip("\n")
+    if stdout_string != b"":
+        logger.info(stdout_string)
+    if stderr_string != b"":
+        logger.error(stderr_string)
+
+    ssh_client.close()

primitive/network/ui.py

@@ -8,12 +8,18 @@ def render_ports_table(ports_dict) -> None:
     table = Table(show_header=True, header_style="bold #FFA800")
     table.add_column("Port")
     table.add_column("Status")
-    table.add_column("MAC Address")
-    table.add_column("IP Address")
+    table.add_column("MAC Address | IP | VLAN")
 
     for k, v in ports_dict.items():
         table.add_row(
-            k, v.get("link_status"), v.get("mac_address"), v.get("ip_address")
+            k,
+            v.get("link_status"),
+            "\n".join(
+                [
+                    f"{key} | {values.get('ip_address')} | VLAN {values.get('vlan')}"
+                    for key, values in v.get("mac_addresses", {}).items()
+                ]
+            ),
         )
 
     console.print(table)