boto3 SSM PortForwarding

ssmportforward.py

"""
How to use boto3 for SSM local port forwarding.

In this example we are connecting to a MongoDB instance, but
the idea is the same regardless of connection.
"""
import os
import subprocess
import json
import boto3
from dotenv import load_dotenv
import pymongo

load_dotenv()

MONGO_INSTANCE_NAME = os.getenv("MONGO_INSTANCE")
MONGO_CONNECTION_STRING = os.getenv("MONGO_CONNECTION_STRING")
AWS_PROFILE = os.getenv("AWS_PROFILE")
MONGO_PORT_NUMBER = "27017"
LOCAL_PORT_NUMBER = "27017"

# Set up to use configured SSO profile
boto3.setup_default_session(profile_name=AWS_PROFILE)

# Get the right MongoDB Instance
ec2_client = boto3.client("ec2")
filters = [{"Name": "tag:Name", "Values": [MONGO_INSTANCE_NAME]}]
instances = ec2_client.describe_instances(Filters=filters)
mongo_instance_id = instances["Reservations"][0]["Instances"][0]["InstanceId"]

# Get an SSM client
ssm_client = boto3.client("ssm")

# Create Local port mapping to MongoDB
# See: https://stackoverflow.com/questions/66222667/how-to-use-session-manager-plugin-command/70311671#70311671
options = {
    "Target": mongo_instance_id,
    "DocumentName": "AWS-StartPortForwardingSession",
    "Reason": "Mongo Access via Python",
    "Parameters": {
        "portNumber": [MONGO_PORT_NUMBER],
        "localPortNumber": [LOCAL_PORT_NUMBER],
    },
}
# This just gives us some session details that need to be passed into session-manager-plugin
ssm_response = ssm_client.start_session(**options)

cmd = [
    "session-manager-plugin",
    json.dumps(ssm_response),
    "us-east-2",  # client region
    "StartSession",
    AWS_PROFILE,  # profile name from aws credentials/config files
    json.dumps(dict(Target=mongo_instance_id)),
    "https://ssm.us-east-2.amazonaws.com",  # endpoint for ssm service
]
pid = subprocess.Popen(cmd)

# Create Mongo Client
print(f"Connecting to: {MONGO_CONNECTION_STRING}")
mongo_client = pymongo.MongoClient(MONGO_CONNECTION_STRING)
for db in mongo_client.list_database_names():
    print(db)

# Gracefully stop the session-manager-plugin port-forwarding session
pid.terminate()