[Nexthop] Distro CLI device update skeleton

fboss-image/distro_cli/cmds/device.py

@@ -7,23 +7,52 @@
 
 """Device command implementation."""
 
+import json
 import logging
-
-from lib.cli import validate_path
+import os
+import sys
+
+from distro_cli.lib.cli import validate_path
+from distro_cli.lib.distro_infra import (
+    DISTRO_INFRA_CONTAINER,
+    GETIP_SCRIPT_CONTAINER_PATH,
+    deploy_image_to_device,
+    get_interface_name,
+)
+from distro_cli.lib.docker import container
+from distro_cli.lib.exceptions import DistroInfraError
 
 logger = logging.getLogger("fboss-image")
 
 
+def print_to_console(message: str) -> None:
+    """Print message to console"""
+    print(message)  # noqa: T201
+
+
 def image_upstream_command(args):
     """Download full image from upstream repository and set it to be loaded onto device"""
     logger.info(f"Setting upstream image for device {args.mac}")
     logger.info("Device image-upstream command (stub)")
 
 
 def image_command(args):
-    """Set device image from file"""
+    """Set device image from file and configure PXE boot"""
     logger.info(f"Setting image for device {args.mac}: {args.image_path}")
-    logger.info("Device image command (stub)")
+
+    try:
+        deploy_image_to_device(args.mac, args.image_path)
+        logger.info(
+            f"Successfully configured device {args.mac} with image {args.image_path}"
+        )
+        logger.info("Device is ready for PXE boot")
+
+    except DistroInfraError as e:
+        logger.error(f"Failed to configure device: {e}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        sys.exit(1)
 
 
 def reprovision_command(args):
@@ -40,16 +69,92 @@ def update_command(args):
     logger.info("Device update command (stub)")
 
 
+def get_device_ip(mac: str) -> str | None:
+    """Get device IP address by querying the distro-infra container.
+
+    Args:
+        mac: Device MAC address
+
+    Returns:
+        IP address string (IPv4 preferred, IPv6 fallback), or None if not found
+    """
+    if not container.container_is_running(DISTRO_INFRA_CONTAINER):
+        logger.error(f"Container '{DISTRO_INFRA_CONTAINER}' is not running")
+        logger.error("Please start the distro-infra container first")
+        return None
+
+    try:
+        interface = get_interface_name()
+    except DistroInfraError as e:
+        logger.error(f"Failed to get interface name: {e}")
+        return None
+
+    cmd = [GETIP_SCRIPT_CONTAINER_PATH, mac, interface]
+
+    # Execute in container
+    exit_code, stdout, stderr = container.exec_in_container(DISTRO_INFRA_CONTAINER, cmd)
+
+    if exit_code != 0:
+        logger.error(f"getip.sh failed with exit code {exit_code}")
+        if stderr:
+            logger.error(f"stderr: {stderr}")
+        if stdout:
+            logger.error(f"stdout: {stdout}")
+        return None
+
+    try:
+        result = json.loads(stdout)
+
+        if "error_code" in result:
+            logger.error(f"Error: {result.get('error', 'Unknown error')}")
+            logger.error(f"Error code: {result['error_code']}")
+            return None
+
+        ipv4 = result.get("ipv4")
+        ipv6 = result.get("ipv6")
+
+        return ipv4 if ipv4 else ipv6
+
+    except json.JSONDecodeError as e:
+        logger.error(f"Failed to parse JSON output: {e}")
+        logger.error(f"Output was: {stdout}")
+        return None
+
+
 def getip_command(args):
     """Get device IP address"""
     logger.info(f"Getting IP for device {args.mac}")
-    logger.info("Device getip command (stub)")
+
+    ip_address = get_device_ip(args.mac)
+
+    if ip_address:
+        print_to_console(ip_address)
+    else:
+        logger.error("No IP address found in response")
 
 
 def ssh_command(args):
     """SSH to device"""
     logger.info(f"SSH to device {args.mac}")
-    logger.info("Device ssh command (stub)")
+
+    ip_address = get_device_ip(args.mac)
+
+    if not ip_address:
+        logger.error("No IP address found for device")
+        return
+
+    logger.info(f"Connecting to {ip_address}")
+    os.execvp(
+        "ssh",
+        [
+            "ssh",
+            "-o",
+            "StrictHostKeyChecking=no",
+            "-o",
+            "UserKnownHostsFile=/dev/null",
+            f"root@{ip_address}",
+        ],
+    )
 
 
 def setup_device_commands(cli):
@@ -103,7 +208,13 @@ def setup_device_commands(cli):
     )
 
     device.add_command(
-        "getip", getip_command, help_text="Get device IP address", arguments=[]
+        "getip",
+        getip_command,
+        help_text="Get device IP address",
     )
 
-    device.add_command("ssh", ssh_command, help_text="SSH to device", arguments=[])
+    device.add_command(
+        "ssh",
+        ssh_command,
+        help_text="SSH to device",
+    )

fboss-image/distro_cli/lib/device_update.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2004-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree. An additional grant
+# of patent rights can be found in the PATENTS file in the same directory.
+
+"""Device update logic for updating FBOSS services on devices."""
+
+import logging
+from pathlib import Path
+
+from distro_cli.lib.exceptions import DistroInfraError
+from distro_cli.lib.manifest import ImageManifest
+
+logger = logging.getLogger(__name__)
+
+# Component to systemd services mapping.
+COMPONENT_SERVICES: dict[str, list[str]] = {
+    "fboss-forwarding-stack": ["wedge_agent", "fsdb", "qsfp_service"],
+    "fboss-platform-stack": [
+        "platform_manager",
+        "sensor_service",
+        "fan_service",
+        "data_corral_service",
+    ],
+}
+
+
+class DeviceUpdateError(DistroInfraError):
+    """Error during device update."""
+
+
+class DeviceUpdater:
+    """Handles updating FBOSS services on a device.
+
+    Workflow:
+    1. Validate component is supported for update
+    2. Acquire artifacts (build OR download)
+    3. Create update package (artifacts + service_update.sh) to scp to the device
+    4. Get device IP
+    5. SCP update package to device
+    6. SSH: extract and run service_update.sh
+    """
+
+    def __init__(
+        self,
+        mac: str,
+        manifest: ImageManifest,
+        component: str,
+        device_ip: str | None = None,
+    ):
+        """Initialize the DeviceUpdater.
+
+        Args:
+            mac: Device MAC address
+            manifest: Parsed image manifest
+            component: Component name to update
+            device_ip: Optional device IP (if already known)
+        """
+        self.mac = mac
+        self.manifest = manifest
+        self.component = component
+        self.device_ip = device_ip
+
+    def _get_services(self) -> list[str]:
+        """Get systemd services for the component."""
+        return COMPONENT_SERVICES.get(self.component, [])
+
+    def validate(self) -> None:
+        """Validate the update request.
+
+        Raises:
+            DeviceUpdateError: If validation fails
+        """
+        if self.component not in COMPONENT_SERVICES:
+            raise DeviceUpdateError(
+                f"Component '{self.component}' is not updatable. "
+                f"Updatable components: {', '.join(COMPONENT_SERVICES.keys())}"
+            )
+
+        if not self.manifest.has_component(self.component):
+            raise DeviceUpdateError(
+                f"Component '{self.component}' not found in manifest"
+            )
+
+        services = self._get_services()
+        if not services:
+            raise DeviceUpdateError(
+                f"Component '{self.component}' has no services defined in COMPONENT_SERVICES"
+            )
+
+        component_data = self.manifest.get_component(self.component)
+        has_download = "download" in component_data
+        has_execute = "execute" in component_data
+
+        if not has_download and not has_execute:
+            raise DeviceUpdateError(
+                f"Component '{self.component}' has neither 'download' nor 'execute'"
+            )
+
+    def _acquire_artifacts(self) -> Path:
+        """Acquire component artifacts via build or download.
+
+        Returns:
+            Path to the component artifact (tarball)
+
+        Raises:
+            DeviceUpdateError: If artifact acquisition fails
+        """
+        raise NotImplementedError("Stub")
+
+    def _create_update_package(self, artifact_path: Path) -> Path:
+        """Create update package with artifacts and update_service.sh script.
+
+        Args:
+            artifact_path: Path to the component artifact tarball
+
+        Returns:
+            Path to the created update package
+
+        Raises:
+            DeviceUpdateError: If package creation fails
+        """
+        raise NotImplementedError("Stub")
+
+    def _transfer_and_execute(self, package_path: Path, services: list[str]) -> None:
+        """Transfer update package to device and execute update_service.sh.
+
+        Args:
+            package_path: Path to the update package tarball
+            services: List of services to update
+
+        Raises:
+            DeviceUpdateError: If transfer or execution fails
+        """
+        raise NotImplementedError("Stub")
+
+    def update(self) -> bool:
+        """Execute the update workflow.
+
+        Returns:
+            True if update succeeded
+
+        Raises:
+            DeviceUpdateError: If update fails
+        """
+        raise NotImplementedError("Stub")