3 files changed, 55 insertions, 1 deletions
diff --git a/ec2init/DataSource.py b/ec2init/DataSource.py
index 3ada110f..c761ae3e 100644
--- a/ec2init/DataSource.py
+++ b/ec2init/DataSource.py
@@ -20,3 +20,6 @@ class DataSource:
 
     def get_public_ssh_keys(self):
         return([])
+
+    def getswap_devs(self):
+        raise Exception("do not know how to generate swap list")
diff --git a/ec2init/DataSourceEc2.py b/ec2init/DataSourceEc2.py
index 9afbe8b9..123ede67 100644
--- a/ec2init/DataSourceEc2.py
+++ b/ec2init/DataSourceEc2.py
@@ -108,3 +108,14 @@ class DataSourceEc2(DataSource.DataSource):
                     keys.append(pkey)
 
         return(keys)
+
+    def getswap_devs(self):
+        if not self.metadata.has_key('block-device-mapping'):
+            raise Exception("no block-device-mapping")
+        list = []
+        for use_t, device in self.metadata['block-device-mapping'].items():
+            if not device.startswith("/dev"):
+                device="/dev/%s" % device
+            if use_t == "swap":
+                list.append(device)
+        return(list)
diff --git a/ec2init/__init__.py b/ec2init/__init__.py
index a0ef24c2..6816de60 100644
--- a/ec2init/__init__.py
+++ b/ec2init/__init__.py
@@ -25,6 +25,7 @@ import sys
 import os.path
 import errno
 import pwd
+import subprocess
 
 datadir = '/var/lib/cloud/data'
 semdir = '/var/lib/cloud/sem'
@@ -106,7 +107,6 @@ class EC2Init:
             return self.config.get(key, default)
 
     def initctl_emit(self):
-        import subprocess
         subprocess.Popen(['initctl', 'emit', 'cloud-config',
             '%s=%s' % (cfg_env_name,cloud_config)]).communicate()
 
@@ -249,6 +249,46 @@ class EC2Init:
 
         setup_user_keys(keys, 'root', key_prefix)
 
+    def enable_swap(self):
+        swaps=[]
+        try:
+            swaps=self.datasource.getswap_devs()
+        except:
+            print "using fstab"
+            process = subprocess.Popen(
+                ['blkid', '-t', 'TYPE=swap', '-o', 'device'],
+                stdout=subprocess.PIPE)
+            (out,err)=process.communicate()
+            swaps=out.strip().split('\n')
+
+        if len(swaps) == 0: return
+
+        fstab="/etc/fstab"
+        f=file(fstab,"rb")
+        lines=f.read().split('\n')
+        f.close()
+        existing=[]
+        for line in lines:
+            try:
+                (dev,mp,type,opts,dump,pss)=line.split()
+                if dev.startswith("#"): continue
+            except:
+                continue
+            existing.append(dev)
+
+        to_add=[]
+        for dev in swaps:
+            if not dev in existing:
+                to_add.append(dev)
+
+        if len(to_add) == 0 : return
+
+        f=file(fstab,"ab")
+        for dev in to_add:
+            f.write("%s\tnone\tswap\tsw\t0\t0\n" % dev)
+        f.close()
+
+        subprocess.Popen(['swapon', '-a']).communicate()
 
 def write_file(file,content,mode=0644):
         try: