From 51c2448ae74c2c8579f4972dd740aef8da7aa19a Mon Sep 17 00:00:00 2001
From: Georgy Okrokvertskhov <gokrokvertskhov@mirantis.com>
Date: Sat, 23 Feb 2013 20:58:50 -0800
Subject: [PATCH] 1. Extended a list of EC2 values 2. Added proper handling of
 non-existent EC2 values 3. Added proper ssh-key handling

---
 cloudbaseinit/metadata/services/ec2service.py | 68 +++++++++++++++----
 1 file changed, 56 insertions(+), 12 deletions(-)

diff --git a/cloudbaseinit/metadata/services/ec2service.py b/cloudbaseinit/metadata/services/ec2service.py
index 146d60bd..83901add 100644
--- a/cloudbaseinit/metadata/services/ec2service.py
+++ b/cloudbaseinit/metadata/services/ec2service.py
@@ -17,6 +17,7 @@
 import posixpath
 import urllib2
 import traceback
+import os
 
 from cloudbaseinit.metadata.services import base
 from cloudbaseinit.openstack.common import cfg
@@ -28,9 +29,12 @@ opts = [
                help='The base URL where the service looks for metadata'),
   ]
 
-ec2nodes = ['ami-id', 'ami-launch-index', 'ami-manifest-path', 'hostname',
-            'instance-action', 'instance-id', 'instance-type',
-            'local-hostname', 'local-ipv4', 'public-hostname', 'public-ipv4']
+ec2nodes = ['ami-id', 'ami-launch-index', 'ami-manifest-path', 'ancestor-ami-ids',
+            'hostname', 'block-device-mapping', 'kernel-id','placement/availability-zone',
+            'instance-action', 'instance-id', 'instance-type', 'product-codes',
+            'local-hostname', 'local-ipv4', 'public-hostname', 'public-ipv4',
+            'ramdisk-id','reservation-id','security-groups',
+            'public-keys/','public-keys/0/','public-keys/0/openssh-key','admin_pass']
 
 CONF = cfg.CONF
 CONF.register_opts(opts)
@@ -41,6 +45,8 @@ class EC2Service(base.BaseMetadataService):
     def __init__(self):
         super(EC2Service, self).__init__()
         self._enable_retry = True
+        self.error_count = 0
+
 
     def load(self):
         super(EC2Service, self).load()
@@ -56,20 +62,58 @@ class EC2Service(base.BaseMetadataService):
 
     def _get_data(self, path):
         data = {}
+        LOG.debug("Check for EC2 interface availability...")
+        if not self._check_EC2():
+        	raise Exception("EC2 interface is not available")
+
         LOG.debug('Getting data for the path: %s' % path)
         if path.endswith('meta_data.json'):
-            meta_path = posixpath.join(CONF.ec2_metadata_base_url, 'meta-data')
             for key in ec2nodes:
-                norm_path = posixpath.join(meta_path, key)
-                LOG.debug('Getting metadata from: %(norm_path)s' % locals())
-                req = urllib2.Request(norm_path)
-                response = urllib2.urlopen(req)
-                data[key] = response.read()
+                LOG.debug('Getting metadata from: %s' % key)
+                try:
+                	data[key] = self._get_EC2_value(key)
+                except:
+                	LOG.info("EC2 value %s is not available. Skip it." % key)
+            # Saving keys to the local folder
+            self._load_public_keys(data)
 
         if path.endswith('user_data'):
             norm_path = posixpath.join(CONF.ec2_metadata_base_url, 'user-data')
             LOG.debug('Getting metadata from: %(norm_path)s' % locals())
-            req = urllib2.Request(norm_path)
-            response = urllib2.urlopen(req)
-            data = response.read()
+            try:
+               req = urllib2.Request(norm_path)
+               response = urllib2.urlopen(req)
+               data = response.read()
+               LOG.debug("Got data: %s" % data)
+            except:
+               LOG.error("EC2 user-data is not available.")
         return data
+
+    def _check_EC2(self):
+        try:
+          data = self._get_EC2_value('')
+          return True
+        except:
+          return False
+
+    def _get_EC2_value(self, key):
+       meta_path = posixpath.join(CONF.ec2_metadata_base_url, 'meta-data',key)
+       req = urllib2.Request(meta_path)
+       response = urllib2.urlopen(req)
+       return response.read()
+
+    def _load_public_keys(self, data):
+       try:
+          key_list = self._get_EC2_value('public-keys/')
+          LOG.debug("Got a list of keys %s" % key_list)
+          data['public_keys'] = {}
+          
+          for key_name in key_list.split('\n'):
+             key_index = key_name.split('=')[0]
+             LOG.debug('Loading key %s' % key_index)
+             key = self._get_EC2_value('public-keys/%s/openssh-key' % key_index)
+             data['public_keys'].update({key_index : key})
+
+       except Exception, ex:
+          LOG.debug("Can't save public key %s" % ex)
+          LOG.debug(traceback.format_exc())