ssm port forwarder wrapper

ssm-ssh.py

#!/usr/local/bin/python3
import argparse
import botocore.exceptions
import socket
import random
import shlex
import boto3
import json
import os


def generate_random_port(host: str):
    for item in range(15):
        random_port = random.randint(11000, 12000)
        if check_port(host, random_port):
            return random_port


def check_port(host: str, port: int):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    result = sock.connect_ex((host, port))
    sock.close()
    if result == 61:
        return True
    else:
        return False


def run_ssh(ssh_user: str, ssh_port: int):
    os.system(f'/usr/bin/ssh -o ConnectionAttempts=10 -o StrictHostKeyChecking=no {ssh_user}@localhost -p {ssh_port}')


def ssm_port_forward(aws_profile: str, aws_region: str, aws_instance_id: str, local_ssh_port: int):
    session = boto3.session.Session(profile_name=aws_profile)
    ssm = session.client('ssm', region_name=aws_region)
    _ssm_response = ssm.start_session(
        Target=aws_instance_id,
        DocumentName='AWS-StartPortForwardingSession',
        Parameters={"portNumber": ["22"], "localPortNumber": [f"{local_ssh_port}"]})
    return json.dumps(_ssm_response)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='SSH to instance via SSM port forward.')
    parser.add_argument('--instance-id', '-i', required=True, help='specify Instance Id')
    parser.add_argument('--profile', default='default', help='specify AWS Profile from ~/.aws/credentials')
    parser.add_argument('--region', default='us-east-1', help='specify AWS Region')
    parser.add_argument('--ssh-user', default='ec2-user', help='specify Linux user')
    try:
        args = parser.parse_args()
    except:
        parser.print_help()
        exit(0)

    random_port = generate_random_port('localhost')

    # create port forward via ssm tunnel
    ssm_response = ssm_port_forward(args.profile, args.region, args.instance_id, random_port)
    ssm_document = json.dumps({"Target": f"{args.instance_id}", "DocumentName": "AWS-StartPortForwardingSession",
                               "Parameters": {"portNumber": ["22"], "localPortNumber": [f"{random_port}"]}})
    ssm_endpoint = f'https://ssm.{args.region}.amazonaws.com'
    os.system(
        f'/usr/local/bin/session-manager-plugin \'{ssm_response}\' {args.region} StartSession {args.profile} \'{ssm_document}\' {ssm_endpoint}  1> /dev/null 2> /dev/null &')
    # run ssh to localhost
    print(run_ssh(args.ssh_user, random_port))