eicproxy

#!/usr/bin/env python3

import argparse
import boto3
import os
import subprocess
import sys

from shutil import which
from nclib import Netcat
from botocore.exceptions import ClientError
from os.path import expanduser, isfile, join

home = expanduser("~")
default_key_file_path_public = f'{home}/.ssh/id_rsa.pub'

default_aws_profile = boto3.session.Session().profile_name
print(str(boto3.session.Session().profile_name))
if not default_aws_profile:
    default_aws_profile = 'default'

default_region = boto3.session.Session().region_name
if not default_region:
    default_region = 'us-east-1'
default_regions = [default_region]

default_ssh_port = 22

default_use_private_ip = False
default_use_tag_name = False

parser = argparse.ArgumentParser(description=f'ssh ProxyCommand script that ec2 instance connect access using your IAM user has rights for and that are reachable.')
parser.add_argument('connection_str', type=str, help='Pass ssh TOKENS as a connection string: `eicproxy %r@%h:%p` ProxyCommand translates it to `eicproxy <User>@<HostName>:<Port>`.')
parser.add_argument('--public-key-file', type=argparse.FileType('r'),
                    default=default_key_file_path_public, help=f'Public key file to use for connection. Default: {default_key_file_path_public}')
parser.add_argument('--regions', type=str, nargs='+', default=default_regions, help=f'Look for the instance in the given regions. Default: {default_regions}')
parser.add_argument('--use-private-ip', action='store_true', help=f'Use private IP even if public IP is available. eicproxy will use private ip automatically, if no public IP is available.', default=default_use_private_ip)
parser.add_argument('--use-tag-name', action='store_true', help=f'Search for instance with Tag Name equal to the host instance_id parameter.', default=default_use_tag_name)
parser.add_argument('--jumphost', action='store', help='Proxy through a defined ssh config Host', type=str, metavar='')
parser.add_argument('--profile', action='store', help='AWS Config Profile', type=str, default=default_aws_profile, metavar='')
parser.add_argument('-t', '--target', action='store', help='Targe Instance ID')
parser.add_argument('-z', '--zone', action='store', help='Availability zone', type=str, metavar='')
args = parser.parse_args()

#print(str(args), file=sys.stderr)

def connect(region, profile_name):
    session = boto3.session.Session(profile_name=profile_name, region_name=region)
    ec2_client = session.client('ec2')
    print('tagname:',str(args.use_tag_name))
    if args.use_tag_name:
        tag_name = host_token
        lookup = session.client('resourcegroupstaggingapi')
        tag_name_result = lookup.get_resources(
            PaginationToken='',
            TagFilters=[
                {
                    'Key': 'Name',
                    'Values': [
                        tag_name,
                    ]
                },
            ],
            ResourcesPerPage=2,
            ResourceTypeFilters=[
                'ec2',
            ],
        )
        if tag_name_result['ResourceTagMappingList']:
            for resource in tag_name_result['ResourceTagMappingList']:
                if 'instance/' in resource['ResourceARN'].split(':')[5] :
                    print(resource['ResourceARN'].split('instance/')[1])
                    instance_id = resource['ResourceARN'].split('instance/')[1]
                    print(str(instance_id))
        else:
            print(f'Could not find instance id tagged Name:{tag_name}')
            return None
    else:
        instance_id = host_token

    try:
        print(f'Looking for instance {instance_id} in region {region}')
        response = ec2_client.describe_instances(InstanceIds=[instance_id])
    except ClientError as e:
        if e.response['Error']['Code'] != 'InvalidInstanceID.NotFound':
            raise
        else:
            print(f'Did not find instance {instance_id} in region {region}')
            return None

    instance = response['Reservations'][0]['Instances'][0]
    zone = instance['Placement']['AvailabilityZone']

    if args.use_private_ip or 'PublicIpAddress' not in instance:
        ip = instance['PrivateIpAddress']
    else:
        ip = instance['PublicIpAddress']

    connect_client = session.client('ec2-instance-connect')
    connect_client.send_ssh_public_key(
        InstanceId=instance_id,
        InstanceOSUser=os_user,
        SSHPublicKey=args.public_key_file.read(),
        AvailabilityZone=zone
    )
    return ip

profile = args.profile

str_id = args.connection_str
if '@' in str_id and ':' in str_id:
    os_user = str_id.split('@')[0]
    host_token = str_id.split('@')[1].split(':')[0]
    ssh_port = str_id.split('@')[1].split(':')[1]
else:
    print("invalid connection string, eicproxy %r@%h:%p or <ssh-user>@<ssh-host>:<ssh-port> is required", file=sys.stderr)


ip_to_connect_to = None

for region in args.regions:
    ip_to_connect_to = connect(region, profile)
    if ip_to_connect_to:
        break
if not ip_to_connect_to:
    print(f'Error: Did not find {host_token} in any region: {args.regions}')
    sys.exit(1)


if args.jumphost:
    str_jumphost = args.jumphost
    print(str_jumphost)
    command_list = ['ssh']
    command_list.extend([f'{str_jumphost}'])
    command_list.extend(['-W'])
    command_list.extend([f'{ip_to_connect_to}:{ssh_port}'])
    print(str(command_list))
    try:
        subprocess.run(command_list)
    except (BrokenPipeError, IOError):
        pass
else:
    # if nc is present use nc, else use nclib netcat.
    # if which('nc') is not None:
    #     print('using', str(which('nc')), str(ip_to_connect_to), str(ssh_port))
    #     command_list = ['nc']
    #     command_list.extend([f'{ip_to_connect_to}'])
    #     command_list.extend([f'{ssh_port}'])
    #     try:
    #         subprocess.run(command_list)
    #     except (BrokenPipeError, IOError):
    #         pass
    # else:
    #     print(f'netcat not found. using nclib Netcat')
    #     options = {'verbose': False, 'listen': False, 'listenmore': False}
    #     Netcat(server=(ip_to_connect_to, int(ssh_port)), verbose=options['verbose']).interact()

    print(f'netcat not found. using nclib Netcat')
    options = {'verbose': False, 'listen': False, 'listenmore': False}
    Netcat(server=(ip_to_connect_to, int(ssh_port)), verbose=options['verbose']).interact()

#sys.stderr.close()