Add SSH host key from EC2 console output

add-host-key.py

#!/usr/bin/env python3
# 
# Copyright (C) 2019 by Hunter Blanks <hblanks@artifex.org>
# 
# Permission to use, copy, modify, and/or distribute this software for
# any purpose with or without fee is hereby granted.
# 
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
# AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
# DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
# PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
# TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

"""
Adds an EC2 instance's SSH host key, as pulled from EC2 console output,
to your SSH known_hosts file.

Only 3rd-party dependency is boto3.
"""

from argparse import ArgumentParser
import subprocess
import logging
import os.path
import socket
import sys

import boto3

parser = ArgumentParser(description=__doc__.strip())
parser.add_argument('public_ip', help='Public IPv4 address to instance')
parser.add_argument('-v', dest='verbose', action='store_true')


def get_instance_by_addr(ec2_client, addr):
    """ Returns an EC2 instance id, given a public IPv4 address. """
    # NB: could do others here, such as:
    #   - dns-name
    #   - private-dns-name
    #   - private-ipaddress
    resp = c.describe_instances(
        Filters=[
            {'Name': 'ip-address', 'Values': [addr]},
        ])
    try:
        return resp['Reservations'][0]['Instances'][0]['InstanceId']    
    except IndexError:
        return

ST_BEFORE_HOST_KEYS = 0
ST_IN_HOST_KEYS = 1

def get_instance_pub_keys(ec2_client, instance_id):
    """ Returns host pubkeys using an instance's console out.  """
    resp = ec2_client.get_console_output(InstanceId=instance_id)
    if 'Output' in resp:
        keys = []
        state = ST_BEFORE_HOST_KEYS
        for line in resp['Output'].splitlines():
            if state == ST_BEFORE_HOST_KEYS:
                if '--BEGIN SSH HOST KEY KEYS' in line:
                    state = ST_IN_HOST_KEYS
                    continue

            if state == ST_IN_HOST_KEYS:
                if '--END SSH HOST KEY KEYS' in line:
                    break
                keys.append(line)

        return keys


def set_host_keys(addr, keys):
    """ Updates a host's pubkeys in ~/.ssh/known_hosts. """
    hostname = socket.gethostbyaddr(addr)[0]
    # Remove existing keys, if any.
    for arg in (addr, hostname):
        try:
            subprocess.check_output(
                ['ssh-keygen', '-R', arg], stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            sys.stderr.buffer.write(e.output)
            sys.stderr.buffer.write(b'\n')
            sys.stderr.buffer.flush()
            raise
    with open(os.path.expanduser('~/.ssh/known_hosts'), 'a') as f:
        for key in keys:
            assert not key.endswith('\n')
            f.write(
                '%s,%s %s\n' % (addr, hostname, key)
            )


if __name__ == '__main__':
    args = parser.parse_args()
    c = boto3.client('ec2')

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    instance_id = get_instance_by_addr(c, args.public_ip)
    if not instance_id:
        print(
            'Failed to find instance for ip %s' % args.public_ip,
            file=sys.stderr)
        sys.exit(1)

    keys = get_instance_pub_keys(c, instance_id)
    if not keys:
        print(
            'Failed to find keys in console output of %s' % instance_id,
            file=sys.stderr)
        sys.exit(1)

    set_host_keys(args.public_ip, keys)