support compressed user data, use cPickle instead of pickle

3 files changed, 51 insertions, 17 deletions

diff --git a/ec2init/DataSource.py b/ec2init/DataSource.py
index da6170fd..b940c564 100644
--- a/ec2init/DataSource.py
+++ b/ec2init/DataSource.py
@@ -1,8 +1,26 @@
 
+import ec2init
+
 class DataSource:
+    userdata = None
+    metadata = None
+    userdata_raw = None
+
     def __init__(self):
        pass
 
+    def store_user_data_raw(self):
+        fp=fopen(user_data_raw,"wb")
+        fp.write(self.userdata_raw)
+        fp.close()
+    
+    def store_user_data(self):
+        fp=fopen(user_data,"wb")
+        fp.write(self.userdata)
+        fp.close()
+
     def get_user_data(self):
-        raise Exception("get_user_data Not-implemented")
+        if self.userdata == None:
+            self.userdata = ec2init.preprocess_user_data(self.userdata_raw)
 
+        return self.userdata
diff --git a/ec2init/DataSourceEc2.py b/ec2init/DataSourceEc2.py
index c3317272..8ee92d29 100644
--- a/ec2init/DataSourceEc2.py
+++ b/ec2init/DataSourceEc2.py
@@ -5,7 +5,7 @@ import boto.utils
 import socket
 import urllib2
 import time
-import pickle
+import cPickle
 
 class DataSourceEc2(DataSource.DataSource):
     api_ver  = '2009-04-04'
@@ -26,26 +26,25 @@ class DataSourceEc2(DataSource.DataSource):
         self.meta_data_base_url = 'http://169.254.169.254/%s/meta-data' % self.api_ver
         self.user_data_base_url = 'http://169.254.169.254/%s/user-data' % self.api_ver
 
-    def get_user_data(self):
-        return("hello")
-
     def get_data(self):
-        print "checking %s" % self.cachedir + "/user-data.pkl"
-        udf = open(self.cachedir + "/user-data.pkl")
-        self.userdata = pickle.load(udf)
-        udf.close()
+        try:
+            udf = open(self.cachedir + "/user-data.pkl")
+            self.userdata_raw = cPickle.load(udf)
+            udf.close()
 
-        mdf = open(self.cachedir + "/meta-data.pkl")
-        self.metadata = pickle.load(mdf)
-        mdf.close()
+            mdf = open(self.cachedir + "/meta-data.pkl")
+            self.metadata = cPickle.load(mdf)
+            mdf.close()
 
-        return True
+            return True
+        except:
+            pass
 
         try:
             if not self.wait_for_metadata_service():
                 return False
             self.metadata = boto.utils.get_instance_userdata(api_ver)
-            self.userdata = boto.utils.get_instance_metadata(api_ver)
+            self.userdata_raw = boto.utils.get_instance_metadata(api_ver)
         except Exception as e:
             print e
             return False
diff --git a/ec2init/__init__.py b/ec2init/__init__.py
index 05404c3a..60a87275 100644
--- a/ec2init/__init__.py
+++ b/ec2init/__init__.py
@@ -21,6 +21,7 @@ import os
 from   configobj import ConfigObj
 
 import boto.utils
+import cPickle
 
 datadir = '/var/lib/cloud/data'
 cachedir = datadir + '/cache'
@@ -36,7 +37,7 @@ class EC2Init:
     def restore_from_cache(self):
         try:
             f=open(cachedir + "/obj.pkl", "rb")
-            data = pickle.load(f)
+            data = cPickle.load(f)
             self.datasource = data
             return True
         except:
@@ -45,7 +46,7 @@ class EC2Init:
     def write_to_cache(self):
         try:
             f=open(cachedir + "/obj.pkl", "wb")
-            data = pickle.dump(self.datasource,f)
+            data = cPickle.dump(self.datasource,f)
             return True
         except:
             return False
@@ -60,7 +61,8 @@ class EC2Init:
                 if s.get_data():
                     self.datasource = s
                     return
-            except:
+            except Exception as e:
+                print e
                 pass
         raise Exception("Could not find data source")
 
@@ -77,3 +79,18 @@ class EC2Init:
         import subprocess
         subprocess.Popen(['initctl', 'CFG_FILE=%s' % user_config]).communicate()
 
+
+# if 'str' is compressed return decompressed otherwise return it
+def decomp_str(str):
+    import StringIO
+    import gzip
+    try:
+        uncomp = gzip.GzipFile(None,"rb",1,StringIO.StringIO(str)).read()
+        return(uncomp)
+    except:
+        return(str)
+
+
+# preprocess the user data (include / uncompress)
+def preprocess_user_data(ud):
+   return(decomp_str(ud))