#!/usr/bin/env python2.7

import argparse
import hashlib
import os
import sys
import urllib2
import boto

def get_deployment_pubkeys():
    pubkeyurl = 'http://instance-data/latest/meta-data/public-keys/'
    for line in urllib2.urlopen(pubkeyurl).readlines():
        line = line.strip()
        if not line:
            continue
        keyid = line.split('=')[0].strip()
        yield urllib2.urlopen(pubkeyurl + keyid + '/openssh-key').read().strip()

def get_s3_pubkeys(pubkey_bucket, pubkey_prefix):
    s3conn = boto.connect_s3()
    try:
        bucket = s3conn.get_bucket(pubkey_bucket)
    except boto.exceptions.S3ResponseError:
        sys.stderr.write("Specified bucket doesn't exist.\n")
        sys.exit(10)
    return dict([(os.path.basename(pk.name), pk) for pk in bucket.get_all_keys(prefix=pubkey_prefix) if pk.name.endswith('.pub')])

def update_authorized_keys(pubkey_bucket, pubkey_prefix, user, include_deployment_keys=False, include_deployment_local_keys=False):
    s3pubkeys = get_s3_pubkeys(pubkey_bucket, pubkey_prefix)

    if include_deployment_local_keys:
        s3pubkeys.update( get_s3_pubkeys(pubkey_bucket, '%s/%s' % (include_deployment_local_keys, pubkey_prefix)) )

    if not s3pubkeys:
        sys.stderr.write("No public keys found in the %s bucket with a prefix of %s.\n" % (pubkey_bucket, pubkey_prefix))
        sys.exit(10)

    localpubkeys = []

    home_dir = os.path.expanduser('~%s' % user)
    if not os.path.exists(home_dir):
        sys.stderr.write("Homedir for user %s, or user doesn't exist." % user)
        sys.exit(10)

    dot_ssh = os.path.join(home_dir, '.ssh')
    authorized_keys = os.path.join(dot_ssh, 'authorized_keys')

    for item in os.listdir(dot_ssh):
        filePath = os.path.join(dot_ssh, item)
        if os.path.isfile(filePath) and item.endswith('.pub'):
            if item not in s3pubkeys or s3pubkeys[item].etag.strip('"') != md5sum(filePath):
                os.remove(filePath)
            else:
                localpubkeys.append(filePath)
                s3pubkeys.pop(item)

    for filename,pkey in s3pubkeys.iteritems():
        filePath = os.path.join(dot_ssh, filename)
        pkey.get_contents_to_filename(filePath)
        if os.path.isfile(filePath):
            if (os.path.getsize(filePath) == pkey.size and
                pkey.etag.strip('"') == md5sum(filePath)):
                localpubkeys.append(filePath)
            else:
                os.remove(filePath)

    with open(authorized_keys, 'w') as f:
        if include_deployment_keys:
            for pubkey in get_deployment_pubkeys():
                f.write(pubkey + '\n')

        for pubkey in localpubkeys:
            with open(pubkey) as pk:
                f.write(pk.read().strip() + '\n')


def md5sum(filePath):
    h = hashlib.md5()
    with open(filePath,'r') as f:
        while True:
            data = f.read(4096)
            if not data: break
            h.update(data)
    return h.hexdigest()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('pubkey_bucket')
    parser.add_argument('pubkey_prefix')
    parser.add_argument('user')
    parser.add_argument('-d', dest="include_deployment_keys", action='store_true')
    parser.add_argument('-l', dest="include_deployment_local_keys", default=False)
    args = parser.parse_args()

    update_authorized_keys(args.pubkey_bucket, args.pubkey_prefix, args.user, args.include_deployment_keys, args.include_deployment_local_keys)


@@ -97,16 +97,45 @@
 ort boto
+%0Afrom tempfile import mkstemp
 %0A%0Adef ge
@@ -2494,37 +2494,93 @@
 h)%0A%0A
+%0A    temp_fd,temp_file = mkstemp(dir=dot_ssh)%0A    try:%0A 
     
+   
 with o
-pen(authorized_keys
+s.fdopen(temp_fd
 , 'w
@@ -2580,32 +2580,36 @@
 _fd, 'w') as f:%0A
+    
         if inclu
@@ -2632,32 +2632,36 @@
 ys:%0A            
+    
 for pubkey in ge
@@ -2692,32 +2692,36 @@
                 
+   
+ 
 f.write(pubkey +
@@ -2724,24 +2724,28 @@
 ey + '%5Cn')%0A%0A
+    
         for 
@@ -2776,24 +2776,28 @@
             
+  
+  
 with open(pu
@@ -2809,16 +2809,20 @@
  as pk:%0A
+    
         
@@ -2864,16 +2864,804 @@
  '%5Cn')%0A%0A
+        temp_file_size = os.path.getsize(temp_file)%0A        combined_size = sum(os.path.getsize(pk) for pk in localpubkeys)%0A        if temp_file_size and temp_file_size %3E= combined_size:%0A            if os.path.isfile(authorized_keys):%0A                orig_stats = os.stat(authorized_keys)%0A                os.chmod(temp_file, orig_stats.st_mode)%0A                os.chown(temp_file, orig_stats.st_uid, orig_stats.st_gid)%0A            os.rename(temp_file, authorized_keys)%0A    finally:%0A        if os.path.exists(temp_file):%0A            try:%0A                #not sure whether the fd is open or not. close it anyway%0A                #if it succeeds, yay. if it fails, whatever.%0A                os.close(temp_fd)%0A            except OSError:%0A                pass%0A            os.remove(temp_file)%0A%0A
 %0Adef md5

write new authorized_keys file to a temporary file first, so we never leave the server in a bad state