[Nexthop] Distro CLI getip and ssh commands

fboss-image/distro_cli/cmds/device.py

@@ -7,13 +7,27 @@
 
 """Device command implementation."""
 
+import json
 import logging
+import os
 
-from lib.cli import validate_path
+from distro_cli.lib.cli import validate_path
+from distro_cli.lib.distro_infra import (
+    DISTRO_INFRA_CONTAINER,
+    GETIP_SCRIPT_CONTAINER_PATH,
+    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}")
@@ -40,16 +54,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 +193,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/distro_infra.py

@@ -0,0 +1,128 @@
+# 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.
+
+"""Distro Infrastructure helper functions."""
+
+import json
+import logging
+import re
+import subprocess
+from pathlib import Path
+
+from distro_cli.lib.docker import container
+from distro_cli.lib.exceptions import DistroInfraError
+
+logger = logging.getLogger("fboss-image")
+
+# This should match DISTRO_CONTAINER_NAME in distro_infra/distro_infra.sh
+DISTRO_INFRA_CONTAINER = "fboss-distro-infra"
+
+GETIP_SCRIPT_CONTAINER_PATH = "/distro_infra/getip.sh"
+
+
+def normalize_mac_address(mac: str) -> tuple[str, str]:
+    """Normalize MAC address to both dash and colon formats.
+
+    Args:
+        mac: MAC address in any format
+
+    Returns:
+        Tuple of (dash_format, colon_format)
+        e.g., ("aa-bb-cc-dd-ee-ff", "aa:bb:cc:dd:ee:ff")
+
+    Raises:
+        DistroInfraError: If MAC address is invalid
+    """
+    # Remove all separators and convert to lowercase
+    mac_clean = re.sub(r"[:\-]", "", mac.lower())
+
+    # Validate MAC address format (12 hex characters)
+    if not re.match(r"^[0-9a-f]{12}$", mac_clean):
+        raise DistroInfraError(
+            f"Invalid MAC address: {mac}. Expected 12 hex characters with optional colons or dashes."
+        )
+
+    # Convert to dash and colon formats
+    dash_mac = "-".join([mac_clean[i : i + 2] for i in range(0, 12, 2)])
+    colon_mac = ":".join([mac_clean[i : i + 2] for i in range(0, 12, 2)])
+
+    return dash_mac, colon_mac
+
+
+def get_interface_name() -> str:
+    """Get the network interface name of the distro-infra container from the persistent directory.
+
+    Returns:
+        Network interface name
+
+    Raises:
+        DistroInfraError: If interface_name.txt not found or empty
+    """
+    persistent_dir = find_persistent_dir()
+    interface_file = persistent_dir / "interface_name.txt"
+
+    if not interface_file.exists():
+        raise DistroInfraError(
+            f"Interface name file not found: {interface_file}. "
+            "The distro-infra container may not have started properly."
+        )
+
+    interface = interface_file.read_text().strip()
+    if not interface:
+        raise DistroInfraError(f"Interface name file is empty: {interface_file}")
+
+    return interface
+
+
+def find_persistent_dir() -> Path:
+    """Find the persistent directory mounted in the distro_infra container.
+
+    Returns:
+        Path to the persistent directory on the host
+
+    Raises:
+        DistroInfraError: If container is not running or persistent dir not found
+    """
+    # Check if container is running
+    if not container.container_is_running(DISTRO_INFRA_CONTAINER):
+        raise DistroInfraError(
+            f"Container '{DISTRO_INFRA_CONTAINER}' is not running. "
+            "Please start it first with distro_infra.sh"
+        )
+
+    try:
+        result = subprocess.run(
+            ["docker", "inspect", DISTRO_INFRA_CONTAINER],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        inspect_data = json.loads(result.stdout)
+
+        if not inspect_data:
+            raise DistroInfraError(
+                f"Container {DISTRO_INFRA_CONTAINER} is not running. "
+                "Please start it first with distro_infra.sh"
+            )
+
+        # Find the volume mount for /distro_infra/persistent
+        mounts = inspect_data[0].get("Mounts", [])
+        for mount in mounts:
+            if mount.get("Destination") == "/distro_infra/persistent":
+                return Path(mount["Source"])
+
+        raise DistroInfraError(
+            f"Could not find persistent directory mount in container {DISTRO_INFRA_CONTAINER}"
+        )
+
+    except subprocess.CalledProcessError as e:
+        raise DistroInfraError(
+            f"Container {DISTRO_INFRA_CONTAINER} is not running. "
+            "Please start it first with distro_infra.sh"
+        ) from e
+    except (json.JSONDecodeError, KeyError, IndexError) as e:
+        raise DistroInfraError(f"Failed to parse container inspect data: {e}") from e