primitive 0.2.66__tar.gz → 0.2.68__tar.gz

{primitive-0.2.66 → primitive-0.2.68}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.66
+Version: 0.2.68
 Project-URL: Documentation, https://github.com//primitivecorp/primitive-cli#readme
 Project-URL: Issues, https://github.com//primitivecorp/primitive-cli/issues
 Project-URL: Source, https://github.com//primitivecorp/primitive-cli

{primitive-0.2.66 → primitive-0.2.68}/src/primitive/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2025-present Dylan Stein <dylan@primitive.tech>
 #
 # SPDX-License-Identifier: MIT
-__version__ = "0.2.66"
+__version__ = "0.2.68"

{primitive-0.2.66 → primitive-0.2.68}/src/primitive/agent/runner.py

@@ -9,6 +9,7 @@ from datetime import datetime, timezone
 
 from loguru import logger
 
+from primitive.network.ssh import wait_for_ssh
 from primitive.utils.cache import get_artifacts_cache, get_logs_cache, get_sources_cache
 from primitive.utils.logging import fmt, log_context
 from primitive.utils.psutil import kill_process_and_children
@@ -231,6 +232,7 @@ class Runner:
 
         for i, cmd in enumerate(commands):
             if cmd == "oobpowercycle":
+                logger.info("Performing out-of-band power cycle")
                 from primitive.network.redfish import RedfishClient
 
                 bmc_host = self.target_hardware_secret.get("bmcHostname", None)
@@ -253,15 +255,16 @@ class Runner:
                 )
                 redfish.compute_system_reset(system_id="1", reset_type="ForceRestart")
                 if self.target_hardware_id:
-                    self.primitive.hardware.update_hardware(
+                    await self.primitive.hardware.aupdate_hardware(
                         hardware_id=self.target_hardware_id,
                         is_online=False,
                         is_rebooting=True,
-                        start_rebooting_at=datetime.now(timezone.utc),
+                        start_rebooting_at=str(datetime.now(timezone.utc)),
                     )
                 continue
 
             if cmd == "pxeboot":
+                logger.info("Setting next boot to PXE and rebooting")
                 from primitive.network.redfish import RedfishClient
 
                 bmc_host = self.target_hardware_secret.get("bmcHostname", None)
@@ -290,12 +293,19 @@ class Runner:
                 )
                 redfish.compute_system_reset(system_id="1", reset_type="ForceRestart")
                 if self.target_hardware_id:
-                    self.primitive.hardware.update_hardware(
+                    await self.primitive.hardware.aupdate_hardware(
                         hardware_id=self.target_hardware_id,
                         is_online=False,
                         is_rebooting=True,
-                        start_rebooting_at=datetime.now(timezone.utc),
+                        start_rebooting_at=str(datetime.now(timezone.utc)),
                     )
+                wait_for_ssh(
+                    hostname=self.target_hardware_secret.get("hostname"),
+                    username=self.target_hardware_secret.get("username"),
+                    password=self.target_hardware_secret.get("password"),
+                    port=22,
+                )
+                logger.info("PXE boot successful, SSH is now available")
                 continue
 
             args = ["/bin/bash", "-c", cmd]

{primitive-0.2.66 → primitive-0.2.68}/src/primitive/hardware/actions.py

@@ -981,3 +981,39 @@ class Hardware(BaseAction):
         logger.info(message)
 
         return result
+
+    @guard
+    async def aupdate_hardware(
+        self,
+        hardware_id: str,
+        is_online: Optional[bool] = None,
+        is_rebooting: Optional[bool] = None,
+        start_rebooting_at: Optional[str] = None,
+    ):
+        new_state: dict = {
+            "id": hardware_id,
+        }
+        if is_online is not None:
+            new_state["isOnline"] = is_online
+        if is_rebooting is not None:
+            new_state["isRebooting"] = is_rebooting
+        if start_rebooting_at is not None:
+            new_state["startRebootingAt"] = start_rebooting_at
+
+        mutation = gql(hardware_update_mutation)
+
+        input = new_state
+        variables = {"input": input}
+        try:
+            result = await self.primitive.session.execute_async(
+                mutation, variable_values=variables, get_execution_result=True
+            )
+        except client_exceptions.ClientConnectorError as exception:
+            message = "Failed to update hardware! "
+            logger.exception(message)
+            raise exception
+
+        message = "Updated hardware successfully! "
+        logger.info(message)
+
+        return result

primitive-0.2.68/src/primitive/network/ssh.py

@@ -0,0 +1,94 @@
+import paramiko
+import socket
+import time
+
+
+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 e:
+        print(f"SSH error connecting to {hostname}: {e}")
+        return False
+    except socket.error as e:
+        print(f"Socket error connecting to {hostname}: {e}")
+        return False
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        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
+        print(f"Waiting for SSH to become available on {hostname}...")
+        time.sleep(10)
+
+    print(
+        f"Timeout reached: Unable to connect to {hostname} via SSH within {timeout} seconds."
+    )
+    return False