Commit 0d0f5ebd authored by Apoorva Srinivas Appadoo's avatar Apoorva Srinivas Appadoo
Browse files

feat: add instance name resolving for ssm

parent 0e5e441a
Loading
Loading
Loading
Loading
+19 −13
Original line number Diff line number Diff line
@@ -127,8 +127,9 @@ This API starts an [AWS Systems Manager (SSM) port forwarding session](https://d
#### Query Parameters

| Name            | Description                                                   | Required              |
|---------------|---------------------------------------------------------------|-----------------------|
| `instance_id` | EC2 instance ID to connect to                                 | yes                   |
|-----------------|---------------------------------------------------------------|-----------------------|
| `instance_id`   | EC2 instance ID to connect to                                 | yes _(or `instance_name`)_ |
| `instance_name` | EC2 instance Name tag to connect to                           | yes _(or `instance_id`)_ |
| `remote_port`   | Port on the remote instance                                   | yes                   |
| `remote_host`   | Remote host for advanced port forwarding scenarios           | no                    |
| `local_port`    | Local port to bind (auto-allocated if not specified)          | no                    |
@@ -137,6 +138,8 @@ This API starts an [AWS Systems Manager (SSM) port forwarding session](https://d
| `env_ctx`       | the [environment context to consider](#the-notion-of-env_ctx) | no _(can be guessed)_ |
| `role_arn`      | AWS IAM role ARN to assume                                    | no _(can be retrieved from env)_ |

> **Note:** You must provide either `instance_id` or `instance_name`. If `instance_name` is provided, the API will look up the instance ID by the EC2 `Name` tag (only running instances are considered).

### `GET /kubeconfig`

This API generates a complete kubeconfig file for an AWS EKS cluster with a valid authentication token.
@@ -177,7 +180,10 @@ You can also **force SSM tunneling** for public clusters by providing the `insta
| `ttl_minutes`  | Time-to-live for the token in minutes (max: 15)               | no                    | `15`            |
| `namespace`      | Default namespace for kubectl context                         | no                    | `default`       |
| `user_name`      | Username in the kubeconfig                                    | no                    | `kubectl-user`  |
| `instance_id`  | EC2 instance ID for SSM port forwarding                       | **yes** _(for private clusters only)_ | -               |
| `instance_id`    | EC2 instance ID for SSM port forwarding                       | **yes** _(for private clusters only, or use `instance_name`)_ | -               |
| `instance_name`  | EC2 instance Name tag for SSM port forwarding                 | **yes** _(for private clusters only, or use `instance_id`)_ | -               |

> **Note:** For private clusters, you must provide either `instance_id` or `instance_name`. If `instance_name` is provided, the API will look up the instance ID by the EC2 `Name` tag.

#### Example Usage

+74 −3
Original line number Diff line number Diff line
@@ -40,6 +40,60 @@ class SsmPortForwardResponse:
    plugin_local_port: Optional[int] = None


def resolve_instance_name_to_id(instance_name: str) -> str:
    """
    Resolve an EC2 instance name (Name tag) to its instance ID.

    Args:
        instance_name: The Name tag value of the EC2 instance

    Returns:
        The instance ID

    Raises:
        HTTPException: If the instance is not found or multiple instances match
    """
    logger.debug(f"Resolving instance name '{instance_name}' to instance ID")
    ec2_client = boto3.client("ec2")
    
    try:
        response = ec2_client.describe_instances(
            Filters=[
                {"Name": "tag:Name", "Values": [instance_name]},
                {"Name": "instance-state-name", "Values": ["running"]},
            ]
        )
        
        instances = []
        for reservation in response.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                instances.append(instance["InstanceId"])
        
        if not instances:
            logger.error(f"No running instance found with name '{instance_name}'")
            raise HTTPException(
                status_code=404,
                detail=f"No running instance found with name '{instance_name}'",
            )
        
        if len(instances) > 1:
            logger.warning(
                f"Multiple instances found with name '{instance_name}': {instances}. Using first one: {instances[0]}"
            )
        
        instance_id = instances[0]
        logger.info(f"Resolved instance name '{instance_name}' to ID '{instance_id}'")
        return instance_id
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error resolving instance name: {str(e)}")
        raise HTTPException(
            status_code=500, detail=f"Error resolving instance name: {str(e)}"
        )


def is_port_available(port: int) -> bool:
    """Check if a port is available for binding."""
    try:
@@ -99,6 +153,7 @@ def start_ssm_port_forward(
    region: str = None,
    role_arn: str = None,
    instance_id: str = None,
    instance_name: str = None,
    remote_port: int = None,
    remote_host: str = None,
    local_port: int = None,
@@ -112,7 +167,8 @@ def start_ssm_port_forward(
        env_ctx: Environment context (PROD, STAGING, INTEG, REVIEW)
        region: AWS region
        role_arn: AWS role ARN to assume
        instance_id: EC2 instance ID to connect to
        instance_id: EC2 instance ID to connect to (either instance_id or instance_name required)
        instance_name: EC2 instance Name tag to connect to (either instance_id or instance_name required)
        remote_port: Remote port on the instance
        remote_host: Remote host (optional)
        local_port: Local port to forward to (optional, will use remote_port if not specified)
@@ -120,7 +176,24 @@ def start_ssm_port_forward(

    Returns:
        SsmPortForwardResponse object with session details

    Raises:
        HTTPException: If neither instance_id nor instance_name is provided
    """
    # Validate that either instance_id or instance_name is provided
    if not instance_id and not instance_name:
        raise HTTPException(
            status_code=400,
            detail="Either instance_id or instance_name must be provided",
        )

    # Configure boto before resolving instance name (need credentials)
    configure_boto(env_ctx, region, role_arn)

    # Resolve instance_name to instance_id if needed
    if not instance_id and instance_name:
        instance_id = resolve_instance_name_to_id(instance_name)

    logger.info(
        f"Starting SSM port forward: instance={instance_id}, remote_port={remote_port}, local_port={local_port}, protocol={protocol}"
    )
@@ -155,8 +228,6 @@ def start_ssm_port_forward(
                        plugin_local_port=session_data.get("plugin_local_port"),
                    )

    configure_boto(env_ctx, region, role_arn)

    # Determine local port: try requested port, then remote port, then find available
    if local_port is None:
        # No local port specified, try to use the same port as remote