tracer-cloud--opensre/platform/deployment/aws/ssm.py

"""SSM Run Command and instance registration helpers for OpenSRE deployments."""

from __future__ import annotations

import logging
import time

from platform.deployment.aws.client import DEFAULT_REGION, get_boto3_client
from platform.deployment.aws.config import (
    SSM_CMD_POLL_ATTEMPTS,
    SSM_CMD_POLL_INTERVAL_SECONDS,
    SSM_REGISTRATION_MAX_ATTEMPTS,
    SSM_REGISTRATION_POLL_INTERVAL_SECONDS,
    SSM_SHELL_DOCUMENT,
    SSM_TERMINAL_STATUSES,
)

logger = logging.getLogger(__name__)


def wait_for_ssm_registration(
    instance_id: str,
    region: str = DEFAULT_REGION,
    poll_interval: int = SSM_REGISTRATION_POLL_INTERVAL_SECONDS,
    max_attempts: int = SSM_REGISTRATION_MAX_ATTEMPTS,
) -> bool:
    """Wait until the SSM agent on the instance registers and becomes online."""
    ssm = get_boto3_client("ssm", region)

    for attempt in range(max_attempts):
        try:
            resp = ssm.describe_instance_information(
                Filters=[{"Key": "InstanceIds", "Values": [instance_id]}]
            )
            instances = resp.get("InstanceInformationList", [])
            if instances and instances[0].get("PingStatus") == "Online":
                logger.info("SSM agent online for %s after %d attempts", instance_id, attempt + 1)
                return True
        except Exception as exc:  # noqa: BLE001
            logger.debug("SSM describe attempt %d: %s", attempt + 1, exc)

        if attempt < max_attempts - 1:
            time.sleep(poll_interval)

    raise TimeoutError(
        f"SSM agent on {instance_id} did not come online after {max_attempts * poll_interval}s"
    )


def run_ssm_shell_command(
    instance_id: str,
    commands: list[str],
    region: str = DEFAULT_REGION,
    poll_interval: int = SSM_CMD_POLL_INTERVAL_SECONDS,
    max_poll_attempts: int = SSM_CMD_POLL_ATTEMPTS,
) -> dict[str, str]:
    """Execute shell commands on an EC2 instance via SSM Run Command."""
    ssm = get_boto3_client("ssm", region)

    resp = ssm.send_command(
        InstanceIds=[instance_id],
        DocumentName=SSM_SHELL_DOCUMENT,
        Parameters={"commands": commands},
    )
    command_id = resp["Command"]["CommandId"]
    logger.debug("SSM command %s sent to %s", command_id, instance_id)

    for attempt in range(max_poll_attempts):
        time.sleep(poll_interval)
        try:
            inv = ssm.get_command_invocation(
                CommandId=command_id,
                InstanceId=instance_id,
            )
        except ssm.exceptions.InvocationDoesNotExist:
            logger.debug("SSM invocation %s not yet available, retrying...", command_id)
            continue

        status = inv["Status"]
        if status in SSM_TERMINAL_STATUSES:
            result = {
                "status": status,
                "stdout": inv.get("StandardOutputContent", ""),
                "stderr": inv.get("StandardErrorContent", ""),
            }
            logger.debug(
                "SSM command %s finished: status=%s stdout=%r",
                command_id,
                status,
                result["stdout"][:200],
            )
            return result

        logger.debug(
            "SSM command %s status=%s attempt=%d/%d",
            command_id,
            status,
            attempt + 1,
            max_poll_attempts,
        )

    raise TimeoutError(
        f"SSM command {command_id} on {instance_id} did not complete "
        f"within {max_poll_attempts * poll_interval}s"
    )