#!/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 SnapMgr(object): 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(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=/dev/sda ' '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 update (self): self._rewrite(self._get_default(), self._get_systems()) 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(default_system, self._get_systems()) def set_default (self, system_name): self._check_system_exists(system_name) self._rewrite(system_name, self._get_systems()) def __init__(self): self._closed = False self.device, self.current_system = self._get_mount_info() self.mount_dir = tempfile.mkdtemp() subprocess.call(['mount', self.device, self.mount_dir]) def close(self): self._closed = True subprocess.call(['umount', self.mount_dir]) def __del__(self): # No exception was raised, so no cleanup is required if self._closed: return self.close() # create the top-level parser parser = argparse.ArgumentParser(prog='snapshot-mgr') 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 "update" command parser_update = subparsers.add_parser('update', help='Update the extlinux.conf file with the present systems') parser_update.set_defaults(action = 'update') # create the parser for the "get-default" command parser_update = subparsers.add_parser('get-default', help='Prints the default system') parser_update.set_defaults(action = 'get-default') # create the parser for the "get-running" command parser_update = subparsers.add_parser('get-running', help='Prints the running system') parser_update.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_remove = subparsers.add_parser('set-default', help='set a system as default') parser_remove.add_argument('system_name', help='name of the system to set') parser_remove.set_defaults(action = 'set-default') args = parser.parse_args() action = args.action snap_manager = SnapMgr() if action == "list": snap_manager.list() elif action == "update": snap_manager.update() elif action == "remove": snap_manager.remove(args.system_name) elif action == "set-default": snap_manager.set_default(args.system_name) elif action == "get-default": snap_manager.get_default() elif action == "get-running": snap_manager.get_running() snap_manager.close()