[sonic_installer] Refactor sonic_installer code (#953)

Add a new Bootloader abstraction.
This makes it easier to add bootloader specific behavior while keeping
the main logic identical.
It is also a step that will ease the introduction of secureboot which
relies on bootloader specific behaviors.

Shuffle code around to get rid of the hacky if/else all over the place.
There are now 3 bootloader classes
 - AbootBootloader
 - GrubBootloader
 - UbootBootloader

There was almost no logic change in any of the implementations.
Only the AbootBootloader saw some small improvements.
More will follow in subsequent changes.

setup.py

@@ -49,6 +49,7 @@
         'pddf_ledutil',
         'show',
         'sonic_installer',
+        'sonic_installer.bootloader',
         'sonic-utilities-tests',
         'undebug',
         'utilities_common',

sonic_installer/bootloader/__init__.py

@@ -0,0 +1,16 @@
+
+from .aboot import AbootBootloader
+from .grub import GrubBootloader
+from .uboot import UbootBootloader
+
+BOOTLOADERS = [
+    AbootBootloader,
+    GrubBootloader,
+    UbootBootloader,
+]
+
+def get_bootloader():
+    for bootloaderCls in BOOTLOADERS:
+        if bootloaderCls.detect():
+            return bootloaderCls()
+    raise RuntimeError('Bootloader could not be detected')

sonic_installer/bootloader/aboot.py

@@ -0,0 +1,125 @@
+"""
+Bootloader implementation for Aboot used on Arista devices
+"""
+
+import collections
+import os
+import re
+import subprocess
+
+import click
+
+from ..common import (
+   HOST_PATH,
+   IMAGE_DIR_PREFIX,
+   IMAGE_PREFIX,
+   run_command,
+)
+from .bootloader import Bootloader
+
+_secureboot = None
+def isSecureboot():
+    global _secureboot
+    if _secureboot is None:
+        with open('/proc/cmdline') as f:
+           m  = re.search(r"secure_boot_enable=[y1]", f.read())
+        _secureboot = bool(m)
+    return _secureboot
+
+class AbootBootloader(Bootloader):
+
+    NAME = 'aboot'
+    BOOT_CONFIG_PATH = os.path.join(HOST_PATH, 'boot-config')
+    DEFAULT_IMAGE_PATH = '/tmp/sonic_image.swi'
+
+    def _boot_config_read(self, path=BOOT_CONFIG_PATH):
+        config = collections.OrderedDict()
+        with open(path) as f:
+            for line in f.readlines():
+                line = line.strip()
+                if not line or line.startswith('#') or '=' not in line:
+                    continue
+                key, value = line.split('=', 1)
+                config[key] = value
+        return config
+
+    def _boot_config_write(self, config, path=BOOT_CONFIG_PATH):
+        with open(path, 'w') as f:
+            f.write(''.join('%s=%s\n' % (k, v) for k, v in config.items()))
+
+    def _boot_config_set(self, **kwargs):
+        path = kwargs.pop('path', self.BOOT_CONFIG_PATH)
+        config = self._boot_config_read(path=path)
+        for key, value in kwargs.items():
+            config[key] = value
+        self._boot_config_write(config, path=path)
+
+    def _swi_image_path(self, image):
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        if isSecureboot():
+           return 'flash:%s/sonic.swi' % image_dir
+        return 'flash:%s/.sonic-boot.swi' % image_dir
+
+    def get_current_image(self):
+        with open('/proc/cmdline') as f:
+            current = re.search(r"loop=/*(\S+)/", f.read()).group(1)
+        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+
+    def get_installed_images(self):
+        images = []
+        for filename in os.listdir(HOST_PATH):
+            if filename.startswith(IMAGE_DIR_PREFIX):
+                images.append(filename.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX))
+        return images
+
+    def get_next_image(self):
+        config = self._boot_config_read()
+        match = re.search(r"flash:/*(\S+)/", config['SWI'])
+        return match.group(1).replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+
+    def set_default_image(self, image):
+        image_path = self._swi_image_path(image)
+        self._boot_config_set(SWI=image_path, SWI_DEFAULT=image_path)
+        return True
+
+    def set_next_image(self, image):
+        image_path = self._swi_image_path(image)
+        self._boot_config_set(SWI=image_path)
+        return True
+
+    def install_image(self, image_path):
+        run_command("/usr/bin/unzip -od /tmp %s boot0" % image_path)
+        run_command("swipath=%s target_path=/host sonic_upgrade=1 . /tmp/boot0" % image_path)
+
+    def remove_image(self, image):
+        nextimage = self.get_next_image()
+        current = self.get_current_image()
+        if image == nextimage:
+            image_path = self._swi_image_path(current)
+            self._boot_config_set(SWI=image_path, SWI_DEFAULT=image_path)
+            click.echo("Set next and default boot to current image %s" % current)
+
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        click.echo('Removing image root filesystem...')
+        subprocess.call(['rm','-rf', os.path.join(HOST_PATH, image_dir)])
+        click.echo('Image removed')
+
+    def get_binary_image_version(self, image_path):
+        try:
+            version = subprocess.check_output(['/usr/bin/unzip', '-qop', image_path, '.imagehash'])
+        except subprocess.CalledProcessError:
+            return None
+        return IMAGE_PREFIX + version.strip()
+
+    def verify_binary_image(self, image_path):
+        try:
+            subprocess.check_call(['/usr/bin/unzip', '-tq', image_path])
+            # TODO: secureboot check signature
+        except subprocess.CalledProcessError:
+            return False
+        return True
+
+    @classmethod
+    def detect(cls):
+        with open('/proc/cmdline') as f:
+            return 'Aboot=' in f.read()

sonic_installer/bootloader/bootloader.py

@@ -0,0 +1,50 @@
+"""
+Abstract Bootloader class
+"""
+
+class Bootloader(object):
+
+    NAME = None
+    DEFAULT_IMAGE_PATH = None
+
+    def get_current_image(self):
+        """returns name of the current image"""
+        raise NotImplementedError
+
+    def get_next_image(self):
+        """returns name of the next image"""
+        raise NotImplementedError
+
+    def get_installed_images(self):
+        """returns list of installed images"""
+        raise NotImplementedError
+
+    def set_default_image(self, image):
+        """set default image to boot from"""
+        raise NotImplementedError
+
+    def set_next_image(self, image):
+        """set next image to boot from"""
+        raise NotImplementedError
+
+    def install_image(self, image_path):
+        """install new image"""
+        raise NotImplementedError
+
+    def remove_image(self, image):
+        """remove existing image"""
+        raise NotImplementedError
+
+    def get_binary_image_version(self, image_path):
+        """returns the version of the image"""
+        raise NotImplementedError
+
+    def verify_binary_image(self, image_path):
+        """verify that the image is supported by the bootloader"""
+        raise NotImplementedError
+
+    @classmethod
+    def detect(cls):
+        """returns True if the bootloader is in use"""
+        return False
+

sonic_installer/bootloader/grub.py

@@ -0,0 +1,86 @@
+"""
+Bootloader implementation for grub based platforms
+"""
+
+import os
+import re
+import subprocess
+
+import click
+
+from ..common import (
+   HOST_PATH,
+   IMAGE_DIR_PREFIX,
+   IMAGE_PREFIX,
+   run_command,
+)
+from .onie import OnieInstallerBootloader
+
+class GrubBootloader(OnieInstallerBootloader):
+
+    NAME = 'grub'
+
+    def get_installed_images(self):
+        images = []
+        config = open(HOST_PATH + '/grub/grub.cfg', 'r')
+        for line in config:
+            if line.startswith('menuentry'):
+                image = line.split()[1].strip("'")
+                if IMAGE_PREFIX in image:
+                    images.append(image)
+        config.close()
+        return images
+
+    def get_next_image(self):
+        images = self.get_installed_images()
+        grubenv = subprocess.check_output(["/usr/bin/grub-editenv", HOST_PATH + "/grub/grubenv", "list"])
+        m = re.search(r"next_entry=(\d+)", grubenv)
+        if m:
+            next_image_index = int(m.group(1))
+        else:
+            m = re.search(r"saved_entry=(\d+)", grubenv)
+            if m:
+                next_image_index = int(m.group(1))
+            else:
+                next_image_index = 0
+        return images[next_image_index]
+
+    def set_default_image(self, image):
+        images = self.get_installed_images()
+        command = 'grub-set-default --boot-directory=' + HOST_PATH + ' ' + str(images.index(image))
+        run_command(command)
+        return True
+
+    def set_next_image(self, image):
+        images = self.get_installed_images()
+        command = 'grub-reboot --boot-directory=' + HOST_PATH + ' ' + str(images.index(image))
+        run_command(command)
+        return True
+
+    def install_image(self, image_path):
+        run_command("bash " + image_path)
+        run_command('grub-set-default --boot-directory=' + HOST_PATH + ' 0')
+
+    def remove_image(self, image):
+        click.echo('Updating GRUB...')
+        config = open(HOST_PATH + '/grub/grub.cfg', 'r')
+        old_config = config.read()
+        menuentry = re.search("menuentry '" + image + "[^}]*}", old_config).group()
+        config.close()
+        config = open(HOST_PATH + '/grub/grub.cfg', 'w')
+        # remove menuentry of the image in grub.cfg
+        config.write(old_config.replace(menuentry, ""))
+        config.close()
+        click.echo('Done')
+
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        click.echo('Removing image root filesystem...')
+        subprocess.call(['rm','-rf', HOST_PATH + '/' + image_dir])
+        click.echo('Done')
+
+        run_command('grub-set-default --boot-directory=' + HOST_PATH + ' 0')
+        click.echo('Image removed')
+
+    @classmethod
+    def detect(cls):
+        return os.path.isfile(os.path.join(HOST_PATH, 'grub/grub.cfg'))

sonic_installer/bootloader/onie.py

@@ -0,0 +1,48 @@
+"""
+Common logic for bootloaders using an ONIE installer image
+"""
+
+import os
+import re
+import signal
+import subprocess
+
+from ..common import (
+   IMAGE_DIR_PREFIX,
+   IMAGE_PREFIX,
+)
+from .bootloader import Bootloader
+
+# Needed to prevent "broken pipe" error messages when piping
+# output of multiple commands using subprocess.Popen()
+def default_sigpipe():
+    signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+
+class OnieInstallerBootloader(Bootloader): # pylint: disable=abstract-method
+
+    DEFAULT_IMAGE_PATH = '/tmp/sonic_image'
+
+    def get_current_image(self):
+        cmdline = open('/proc/cmdline', 'r')
+        current = re.search(r"loop=(\S+)/fs.squashfs", cmdline.read()).group(1)
+        cmdline.close()
+        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+
+    def get_binary_image_version(self, image_path):
+        """returns the version of the image"""
+        p1 = subprocess.Popen(["cat", "-v", image_path], stdout=subprocess.PIPE, preexec_fn=default_sigpipe)
+        p2 = subprocess.Popen(["grep", "-m 1", "^image_version"], stdin=p1.stdout, stdout=subprocess.PIPE, preexec_fn=default_sigpipe)
+        p3 = subprocess.Popen(["sed", "-n", r"s/^image_version=\"\(.*\)\"$/\1/p"], stdin=p2.stdout, stdout=subprocess.PIPE, preexec_fn=default_sigpipe)
+
+        stdout = p3.communicate()[0]
+        p3.wait()
+        version_num = stdout.rstrip('\n')
+
+        # If we didn't read a version number, this doesn't appear to be a valid SONiC image file
+        if not version_num:
+            return None
+
+        return IMAGE_PREFIX + version_num
+
+    def verify_binary_image(self, image_path):
+        return os.path.isfile(image_path)