#!/usr/bin/env python
"""Download a given SSH certificate and put it in the right place.

Downloads the certificate specified on argv and then searches through the
user's .ssh directory looking for a matching private key. Puts the certificate
next to that one.
"""
import os
import subprocess
import sys
import tempfile
import urllib


def download_cert_to_tempfile(url):
    resp = urllib.urlopen(url)
    temp_file = tempfile.NamedTemporaryFile(delete=False)
    with temp_file.file:
        temp_file.write(resp.read())
    return temp_file.name


def get_public_key_fingerprint(cert_path):
    proc = subprocess.Popen(['/usr/bin/ssh-keygen', '-L', '-f', cert_path],
        stdout=subprocess.PIPE)
    for line in proc.stdout.readlines():
        if 'Public key:' in line:
            fingerprint = line[line.find('RSA-CERT') + 9:]
            fingerprint = fingerprint.strip()
            return fingerprint


def find_private_key_for_public_key(pub_fingerprint):
    ssh_dir = os.getenv('HOME') + '/.ssh'
    for filename in os.listdir(ssh_dir):
        key_filename = ssh_dir + '/' + filename
        if key_filename.endswith('pub'):
            continue
        proc = subprocess.Popen(['/usr/bin/ssh-keygen', '-l', '-f',
            key_filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        for line in proc.stdout.readlines():
            if pub_fingerprint in line:
                return key_filename


def move_cert_into_place(cert_path, private_key_filename):
    new_cert_filename = private_key_filename + '-cert.pub'
    os.rename(cert_path, new_cert_filename)


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print 'Usage: %s <URL to cert>' % (sys.argv[0],)
        sys.exit(1)

    cert_filename = download_cert_to_tempfile(sys.argv[1])
    key_fingerprint = get_public_key_fingerprint(cert_filename)
    private_key_filename = find_private_key_for_public_key(key_fingerprint)
    if not private_key_filename:
        print 'Unable to find private key matching certificate.'
        sys.exit(1)

    move_cert_into_place(cert_filename, private_key_filename)
