primitive 0.2.61__py3-none-any.whl → 0.2.64__py3-none-any.whl

primitive/hardware/graphql/queries.py

@@ -30,6 +30,47 @@ query hardwareList(
 """
 )
 
+hardware_with_parent_list = (
+    hardware_fragment
+    + """
+query hardwareWithParentList(
+  $before: String
+  $after: String
+  $first: Int
+  $last: Int
+  $filters: HardwareFilters
+) {
+  hardwareList(
+    before: $before
+    after: $after
+    first: $first
+    last: $last
+    filters: $filters
+  ) {
+    totalCount
+    edges {
+      cursor
+      node {
+        ...HardwareFragment
+        parent {
+          id
+          pk
+          name
+          slug
+          manufacturer {
+            id
+            pk
+            name
+            slug
+          }
+        }
+      }
+    }
+  }
+}
+"""
+)
+
 nested_children_hardware_list = (
     hardware_fragment
     + """
@@ -97,3 +138,10 @@ query hardwareDetails(
 }
 """
 )
+
+
+hardware_secret = """
+query hardwareSecret($hardwareId: ID!) {
+  hardwareSecret(hardwareId: $hardwareId)
+}
+"""

primitive/jobs/graphql/fragments.py

@@ -38,6 +38,10 @@ fragment JobRunFragment on JobRun {
     branch
     repoFullName
   }
+  executionHardware {
+    id
+    pk
+  }
 }
 """
 

primitive/jobs/graphql/queries.py

@@ -74,7 +74,7 @@ query jobRuns(
 job_run_query = (
     job_run_fragment
     + """
-query jobRun($id: GlobalID!) {
+query jobRun($id: ID!) {
     jobRun(id: $id) {
         ...JobRunFragment
     }
@@ -83,7 +83,7 @@ query jobRun($id: GlobalID!) {
 )
 
 github_app_token_for_job_run_query = """
-query ghAppTokenForJobRun($jobRunId: GlobalID!) {
+query ghAppTokenForJobRun($jobRunId: ID!) {
     ghAppTokenForJobRun(jobRunId: $jobRunId)
 }
 """
@@ -91,7 +91,7 @@ query ghAppTokenForJobRun($jobRunId: GlobalID!) {
 job_run_status_query = (
     job_run_status_fragment
     + """
-query jobRun($id: GlobalID!) {
+query jobRun($id: ID!) {
     jobRun(id: $id) {
         ...JobRunStatusFragment
     }
@@ -101,7 +101,7 @@ query jobRun($id: GlobalID!) {
 
 
 job_secrets_for_job_run_query = """
-query jobSecretsForJobRun($jobRunId: GlobalID!) {
+query jobSecretsForJobRun($jobRunId: ID!) {
     jobSecretsForJobRun(jobRunId: $jobRunId)
 }
 """

primitive/messaging/provider.py

@@ -3,16 +3,18 @@ import json
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING
 from uuid import uuid4
+
 import pika
+from loguru import logger
 from pika import credentials
+
+from primitive.utils.actions import BaseAction
+from primitive.__about__ import __version__
 from ..utils.x509 import (
     are_certificate_files_present,
     create_ssl_context,
     read_certificate_common_name,
 )
-from loguru import logger
-
-from primitive.utils.actions import BaseAction
 
 if TYPE_CHECKING:
     import primitive.client
@@ -26,6 +28,8 @@ CELERY_TASK_NAME = "hardware.tasks.task_receive_hardware_message"
 class MESSAGE_TYPES(enum.Enum):
     CHECK_IN = "CHECK_IN"
     METRICS = "METRICS"
+    SWITCH_AND_INTERFACES_INFO = "SWITCH_AND_INTERFACES_INFO"
+    OWN_NETWORK_INTERFACES = "OWN_NETWORK_INTERFACES"
 
 
 class MessagingProvider(BaseAction):
@@ -52,7 +56,7 @@ class MessagingProvider(BaseAction):
         elif primitive.host == "api.test.primitive.tech":
             rabbitmq_host = "rabbitmq-cluster.test.primitive.tech"
         elif primitive.host == "localhost:8000":
-            rabbitmq_host = "localhost"
+            rabbitmq_host = primitive.host.split(":")[0]
             RABBITMQ_PORT = 5671
 
         if not are_certificate_files_present():
@@ -99,6 +103,7 @@ class MessagingProvider(BaseAction):
 
             headers = {
                 "fingerprint": self.fingerprint,
+                "version": __version__,
                 "token": self.token,
                 "argsrepr": "()",
                 "id": message_uuid,

primitive/monitor/actions.py

@@ -39,7 +39,7 @@ class Monitor(BaseAction):
                 # only set is_available after we've checked that no active reservation is present
                 # setting is_available of the parent also effects the children,
                 # which may have active reservations as well
-                self.primitive.hardware.check_in_http(is_online=True)
+                self.primitive.hardware.check_in(is_online=True)
             except Exception as exception:
                 logger.exception(f"Error checking in hardware: {exception}")
                 sys.exit(1)
@@ -59,6 +59,8 @@ class Monitor(BaseAction):
             active_reservation_data = None
             previous_reservation_id = None
             active_reservation_id = None
+            hardware = None
+            job_run_data = None
 
             last_provisioned_reservation_id = None
 
@@ -68,8 +70,8 @@ class Monitor(BaseAction):
                 # obtains an active JobRun's ID
                 if not RUNNING_IN_CONTAINER:
                     # self.primitive.hardware.push_metrics()
-
                     hardware = self.primitive.hardware.get_own_hardware_details()
+
                     # fetch the latest hardware and activeReservation details
                     if active_reservation_data := hardware["activeReservation"]:
                         active_reservation_id = active_reservation_data.get("id", None)
@@ -85,9 +87,14 @@ class Monitor(BaseAction):
                         and active_reservation_id is None
                         and previous_reservation_id is None
                     ):
-                        self.primitive.hardware.check_in_http(
+                        self.primitive.hardware.check_in(
                             is_available=True, is_online=True
                         )
+                        # if the hardware is a control node, get the latest switch info
+                        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._sync_children(hardware=hardware)
 
                         sleep_amount = 5
@@ -138,7 +145,7 @@ class Monitor(BaseAction):
 
                     # Golden state for normal reservation
                     if not job_run_id and active_reservation_id:
-                        self.primitive.hardware.check_in_http(
+                        self.primitive.hardware.check_in(
                             is_available=False, is_online=True
                         )
                         sleep_amount = 5
@@ -148,7 +155,7 @@ class Monitor(BaseAction):
                         sleep(sleep_amount)
                         continue
 
-                # job_run_data can come from 3 places:
+                # job_run_id can come from 3 places:
                 # 1. an explicitly passed job_run_id
                 # 2. the previous reservation has an job_run_id (kill old PIDs)
                 # 3. the active reservation has an job_run_id (check status)
@@ -183,6 +190,21 @@ class Monitor(BaseAction):
                         status_value = None
                         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
+                        )
+
+                        if (
+                            hardware_id is not None
+                            and execution_hardware_id is not None
+                            and (hardware_id != execution_hardware_id)
+                        ):
+                            logger.info(
+                                f"Job Run {job_run_id} is being executed by the controller. Monitoring may stop."
+                            )
                         logger.info(
                             f"Job Run {job_run_id} with Status {status_value} with PID {parent_pid}. [sleeping {sleep_amount} seconds]"
                         )
@@ -193,7 +215,7 @@ class Monitor(BaseAction):
             logger.info("Stopping primitive monitor...")
             try:
                 if not RUNNING_IN_CONTAINER:
-                    self.primitive.hardware.check_in_http(
+                    self.primitive.hardware.check_in(
                         is_available=False, is_online=False, stopping_agent=True
                     )
 

primitive/network/actions.py

@@ -0,0 +1,436 @@
+from pathlib import Path
+
+import requests
+import serial
+from loguru import logger
+from paramiko import SSHClient
+from typing import TypedDict
+
+from primitive.messaging.provider import MESSAGE_TYPES
+from primitive.utils.actions import BaseAction
+
+# 0. see if this controller node hardware has a switch connected to it
+# see if there is any information about the switch from the database
+# 1. check if the switch is reachable by an IP
+# 2. if reachable, check if the API is enabled
+# 3. if API is enabled, return True
+# 4. if API is not enabled
+# - check if we can do ssh connection
+# - if ssh connection is possible, enable the API, return True
+# - if ssh connection is not possible, check for serial connection
+# - if serial connection is possible, enable the API, return True
+# - if serial connection is not possible, return False
+# 5. if switch is not reachable by IP, SSH or Serial, return False
+# 6. if the switch is reachable, ask for the switch information
+# 7. json form that switch information and send it over rabbitmq
+
+
+def mac_address_manufacturer_style_to_ieee(mac: str) -> str:
+    """
+    Convert Arista-style MAC (xxxx.xxxx.xxxx) into IEEE style (xx:xx:xx:xx:xx:xx).
+    Example: '54b2.0319.7692' -> '54:b2:03:19:76:92'
+    """
+    # Remove dots
+    mac = mac.replace(".", "")
+    # Ensure correct length
+    if len(mac) != 12:
+        raise ValueError("Invalid Arista MAC format")
+    # Split into pairs
+    return ":".join(mac[i : i + 2] for i in range(0, 12, 2))
+
+
+class SwitchConnectionInfo(TypedDict):
+    vendor: str
+    hostname: str
+    username: str
+    password: str
+
+
+class Network(BaseAction):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.switch_tty_name = None
+
+        self.switch_connection_info: SwitchConnectionInfo | None = None
+        self.remote_switch = None
+        self.local_switch = None
+
+        # session info
+        # if it is set to None it has not been checked yet
+        self.switch_api_available: bool | None = None
+        self.switch_ssh_available: bool | None = None
+
+    def is_switch_api_enabled(self) -> bool:
+        if self.switch_api_available is not None:
+            return self.switch_api_available
+
+        self.switch_api_available = False
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            response = requests.get(
+                f"http://{self.switch_connection_info['hostname']}/eapi"
+            )
+            if response.ok:
+                self.switch_api_available = True
+        return self.switch_api_available
+
+    def is_switch_ssh_enabled(self) -> bool:
+        if self.switch_ssh_available is not None:
+            return self.switch_ssh_available
+
+        self.switch_ssh_available = False
+        if self.switch_connection_info:
+            ssh_client = SSHClient()
+            ssh_client.load_system_host_keys()
+
+            ssh_hostname = self.switch_connection_info.get("hostname", None)
+            ssh_username = self.switch_connection_info.get("username", None)
+            ssh_password = self.switch_connection_info.get("password", None)
+            if ssh_hostname and ssh_username and ssh_password:
+                try:
+                    ssh_client.connect(
+                        hostname=ssh_hostname,
+                        username=ssh_username,
+                        password=ssh_password,
+                        timeout=5,
+                    )
+                    self.switch_ssh_available = True
+                except Exception as exception:
+                    logger.error(f"Error connecting to switch via SSH: {exception}")
+        return self.switch_ssh_available
+
+    def arista_eapi_request_handler(self, command):
+        # example commands "show version" "show interfaces status"
+        if not self.switch_connection_info:
+            return None
+
+        url = f"http://{self.switch_connection_info['hostname']}/command-api"
+        with requests.Session() as session:
+            session.auth = (
+                self.switch_connection_info["username"],
+                self.switch_connection_info["password"],
+            )
+            session.headers.update({"Content-type": "application/json-rpc"})
+            payload = {
+                "jsonrpc": "2.0",
+                "method": "runCmds",
+                "params": {
+                    "version": "latest",
+                    "format": "json",
+                    "cmds": [command],
+                },
+                "id": 1,
+            }
+            with session:
+                response = session.post(
+                    url=url,
+                    json=payload,
+                )
+                if response.ok:
+                    return response.json()
+                else:
+                    logger.error(
+                        f"Error connecting to eAPI: {response.status_code} {response.text}"
+                    )
+                return None
+
+    def get_switch_info(self):
+        if self.is_switch_api_enabled():
+            switch_info = self.get_switch_info_via_api()
+            if switch_info:
+                return switch_info
+
+        return None
+
+    def get_interfaces_info(self):
+        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()
+
+            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][
+                                    "ip_address"
+                                ] = ip_arp.get("ip_address", "")
+
+                return interfaces_info
+
+        return None
+
+    def get_mac_address_info(self):
+        if self.is_switch_api_enabled():
+            mac_address_info = self.get_mac_address_info_via_api()
+            if mac_address_info:
+                return mac_address_info
+
+        return None
+
+    def get_switch_info_via_api(self):
+        formatted_switch_info = None
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            if response := self.arista_eapi_request_handler("show version"):
+                arista_version_info = response.get("result", [])[0]
+                # example output:
+                # {
+                #     "imageFormatVersion": "3.0",
+                #     "uptime": 6038503.06,
+                #     "modelName": "DCS-7050TX-64-R",
+                #     "internalVersion": "4.28.5.1M-30127723.42851M",
+                #     "memTotal": 3982512,
+                #     "mfgName": "Arista",
+                #     "serialNumber": "JPE16121065",
+                #     "systemMacAddress": "44:4c:a8:a3:61:77",
+                #     "bootupTimestamp": 1753302953.365648,
+                #     "memFree": 2546096,
+                #     "version": "4.28.5.1M",
+                #     "configMacAddress": "00:00:00:00:00:00",
+                #     "isIntlVersion": false,
+                #     "imageOptimization": "Strata-4GB",
+                #     "internalBuildId": "9adca383-a3bd-4507-b53b-d99ca7d61291",
+                #     "hardwareRevision": "01.11",
+                #     "hwMacAddress": "44:4c:a8:a3:61:77",
+                #     "architecture": "i686"
+                # },
+                formatted_switch_info = {
+                    "vendor": arista_version_info.get("mfgName", ""),
+                    "model": arista_version_info.get("modelName", ""),
+                    "serial_number": arista_version_info.get("serialNumber", ""),
+                    "mac_address": arista_version_info.get("systemMacAddress", ""),
+                    # "raw_output": arista_version_info,
+                }
+        return formatted_switch_info
+
+    def get_interfaces_via_api(self):
+        formatted_interfaces_info = None
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            if response := self.arista_eapi_request_handler("show interfaces status"):
+                # example output:
+                # {
+                #     "interfaceStatuses": {
+                #         "Ethernet10": {
+                #             "vlanInformation": {
+                #                 "interfaceMode": "bridged",
+                #                 "vlanId": 1,
+                #                 "interfaceForwardingModel": "bridged"
+                #             },
+                #             "bandwidth": 0,
+                #             "interfaceType": "10GBASE-T",
+                #             "description": "george-michael-oob",
+                #             "autoNegotiateActive": true,
+                #             "duplex": "duplexUnknown",
+                #             "autoNegotigateActive": true,
+                #             "linkStatus": "notconnect",
+                #             "lineProtocolStatus": "down"
+                #         },
+                #     }
+                # }
+                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,
+                }
+        return formatted_interfaces_info
+
+    def get_mac_address_info_via_api(self):
+        interface_to_mac_address_info = None
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            if response := self.arista_eapi_request_handler("show mac address-table"):
+                # {
+                #     "multicastTable": {"tableEntries": []},
+                #     "unicastTable": {
+                #         "tableEntries": [
+                #             {
+                #                 "macAddress": "20:37:f0:6b:d6:8c",
+                #                 "lastMove": 1759519474.903184,
+                #                 "interface": "Ethernet46",
+                #                 "moves": 1,
+                #                 "entryType": "dynamic",
+                #                 "vlanId": 1,
+                #             },
+                #         ]
+                #     },
+                #     "disabledMacLearningVlans": [],
+                # }
+
+                interface_to_mac_address_info = {}
+                table_entries = (
+                    response.get("result", [])[0]
+                    .get("unicastTable", [])
+                    .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
+
+        return interface_to_mac_address_info
+
+    def get_ip_arp_table_via_api(self):
+        ip_to_mac_address_info = {}
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            if response := self.arista_eapi_request_handler("show ip arp"):
+                table_entries = response.get("result", [])[0].get("ipV4Neighbors", [])
+                for entry in table_entries:
+                    interface = entry["interface"].split(", ")[
+                        1
+                    ]  # 'Vlan1, Ethernet46',
+                    mac_address = mac_address_manufacturer_style_to_ieee(
+                        entry["hwAddress"]
+                    )
+                    if interface not in ip_to_mac_address_info:
+                        ip_to_mac_address_info[interface] = []
+                    ip_to_mac_address_info[interface].append(
+                        {
+                            "ip_address": entry["address"],
+                            "mac_address": mac_address,
+                            "age": entry["age"],
+                        }
+                    )
+        return ip_to_mac_address_info
+
+    def serial_connect(self):
+        self.ser = serial.Serial()
+        self.ser.port = self.switch_tty_name
+        self.ser.baudrate = 9600
+        self.ser.open()
+
+    def serial_disconnect(self):
+        if self.ser.is_open:
+            self.ser.close()
+
+    def send_serial_command(self, command):
+        if not self.ser.is_open:
+            self.serial_connect()
+        self.ser.write(command.encode("utf-8") + b"\n")
+        response = self.ser.read_all().decode("utf-8")
+        return response
+
+    def get_tty_devices(self):
+        tty_devices = list(Path("/dev").glob("tty.*"))
+        return [str(device) for device in tty_devices]
+
+    def get_switch_tty_device_name(self):
+        if self.switch_tty_name is not None:
+            return self.switch_tty_name
+
+        tty_devices = self.get_tty_devices()
+        for device in tty_devices:
+            if "usbserial" in device or "usbmodem" in device:
+                self.switch_tty_name = device
+                break
+            if "ttyUSB0" in device:
+                self.switch_tty_name = device
+                break
+        return self.switch_tty_name
+
+    def arista_enable_api_via_ssh(self):
+        # TODO: implement this function
+        return True
+
+    def arista_enable_api_via_tty(self):
+        # TODO: implement this function
+        return True
+        # configure terminal
+        # management api http-commands
+        # protocol http
+        # no protocol https
+        # exit
+        # write memory
+
+    def enable_switch_api_via_ssh(self):
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            return self.arista_enable_api_via_ssh()
+        return False
+
+    def enable_switch_api_via_tty(self):
+        if (
+            self.switch_connection_info
+            and self.switch_connection_info["vendor"] == "arista"
+        ):
+            return self.arista_enable_api_via_tty()
+        return False
+
+    def enable_switch_api(self):
+        if self.is_switch_ssh_enabled():
+            result = self.enable_switch_api_via_ssh()
+            return result
+
+        elif self.switch_tty_name or self.get_switch_tty_device_name():
+            self.switch_tty_name = (
+                self.switch_tty_name or self.get_switch_tty_device_name()
+            )
+            result = self.enable_switch_api_via_tty()
+            return result
+
+        return False
+
+    def push_switch_and_interfaces_info(self):
+        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()
+
+            message = {"switch_info": {}, "interfaces_info": {}}
+            if switch_info:
+                message["switch_info"] = switch_info
+            if interfaces_info:
+                message["interfaces_info"] = interfaces_info
+
+            if message:
+                self.primitive.messaging.send_message(
+                    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,
+            )