primitive 0.2.66__py3-none-any.whl → 0.2.70__py3-none-any.whl

primitive/network/actions.py

@@ -1,10 +1,12 @@
 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.messaging.provider import MESSAGE_TYPES
 from primitive.utils.actions import BaseAction
@@ -39,6 +41,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 +53,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 +150,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 +160,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 +280,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 +334,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 +371,13 @@ class Network(BaseAction):
                     )
         return ip_to_mac_address_info
 
+    def get_ip_arp_table_via_ip_command(self):
+        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 +457,10 @@ class Network(BaseAction):
 
         return False
 
-    def push_switch_and_interfaces_info(self):
+    def push_switch_and_interfaces_info(self, interfaces_info: dict | None = None):
         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:

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

@@ -0,0 +1,126 @@
+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):
+    """
+    Tests an SSH connection to a remote host.
+
+    Args:
+        hostname (str): The hostname or IP address of the remote SSH server.
+        username (str): The username for authentication.
+        password (str, optional): The password for authentication. Defaults to None.
+        key_filename (str, optional): Path to the private key file for authentication. Defaults to None.
+        port (int, optional): The SSH port. Defaults to 22.
+
+    Returns:
+        bool: True if the connection is successful, False otherwise.
+    """
+    ssh_client = paramiko.SSHClient()
+    ssh_client.set_missing_host_key_policy(
+        paramiko.AutoAddPolicy()
+    )  # Auto-add new host keys
+
+    try:
+        if password:
+            ssh_client.connect(
+                hostname=hostname, port=port, username=username, password=password
+            )
+        elif key_filename:
+            ssh_client.connect(
+                hostname=hostname,
+                port=port,
+                username=username,
+                key_filename=key_filename,
+            )
+        else:
+            print(
+                "Error: Either password or key_filename must be provided for authentication."
+            )
+            return False
+
+        print(f"Successfully connected to {hostname} as {username}")
+        return True
+    except paramiko.AuthenticationException:
+        print(f"Authentication failed for {username} on {hostname}")
+        return False
+    except paramiko.SSHException as exception:
+        print(f"SSH error connecting to {hostname}: {exception}")
+        return False
+    except socket.error as exception:
+        print(f"Socket error connecting to {hostname}: {exception}")
+        return False
+    except Exception as exception:
+        print(f"An unexpected error occurred: {exception}")
+        return False
+    finally:
+        ssh_client.close()
+
+
+TEN_MINUTES = 60 * 10
+
+
+def wait_for_ssh(
+    hostname, username, password=None, key_filename=None, port=22, timeout=TEN_MINUTES
+):
+    """
+    Waits until an SSH connection to a remote host can be established.
+
+    Args:
+        hostname (str): The hostname or IP address of the remote SSH server.
+        username (str): The username for authentication.
+        password (str, optional): The password for authentication. Defaults to None.
+        key_filename (str, optional): Path to the private key file for authentication. Defaults to None.
+        port (int, optional): The SSH port. Defaults to 22.
+        timeout (int, optional): Maximum time to wait in seconds. Defaults to 300.
+
+    Returns:
+        bool: True if the connection is successful within the timeout, False otherwise.
+    """
+
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if test_ssh_connection(
+            hostname, username, password=password, key_filename=key_filename, port=port
+        ):
+            return True
+        logger.debug(f"Waiting for SSH to become available on {hostname}...")
+        time.sleep(10)
+
+    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)

primitive/operating_systems/actions.py

@@ -0,0 +1,260 @@
+from enum import Enum
+
+from gql import gql
+import requests
+
+from primitive.operating_systems.graphql.mutations import (
+    operating_system_create_mutation,
+)
+from primitive.operating_systems.graphql.queries import operating_system_list_query
+from primitive.utils.actions import BaseAction
+from primitive.utils.auth import guard
+
+from primitive.utils.cache import get_operating_systems_cache
+from pathlib import Path
+from urllib.request import urlopen
+import os
+from loguru import logger
+
+from primitive.utils.checksums import get_checksum_from_file, calculate_sha256
+
+
+class OperatingSystems(BaseAction):
+    def __init__(self, primitive):
+        super().__init__(primitive)
+        self.remote_operating_systems = {
+            "ubuntu-24-04-3": {
+                "iso": "https://releases.ubuntu.com/24.04.3/ubuntu-24.04.3-desktop-amd64.iso",
+                "checksum": "https://releases.ubuntu.com/24.04.3/SHA256SUMS",
+                "checksum_file_type": self.OperatingSystemChecksumFileType.SHA256SUMS,
+            },
+        }
+
+    class OperatingSystemChecksumFileType(Enum):
+        SHA256SUMS = "SHA256SUMS"
+
+    def get_remote_operating_system_names(self):
+        return list(self.remote_operating_systems.keys())
+
+    def _download_remote_operating_system_iso(self, remote_operating_system_name):
+        operating_system_dir = Path(
+            get_operating_systems_cache() / remote_operating_system_name
+        )
+        iso_dir = Path(operating_system_dir / "iso")
+        os.makedirs(iso_dir, exist_ok=True)
+
+        operating_system_info = self.remote_operating_systems[
+            remote_operating_system_name
+        ]
+        iso_remote_url = operating_system_info["iso"]
+        iso_filename = iso_remote_url.split("/")[-1]
+        iso_file_path = Path(iso_dir / iso_filename)
+
+        if iso_file_path.exists() and iso_file_path.is_file():
+            logger.info("Operating system iso already downloaded.")
+            return iso_file_path
+
+        logger.info(
+            f"Downloading operating system '{remote_operating_system_name}' iso. This may take a few minutes..."
+        )
+
+        session = requests.Session()
+        with session.get(iso_remote_url, stream=True) as response:
+            response.raise_for_status()
+            with open(iso_file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+                        f.flush()
+
+        logger.info(
+            f"Successfully downloaded operating system iso to '{iso_file_path}'."
+        )
+
+        return iso_file_path
+
+    def _download_remote_operating_system_checksum(self, remote_operating_system_name):
+        operating_system_dir = Path(
+            get_operating_systems_cache() / remote_operating_system_name
+        )
+        checksum_dir = Path(operating_system_dir / "checksum")
+        os.makedirs(checksum_dir, exist_ok=True)
+
+        operating_system_info = self.remote_operating_systems[
+            remote_operating_system_name
+        ]
+        checksum_filename = operating_system_info["checksum"].split("/")[-1]
+
+        checksum_file_path = Path(checksum_dir / checksum_filename)
+        if checksum_file_path.exists() and checksum_file_path.is_file():
+            logger.info("Operating system checksum already downloaded.")
+            return checksum_file_path
+
+        logger.info(
+            f"Downloading operating system '{remote_operating_system_name}' checksum."
+        )
+
+        checksum_response = urlopen(operating_system_info["checksum"])
+        checksum_file_content = checksum_response.read()
+        with open(checksum_file_path, "wb") as f:
+            f.write(checksum_file_content)
+
+        logger.info(f"Successfully downloaded checksum to '{checksum_file_path}'.")
+
+        return checksum_file_path
+
+    def download_from_remote(self, remote_operating_system_name: str):
+        remote_operating_system_names = self.get_remote_operating_system_names()
+
+        if remote_operating_system_name not in remote_operating_system_names:
+            logger.error(
+                f"No such operating system '{remote_operating_system_name}'. Run 'primitive operating-systems list' for available operating systems."
+            )
+            raise ValueError(
+                f"No such operating system '{remote_operating_system_name}'."
+            )
+
+        iso_file_path = self._download_remote_operating_system_iso(
+            remote_operating_system_name
+        )
+        checksum_file_path = self._download_remote_operating_system_checksum(
+            remote_operating_system_name
+        )
+
+        return iso_file_path, checksum_file_path
+
+    def get_checksum_file_type(self, operating_system_name: str):
+        return self.remote_operating_systems[operating_system_name][
+            "checksum_file_type"
+        ]
+
+    def validate_checksum(
+        self,
+        operating_system_name: str,
+        iso_file_path: str,
+        checksum_file_path: str,
+        checksum_file_type: OperatingSystemChecksumFileType | None = None,
+    ):
+        checksum_file_type = (
+            checksum_file_type
+            if checksum_file_type
+            else self.get_checksum_file_type(operating_system_name)
+        )
+
+        match checksum_file_type:
+            case self.OperatingSystemChecksumFileType.SHA256SUMS:
+                return self._validate_sha256_sums_checksum(
+                    iso_file_path, checksum_file_path
+                )
+            case _:
+                logger.error(f"Invalid checksum file type: {checksum_file_type}")
+                raise ValueError(f"Invalid checksum file type: {checksum_file_type}")
+
+    def _validate_sha256_sums_checksum(self, iso_file_path, checksum_file_path):
+        iso_file_name = Path(iso_file_path).name
+
+        remote_checksum = get_checksum_from_file(checksum_file_path, iso_file_name)
+        local_checksum = calculate_sha256(iso_file_path)
+        return remote_checksum == local_checksum
+
+    @guard
+    def create_operating_system(
+        self,
+        slug: str,
+        organization_id: str,
+        checksum_file_id: str,
+        checksum_file_type: str,
+        iso_file_id: str,
+    ):
+        mutation = gql(operating_system_create_mutation)
+        input = {
+            "slug": slug,
+            "organization": organization_id,
+            "checksumFile": checksum_file_id,
+            "checksumFileType": checksum_file_type,
+            "isoFile": iso_file_id,
+        }
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result.data.get("operatingSystemCreate")
+
+    @guard
+    def get_operating_system_list(
+        self,
+        organization_id: str,
+        slug: str | None = None,
+        id: str | None = None,
+    ):
+        query = gql(operating_system_list_query)
+
+        variables = {
+            "filters": {
+                "organization": {"id": organization_id},
+            }
+        }
+
+        if slug:
+            variables["filters"]["slug"] = {"exact": slug}
+
+        if id:
+            variables["filters"]["id"] = id
+
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+
+        return result
+
+    @guard
+    def get_operating_system(
+        self, organization_id: str, slug: str | None = None, id: str | None = None
+    ):
+        if not (slug or id):
+            raise Exception("Slug or id must be provided.")
+        if slug and id:
+            raise Exception("Only one of slug or id must be provided.")
+
+        operating_system_list_result = self.get_operating_system_list(
+            organization_id=organization_id, slug=slug, id=id
+        )
+
+        edges = operating_system_list_result.data.get("operatingSystemList").get(
+            "edges", []
+        )
+
+        if len(edges) == 0:
+            if slug:
+                logger.error(f"No operating system found for slug '{slug}'.")
+                raise Exception(f"No operating system  found for slug {slug}.")
+            else:
+                logger.error(f"No operating system found for ID {id}.")
+                raise Exception(f"No operating system  found for ID {id}.")
+
+        return edges[0].get("node")
+
+    @guard
+    def is_slug_available(self, slug: str, organization_id: str):
+        query = gql(operating_system_list_query)
+
+        variables = {
+            "filters": {
+                "slug": {"exact": slug},
+                "organization": {"id": organization_id},
+            }
+        }
+
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+
+        count = result.data.get("operatingSystemList").get("totalCount")
+
+        return count == 0
+
+    def is_operating_system_cached(self, slug: str, directory: str | None = None):
+        cache_dir = Path(directory) if directory else get_operating_systems_cache()
+        cache_path = cache_dir / slug
+
+        return cache_path.exists()