Add script to modify the bootloader and manage different parallel OS.

This is part of the upgrades work. With this tool you
can now switch between versions of the OS, remove a
version, list all the versions present in the system, get
the default version and the running version, and deploy
a new system.

All of the above is possible with the following subcommands:
 - list
 - deploy
 - get-default
 - get-running
 - remove
 - set-default

It also activates a bootloader menu to choose a version
to boot. The menu is important to make sure the user can
boot the old OS if the new kernel doesn't work.

4 files changed, 298 insertions, 2 deletions

diff --git a/Makefile.am b/Makefile.am
index 746bd9c..e3d38ef 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -1,6 +1,6 @@
 # vi:set ts=8 sw=8 noet ai nocindent:
 # -
-# Copyright (c) 2011-2013 Codethink Ltd.
+# Copyright (c) 2011-2014 Codethink Ltd.
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License Version 2 as
@@ -22,6 +22,7 @@ SUBDIRS =								\
 	tb-switch							\
 	tb-update							\
 	baserock-system-config-sync					\
+	system-version-manager						\
 	tests
 
 .PHONY: ChangeLog
diff --git a/configure.ac b/configure.ac
index 17ec2dd..0e12f62 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1,6 +1,6 @@
 # vi:set et ai sw=2 sts=2 ts=2: */
 # -
-# Copyright (c) 2011-2012 Codethink Ltd.
+# Copyright (c) 2011-2014 Codethink Ltd.
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License Version 2 as
@@ -122,6 +122,7 @@ tbdiff-deploy/Makefile
 tb-switch/Makefile
 tb-update/Makefile
 baserock-system-config-sync/Makefile
+system-version-manager/Makefile
 tests/Makefile
 ])
 
diff --git a/system-version-manager/Makefile.am b/system-version-manager/Makefile.am
new file mode 100644
index 0000000..ddd0588
--- /dev/null
+++ b/system-version-manager/Makefile.am
@@ -0,0 +1,20 @@
+# vi:set ts=8 sw=8 noet ai nocindent:
+# -
+# Copyright (c) 2014 Codethink Ltd.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License Version 2 as
+# published by the Free Software Foundation.
+#
+# 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.
+# vi:set ts=8 sw=8 noet ai nocindent:
+
+bin_SCRIPTS =								\
+	system-version-manager
diff --git a/system-version-manager/system-version-manager b/system-version-manager/system-version-manager
new file mode 100755
index 0000000..e553562
--- /dev/null
+++ b/system-version-manager/system-version-manager
@@ -0,0 +1,274 @@
+#!/usr/bin/python
+#
+# Copyright (C) 2014  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 argparse
+import subprocess
+import tempfile
+import os
+import sys
+
+
+class AppException(Exception):
+    pass
+
+class SystemNotCompatibleError(Exception):
+    pass
+
+class SystemVersionManager(object):
+    def status(self, msg, *args):
+        print msg % args
+
+    def _check_system_exists(self, system_name):
+        systems_list = self._get_systems()
+
+        if system_name not in systems_list:
+            print "ERROR: the system " + system_name + " doesn't exist"
+            sys.exit(1)
+
+
+    # To get the systems the script lists the systems under the 'systems'
+    # folder which are directories and are not symlinks
+    def _get_systems(self):
+        systems = os.path.join(self.mount_dir, 'systems')
+        return [filename for filename in os.listdir(systems)
+            if os.path.isdir(os.path.join(systems,filename))
+            and not os.path.islink(os.path.join(systems,filename))]
+
+    # To check which system is the default one, it checks the 'ontimeout'
+    # value in the extlinux.conf file. If it's not present, then pick
+    # the first of the present systems.
+    def _get_default(self):
+        extlinux = os.path.join(self.mount_dir, 'extlinux.conf')
+        for line in open(extlinux,'r'):
+            line = line.rstrip('\n')
+            key, value= line.split(' ', 1)
+            if key == "ontimeout":
+                return value
+            if key == "label":
+                break
+
+        return self.current_system
+
+    def _rewrite_boot_menu(self, default, systems):
+        temp_config = os.path.join(self.mount_dir, 'extlinux.tmp')
+        config = os.path.join(self.mount_dir, 'extlinux.conf')
+        with open(temp_config, 'w') as f:
+            f.write('default menu.c32\n')
+            f.write('timeout 50\n')
+            f.write('prompt 0\n')
+            f.write('ontimeout ' + default +'\n')
+            for system in systems:
+                f.write('label ' + system +'\n')
+                f.write('kernel /systems/'+ system +'/kernel\n')
+                f.write('append root='+ self.device +' '
+                        'rootflags=subvol=systems/'+ system +'/run '
+                        'init=/sbin/init rw\n')
+        os.rename(temp_config,config)
+
+        default_path = os.path.join(self.mount_dir, 'systems', 'default')
+        if os.path.islink(default_path):
+            subprocess.call(['ln', '-sfn', default, default_path])
+
+    def list (self):
+        systems = self._get_systems()
+        for system in systems:
+            print system
+
+    def get_default(self):
+        print self._get_default()
+
+    def get_running(self):
+        print self.current_system
+
+    def _parse_deploy_location(self, location):
+        label, snapshot = os.path.split(location)
+        root, label = os.path.split(label)
+        if root != '/systems' or snapshot != 'orig':
+            raise AppException(
+                    "Invalid deploy location %s. Upgrades should be deployed "
+                    "from a Btrfs snapshot following the pattern: "
+                    "/systems/$LABEL/orig" % location)
+        return label
+
+    def deploy(self, location):
+        '''Client-side deployment of a new base OS version.
+
+        This code assumes that the 'orig' subvolume has been correctly created
+        already. In future it could be extended to receive a delta of the
+        upgraded version over network somehow.
+        '''
+        label = self._parse_deploy_location(location)
+        version_root = os.path.join(self.mount_dir, 'systems', label)
+
+        orig_dir = os.path.join(version_root, 'orig')
+        run_dir = os.path.join(version_root, 'run')
+        try:
+            self.status(msg="Creating 'run' subvolume")
+            subprocess.check_call(
+                    ['btrfs', 'subvolume', 'snapshot', orig_dir, run_dir])
+
+            self.status(msg='Updating system configuration')
+            log = os.path.join('/var', 'log', 'baserock-system-config-sync.log')
+
+            with open(log, 'w') as f:
+                subprocess.check_call(
+                        ['baserock-system-config-sync', 'merge', label],
+                        stdout=f)
+
+            self.status(msg="Installing the kernel")
+            self._install_kernel(version_root)
+
+        except Exception as e:
+            subprocess.call(['btrfs', 'subvolume', 'delete', run_dir])
+            raise
+
+        self._rewrite_boot_menu(self._get_default(), self._get_systems())
+
+    def _install_kernel(self, version_root):
+        '''Install the kernel outside of 'orig' or 'run' subvolumes
+
+        This code is kind of duplicated in morphlib/writeexts.py.
+
+        '''
+        image_names = ['vmlinuz', 'zImage', 'uImage']
+        kernel_dest = os.path.join(version_root, 'kernel')
+        for name in image_names:
+            try_path = os.path.join(version_root, 'run', 'boot', name)
+            if os.path.exists(try_path):
+                subprocess.check_call(['cp', '-a', try_path, kernel_dest])
+                break
+
+    def _get_mount_info(self):
+        mountpoint = subprocess.check_output(['findmnt', '/', '-l', '-n',
+            '-o', 'SOURCE'])
+        device, subvolume = mountpoint.split('[', 1)
+        subvolume = subvolume.split('/run',1)[0]
+        current_system = os.path.basename(subvolume)
+        return device, current_system
+
+    def remove (self, system_name):
+        self._check_system_exists(system_name)
+
+        default_system = self._get_default()
+
+        if system_name == default_system:
+            print "ERROR: you can't remove the default system"
+            sys.exit(1)
+        if system_name == self.current_system:
+            print "ERROR: you can't remove the running system"
+            sys.exit(1)
+
+        system_root = os.path.join(self.mount_dir, 'systems', system_name)
+
+        subprocess.call(['btrfs', 'subvolume', 'delete',
+            os.path.join(system_root, 'run')])
+        subprocess.call(['btrfs', 'subvolume', 'delete',
+            os.path.join(system_root, 'orig')])
+        subprocess.call(['rm', '-r', system_root])
+
+        self._rewrite_boot_menu(default_system, self._get_systems())
+
+    def set_default (self, system_name):
+        self._check_system_exists(system_name)
+        self._rewrite_boot_menu(system_name, self._get_systems())
+
+    def __init__(self):
+        self.device, self.current_system = self._get_mount_info()
+        self.mount_dir = tempfile.mkdtemp()
+
+    def mount_fs(self):
+        subprocess.check_call(['mount', self.device, self.mount_dir])
+
+    def umount_fs(self):
+        subprocess.call(['umount', self.mount_dir])
+
+    def _check_system_compatibility(self):
+        menu_file = os.path.join(self.mount_dir, 'menu.c32')
+        if not os.path.isfile(menu_file):
+            raise SystemNotCompatibleError("menu.c32 not found")
+
+    def main(self, args):
+
+        # create the top-level parser
+        parser = argparse.ArgumentParser(prog='system-version-manager')
+        subparsers = parser.add_subparsers(help='sub-command help')
+
+        # create the parser for the "list" command
+        parser_list = subparsers.add_parser('list',
+            help='list show you a list of systems')
+        parser_list.set_defaults(action = 'list')
+
+        # create the parser for the "deploy" command
+        parser_deploy= subparsers.add_parser('deploy',
+            help='deploy an updated base OS version to the running Baserock system')
+        parser_deploy.set_defaults(action = 'deploy')
+        parser_deploy.add_argument('location')
+
+        # create the parser for the "get-default" command
+        parser_get_default = subparsers.add_parser('get-default',
+            help='prints the default system')
+        parser_get_default.set_defaults(action = 'get-default')
+
+        # create the parser for the "get-running" command
+        parser_get_running = subparsers.add_parser('get-running',
+            help='prints the running system')
+        parser_get_running.set_defaults(action = 'get-running')
+
+        # create the parser for the "remove" command
+        parser_remove = subparsers.add_parser('remove', help='remove a system')
+        parser_remove.add_argument('system_name', help='name of the system to remove')
+        parser_remove.set_defaults(action = 'remove')
+
+        # create the parser for the "set-default" command
+        parser_set_default = subparsers.add_parser('set-default',
+            help='set a system as default')
+        parser_set_default.add_argument('system_name', help='name of the system to set')
+        parser_set_default.set_defaults(action = 'set-default')
+
+        args = parser.parse_args(args)
+        action = args.action
+
+        self.mount_fs()
+        try:
+            self._check_system_compatibility()
+
+            if action == "list":
+                self.list()
+            elif action == "deploy":
+                self.deploy(args.location)
+            elif action == "remove":
+                self.remove(args.system_name)
+            elif action == "set-default":
+                self.set_default(args.system_name)
+            elif action == "get-default":
+                self.get_default()
+            elif action == "get-running":
+                self.get_running()
+            else:
+                raise NotImplementedError("Unknown command %s" % action)
+        except SystemNotCompatibleError, e:
+            print "ERROR, system not compatible: %s" % e.args[0]
+            raise
+        finally:
+            self.umount_fs()
+
+
+if __name__ == "__main__":
+
+    sys_mgr = SystemVersionManager()
+    sys_mgr.main(sys.argv[1:])