support getting public ssh keys from ec2 metadata service

5 files changed, 61 insertions, 4 deletions

diff --git a/ec2-init.py b/ec2-init.py
index 464bf568..c335e7ca 100755
--- a/ec2-init.py
+++ b/ec2-init.py
@@ -37,7 +37,7 @@ def main():
     # TODO: cloud.set_defaults()
 
     # set the ssh keys up
-    # TODO: cloud.enable_authorized_keys()
+    cloud.apply_credentials()
 
     # finish, send the cloud-config event
     cloud.initctl_emit()
diff --git a/ec2init/DataSource.py b/ec2init/DataSource.py
index af5e9208..3ada110f 100644
--- a/ec2init/DataSource.py
+++ b/ec2init/DataSource.py
@@ -17,3 +17,6 @@ class DataSource:
 
     def get_userdata_raw(self):
         return(self.userdata_raw)
+
+    def get_public_ssh_keys(self):
+        return([])
diff --git a/ec2init/DataSourceEc2.py b/ec2init/DataSourceEc2.py
index cc12c97c..5d3bab88 100644
--- a/ec2init/DataSourceEc2.py
+++ b/ec2init/DataSourceEc2.py
@@ -120,5 +120,14 @@ class DataSourceEc2(DataSource.DataSource):
         elif availability_zone.startswith('eu-'):
             return 'eu'
         raise Exception('Could not determine location')
- 
 
+    def get_public_ssh_keys(self):
+        keys = []
+        if not self.metadata.has_key('public-keys'): return([])
+        for keyname, klist in self.metadata['public-keys'].items():
+            for pkey in klist:
+                # there is an empty string at the end of the keylist, trim it
+                if pkey:
+                    keys.append(pkey)
+
+        return(keys)
diff --git a/ec2init/UserDataHandler.py b/ec2init/UserDataHandler.py
index f7c56c69..71bc3203 100644
--- a/ec2init/UserDataHandler.py
+++ b/ec2init/UserDataHandler.py
@@ -111,7 +111,6 @@ def walk_userdata(str, callbacks, data = None):
         if not filename:
             filename = 'part-%03d' % partnum
 
-        print ":::::::: %s,%s :::::::" % (ctype,filename)
         if callbacks.has_key(ctype):
             callbacks[ctype](data,ctype,filename,part.get_payload())
 
diff --git a/ec2init/__init__.py b/ec2init/__init__.py
index 3d0ddbaa..e1ae87b0 100644
--- a/ec2init/__init__.py
+++ b/ec2init/__init__.py
@@ -24,6 +24,7 @@ import cPickle
 import sys
 import os.path
 import errno
+import pwd
 
 datadir = '/var/lib/cloud/data'
 semdir = '/var/lib/cloud/sem'
@@ -41,6 +42,7 @@ import UserDataHandler
 class EC2Init:
     datasource_list = [ DataSourceEc2.DataSourceEc2 ]
     part_handlers = { }
+    conffile = '/etc/ec2-init/ec2-config.cfg'
 
     def __init__(self):
         self.part_handlers = {
@@ -49,7 +51,8 @@ class EC2Init:
             'text/upstart-job' : self.handle_upstart_job,
             'text/part-handler' : self.handle_handler
         }
-
+        
+        self.config = ConfigObj(self.conffile)
 
     def restore_from_cache(self):
         try:
@@ -105,6 +108,9 @@ class EC2Init:
             return True
         return False
 
+    def get_cfg_option_str(self, key, default=None):
+            return self.config.get(key, default)
+
     def initctl_emit(self):
         import subprocess
         subprocess.Popen(['initctl', 'emit', 'cloud-config',
@@ -221,6 +227,26 @@ class EC2Init:
 
         self.cloud_config_str+="\n#%s\n%s" % (filename,payload)
 
+    def get_public_ssh_keys(self):
+        return(self.datasource.get_public_ssh_keys())
+
+    def apply_credentials(self):
+        user = self.get_cfg_option_str('user')
+        disable_root = self.get_cfg_option_bool('disable_root', True)
+        
+        keys = self.get_public_ssh_keys()
+
+        if user:
+            setup_user_keys(keys, user, '')
+     
+        if disable_root:
+            key_prefix = 'command="echo \'Please login as the ubuntu user rather than root user.\';echo;sleep 10" ' 
+        else:
+            key_prefix = ''
+
+        setup_user_keys(keys, 'root', key_prefix)
+
+
 def write_file(file,content,mode=0644):
         try:
             os.makedirs(os.path.dirname(file))
@@ -232,3 +258,23 @@ def write_file(file,content,mode=0644):
         f.write(content)
         f.close()
         os.chmod(file,mode)
+
+def setup_user_keys(keys, user, key_prefix):
+    saved_umask = os.umask(077)
+
+    pwent = pwd.getpwnam(user)
+
+    ssh_dir = '%s/.ssh' % pwent.pw_dir
+    if not os.path.exists(ssh_dir):
+        os.mkdir(ssh_dir)
+        os.chown(ssh_dir, pwent.pw_uid, pwent.pw_gid)
+
+    authorized_keys = '%s/.ssh/authorized_keys' % pwent.pw_dir
+    fp = open(authorized_keys, 'a')
+    fp.write(''.join(['%s%s\n' % (key_prefix, key) for key in keys]))
+    fp.close()
+
+    os.chown(authorized_keys, pwent.pw_uid, pwent.pw_gid)
+
+    os.umask(saved_umask)
+