Improve AWS rate limit handling

* Start AWS SSM port-forwarding session using boto3 client instead of AWS CLI
* Bump version to 0.1.2

---------

Signed-off-by: Dennis Zagiansky <dennis@entitle.io>
Co-authored-by: avizetser <avi@entitle.io>

CHANGELOG.md

@@ -1,8 +1,12 @@
 # Change Log
 
+## [0.1.2] - 2023-11-15
+
+- Improve AWS rate limit handling by starting AWS SSM port-forwarding session using boto3 client instead of AWS CLI
+
 ## [0.1.1] - 2023-09-04
 
-Added basic documentation
+- Added basic documentation
 
 ## [0.1.0] - 2023-08-31
 
@@ -23,3 +27,4 @@ Initial release
 [Unreleased]: https://github.com/entitleio/beam/compare/0.1.0...master
 [0.1.0]: https://github.com/entitleio/beam/releases/tag/0.1.0
 [0.1.1]: https://github.com/entitleio/beam/releases/tag/0.1.1
+[0.1.2]: https://github.com/entitleio/beam/releases/tag/0.1.2

beam/ssm.py

@@ -1,10 +1,10 @@
-import socket
+import json
 import subprocess
 from typing import Optional
 
 import boto3
 
-from beam.utils import execute, logger
+from beam.utils import logger, execute
 
 
 def start_ssm_forwarding_session(region: str, instance_id: str, host: str, remote_port: int,
@@ -33,8 +33,7 @@ def start_ssm_forwarding_session(region: str, instance_id: str, host: str, remot
         This function requires AWS CLI and boto3 to be properly installed and configured with valid credentials.
 
     Example:
-        >>> start_ssm_forwarding_session('us-west-2', 'i-0123456789abcdef0', 'localhost', 22, 2222, 'my-profile')
-        True
+        start_ssm_forwarding_session('us-west-2', 'i-0123456789abcdef0', 'localhost', 22, 2222, 'my-profile')
     """
     logger.debug(f"Starting SSM session  (instance_id='{instance_id}', remote_port={remote_port}, local_port={local_port})")
 
@@ -55,26 +54,65 @@ def start_ssm_forwarding_session(region: str, instance_id: str, host: str, remot
 
     logger.debug(
         f' Starting SSM session to instance_id {instance_id} on port {remote_port} and local port {local_port}')
-
     try:
-        process = execute(
-            f'aws ssm --profile {profile} --region {region} start-session --target'
-            f' {instance_id} --document-name AWS-StartPortForwardingSessionToRemoteHost '
-            f'--parameters "{{\\"host\\": [ \\"{host}\\" ], \\"portNumber\\": [ \\"{remote_port}\\" ],'
-            f' \\"localPortNumber\\": [ \\"{local_port}\\" ] }}"')
-        return process
+        session = boto3.Session(profile_name=profile, region_name=region)
+        ssm_client = session.client('ssm')
+        ssm_parameters = {
+            'host': [host],
+            'portNumber': [str(remote_port)],
+            'localPortNumber': [str(local_port)],
+        }
+        response = ssm_client.start_session(
+            Target=instance_id,
+            DocumentName='AWS-StartPortForwardingSessionToRemoteHost',
+            Parameters=ssm_parameters
+        )
+
+        return start_aws_ssm_plugin(response, ssm_parameters, profile, region, instance_id)
     except subprocess.CalledProcessError as e:
         logger.error(f'Error executing command: {e.cmd}. Return code: {e.returncode}. Output: {e.output}')
 
-    ec2 = boto3.client('ec2', region_name=region)
-    try:
-        ec2.describe_instances(InstanceIds=[instance_id])
-    except ec2.exceptions.ClientError as e:
-        raise ValueError('instance_id is not a valid AWS identifier') from e
+    return None
+
+
+def start_aws_ssm_plugin(create_session_response: dict, parameters: dict, profile: str, region: str, instance_id: str) \
+        -> Optional[subprocess.Popen]:
+    """
+    Start the AWS SSM plugin to create a session and forward a local port to a remote port on an EC2 instance.
+
+    Args:
+        create_session_response: The response from creating an SSM session.
+        parameters: The parameters for the SSM session.
+        profile: The AWS CLI profile to be used for the SSM session.
+        region: The AWS region where the instance is located.
+        instance_id: The identifier of the EC2 instance to connect to.
+
+    Returns:
+        subprocess.Popen: The process for the SSM plugin command.
+
+    Raises:
+        subprocess.CalledProcessError: If there is an error executing the SSM plugin command.
+    """
+    plugin_parameters = {
+        'Target': instance_id,
+        'DocumentName': 'AWS-StartPortForwardingSessionToRemoteHost',
+        'Parameters': parameters
+    }
+
+    command = [
+        'session-manager-plugin',
+        f"'{json.dumps(create_session_response)}'",
+        region,
+        'StartSession',
+        profile,
+        f"'{json.dumps(plugin_parameters)}'",
+        f'https://ssm.{region}.amazonaws.com'
+    ]
 
     try:
-        socket.gethostbyname(host)
-    except socket.gaierror as e:
-        raise ValueError('Invalid hostname or IP address') from e
+        process = execute(' '.join(command))
+        return process
+    except subprocess.CalledProcessError as e:
+        logger.exception(f'Error executing command: {e.cmd} (return code: {e.returncode}) | Output: {e.output}')
 
     return None