Merge branch 'liw/deployment-refactor' of git://git.baserock.org/baserock/baserock/morph

1 files changed, 185 insertions, 0 deletions

diff --git a/morphlib/writeexts.py b/morphlib/writeexts.py
new file mode 100755
index 00000000..60848345
--- /dev/null
+++ b/morphlib/writeexts.py
@@ -0,0 +1,185 @@
+# Copyright (C) 2012-2013  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+
+import cliapp
+import os
+import re
+import time
+import tempfile
+
+
+class WriteExtension(cliapp.Application):
+
+    '''A base class for deployment write extensions.
+    
+    A subclass should subclass this class, and add a
+    ``process_args`` method.
+    
+    Note that it is not necessary to subclass this class for write
+    extensions. This class is here just to collect common code for
+    write extensions.
+    
+    '''
+    
+    def process_args(self, args):
+        raise NotImplementedError()
+
+    def status(self, **kwargs):
+        '''Provide status output.
+        
+        The ``msg`` keyword argument is the actual message,
+        the rest are values for fields in the message as interpolated
+        by %.
+        
+        '''
+        
+        self.output.write('%s\n' % (kwargs['msg'] % kwargs))
+    
+    def create_local_system(self, temp_root, raw_disk):
+        '''Create a raw system image locally.'''
+        
+        size = self.get_disk_size()
+        self.create_raw_disk_image(raw_disk, size)
+        try:
+            self.mkfs_btrfs(raw_disk)
+            mp = self.mount(raw_disk)
+        except BaseException:
+            sys.stderr.write('Error creating disk image')
+            os.remove(raw_disk)
+            raise
+        try:
+            self.create_factory(mp, temp_root)
+            self.create_fstab(mp)
+            self.create_factory_run(mp)
+            self.install_extlinux(mp)
+        except BaseException, e:
+            sys.stderr.write('Error creating disk image')
+            self.unmount(mp)
+            os.remove(raw_disk)
+            raise
+        else:
+            self.unmount(mp)
+
+    def get_disk_size(self):
+        '''Parse disk size from environment.'''
+        
+        size = os.environ.get('DISK_SIZE', '1G')
+        m = re.match('^(\d+)([kmgKMG]?)$', size)
+        if not m:
+            raise morphlib.Error('Cannot parse disk size %s' % size)
+
+        factors = {
+            '': 1,
+            'k': 1024,
+            'm': 1024**2,
+            'g': 1024**3,
+        }
+        factor = factors[m.group(2).lower()]
+
+        return int(m.group(1)) * factor
+
+    def create_raw_disk_image(self, filename, size):
+        '''Create a raw disk image.'''
+
+        self.status(msg='Creating empty disk image')
+        with open(filename, 'wb') as f:
+            if size > 0:
+                f.seek(size-1)
+                f.write('\0')
+
+    def mkfs_btrfs(self, location):
+        '''Create a btrfs filesystem on the disk.'''
+        self.status(msg='Creating btrfs filesystem')
+        cliapp.runcmd(['mkfs.btrfs', '-L', 'baserock', location])
+        
+    def mount(self, location):
+        '''Mount the filesystem so it can be tweaked.
+        
+        Return path to the mount point.
+        The mount point is a newly created temporary directory.
+        The caller must call self.unmount to unmount on the return value.
+        
+        '''
+
+        self.status(msg='Mounting filesystem')        
+        tempdir = tempfile.mkdtemp()
+        cliapp.runcmd(['mount', '-o', 'loop', location, tempdir])
+        return tempdir
+        
+    def unmount(self, mount_point):
+        '''Unmount the filesystem mounted by self.mount.
+        
+        Also, remove the temporary directory.
+        
+        '''
+        
+        self.status(msg='Unmounting filesystem')
+        cliapp.runcmd(['umount', mount_point])
+        os.rmdir(mount_point)
+
+    def create_factory(self, real_root, temp_root):
+        '''Create the default "factory" system.'''
+
+        factory = os.path.join(real_root, 'factory')
+
+        self.status(msg='Creating factory subvolume')
+        cliapp.runcmd(['btrfs', 'subvolume', 'create', factory])
+        self.status(msg='Copying files to factory subvolume')
+        cliapp.runcmd(['cp', '-a', temp_root + '/.', factory + '/.'])
+
+        # The kernel needs to be on the root volume.
+        self.status(msg='Copying boot directory to root subvolume')
+        factory_boot = os.path.join(factory, 'boot')
+        root_boot = os.path.join(real_root, 'boot')
+        cliapp.runcmd(['cp', '-a', factory_boot, root_boot])
+        
+    def create_factory_run(self, real_root):
+        '''Create the 'factory-run' snapshot.'''
+
+        self.status(msg='Creating factory-run subvolume')
+        factory = os.path.join(real_root, 'factory')
+        factory_run = factory + '-run'
+        cliapp.runcmd(
+            ['btrfs', 'subvolume', 'snapshot', factory, factory_run])
+
+    def create_fstab(self, real_root):
+        '''Create an fstab.'''
+
+        self.status(msg='Creating fstab')        
+        fstab = os.path.join(real_root, 'factory', 'etc', 'fstab')
+        with open(fstab, 'w') as f:
+            f.write('/dev/sda  /     btrfs defaults,rw,noatime 0 1\n')
+
+    def install_extlinux(self, real_root):
+        '''Install extlinux on the newly created disk image.'''
+
+        self.status(msg='Creating extlinux.conf')
+        config = os.path.join(real_root, 'extlinux.conf')
+        with open(config, 'w') as f:
+            f.write('default linux\n')
+            f.write('timeout 1\n')
+            f.write('label linux\n')
+            f.write('kernel /boot/vmlinuz\n')
+            f.write('append root=/dev/sda rootflags=subvol=factory-run '
+                    'init=/sbin/init rw\n')
+
+        self.status(msg='Installing extlinux')
+        cliapp.runcmd(['extlinux', '--install', real_root])
+
+        # FIXME this hack seems to be necessary to let extlinux finish
+        cliapp.runcmd(['sync'])
+        time.sleep(2)
+