#!/usr/bin/env python

from __future__ import print_function

import boto3
import argparse
import os


def main():
    args = parse_args()
    orignal_hostname = args.hostname.lower()
    hostname = orignal_hostname

    if not(args.suffix is None) and hostname.endswith(args.suffix.lower()):
        hostname = hostname[:-len(args.suffix)]

    if not(args.prefix is None) and hostname.startswith(args.prefix.lower()):
        hostname = hostname[len(args.prefix):]

    if args.verbose:
        print("Searching for", hostname)

    # Get all running instances from EC2
    ec2 = boto3.client(
        'ec2',
        aws_access_key_id=args.aws_access_key_id,
        aws_secret_access_key=args.aws_secret_access_key,
        aws_session_token=args.aws_session_token,
    )
    results = ec2.describe_instances(
        Filters=[
            {
                'Name': 'instance-state-name',
                'Values': [ 'running' ],
            },
        ]
    )

    # Find the instance that matches the hostname
    dns_name = find_dns_name(results, hostname)
    if args.verbose:
        if dns_name is None:
            print("dns name not found, defaulting to", orignal_hostname)
        else:
            print("dns name:", dns_name)

    # Ensure we always have a hostname to use
    if dns_name is None:
        dns_name = orignal_hostname

    # Execute the proxy command that will connect us to the host
    command = args.command.format(host=dns_name, port=args.port)
    exec_command = [ 'bash', '-c', command ]
    os.execvp(exec_command[0], exec_command)


def find_dns_name(results, name):
    if not results.get('Reservations'):
        return

    for reservation in results['Reservations']:
        for instances in reservation['Instances']:

            # Search for the Tags for: { Key: 'Name', Value: $name }
            for tag in instances['Tags']:
                if tag['Key'] == 'Name':
                    aws_name = tag['Value']
                    if aws_name is not None and str(aws_name).lower() == name:
                        return instances['PublicDnsName']


def parse_args():
    parser = argparse.ArgumentParser(
        description='%(prog)s [options] hostname [port]',
    )

    parser.add_argument(
        'hostname',
        help="the hostame",
    )

    parser.add_argument(
        'port',
        type=int,
        help="the port",
        nargs='?',
        default=22,
    )

    parser.add_argument(
        '--suffix', '-s',
        action="store",
        dest="suffix",
        help="suffix added to the hostame",
    )

    parser.add_argument(
        '--prefix', '-p',
        action="store",
        dest="prefix",
        help="prefix added to the hostame",
    )

    parser.add_argument(
        '--proxy-command', '-c',
        action="store",
        dest="command",
        help="comand to execute",
        default='nc {host} {port}',
    )

    parser.add_argument(
        '--verbose', '-v',
        action="store_true",
        dest="verbose",
        help="enable verbose mode",
    )

    parser.add_argument(
        '--aws-access-key-id', '-A',
        action="store",
        dest="aws_access_key_id",
        help="AWS access key id",
        default=os.environ.get('AWS_ACCESS_KEY_ID'),
    )

    parser.add_argument(
        '--aws-secret-access-key', '-S',
        action="store",
        dest="aws_secret_access_key",
        help="AWS secret access key",
        default=os.environ.get('AWS_SECRET_ACCESS_KEY'),
    )

    parser.add_argument(
        '--aws-session-token', '-T',
        action="store",
        dest="aws_session_token",
        help="AWS session token",
        default=os.environ.get('AWS_SESSION_TOKEN'),
    )

    return parser.parse_args()


if __name__ == '__main__':
    main()
