rpi-eeprom-config: Add new option to apply config and update in one operation

* Add -a/--apply parameter which provides a one shot image for applying
  a new configuration to the latest bootloader image and installs it
  via rpi-eeprom-update.
* Print the live configuration if no arguments are specified.
* Add short flags instead of requiring verbose names.
* Make the errors more human readable instead of Exceptions

rpi-eeprom-config

@@ -1,12 +1,17 @@
 #!/usr/bin/env python
 
-# rpi-eeprom-config
+"""
-# Utility for reading and writing the configuration file in the
+rpi-eeprom-config
-# Raspberry Pi 4 bootloader EEPROM image.
+"""
 
 import argparse
+import atexit
+import os
+import subprocess
 import struct
 import sys
+import tempfile
+import time
 
 IMAGE_SIZE = 512 * 1024
 
@@ -22,17 +27,96 @@ MAGIC_MASK = 0xfffff00f
 FILE_MAGIC = 0x55aaf11f # id for modifiable file, currently only bootconf.txt
 FILE_HDR_LEN = 20
 FILENAME_LEN = 12
+TEMP_DIR = None
+
+def exit_handler():
+    """
+    Delete any temporary files.
+    """
+    if TEMP_DIR is not None and os.path.exists(TEMP_DIR):
+        tmp_image = '%s/pieeprom.upd' % TEMP_DIR
+        if os.path.exists(tmp_image):
+            os.remove(tmp_image)
+        os.rmdir(TEMP_DIR)
+
+def exit_error(msg):
+    """
+    Trapped a fatal arror, output message to stderr and exit with non-zero
+    return code.
+    """
+    sys.stderr.write("ERROR: %s\n" % msg)
+    sys.exit(1)
+
+def shell_cmd(args):
+    start = time.time()
+    arg_str = ' '.join(args)
+    result = subprocess.Popen(args, stdout=subprocess.PIPE, stdin=subprocess.PIPE)
+
+    while time.time() - start < 5:
+        if result.poll() is not None:
+            break
+
+    if result.poll is None:
+        exit_error("%s timeout" % arg_str)
+
+    if result.returncode != 0:
+        exit_error("%s failed: %d\n %s\n %s\n" %
+                   (arg_str, result.returncode, result.stdout.read(), result.stderr.read()))
+    else:
+        return result.stdout.read()
+
+def get_latest_eeprom():
+    """
+    Returns the path of the latest EEPROM image file if it exists.
+    """
+    latest = shell_cmd(['rpi-eeprom-update', '-l']).rstrip()
+    if not os.path.exists(latest):
+        exit_error("EEPROM image '%s' not found" % latest)
+    return latest
+
+def apply_update(config):
+    """
+    Applies the config file to the latest available EEPROM image and spawns
+    rpi-eeprom-update to schedule the update at the next reboot.
+    """
+    global TEMP_DIR
+    latest = get_latest_eeprom()
+    TEMP_DIR = tempfile.mkdtemp()
+    tmp_update = "%s/%s" % (TEMP_DIR, 'pieeprom.upd')
+    image = BootloaderImage(latest, tmp_update)
+    image.write(config)
+
+    sys.stdout.write("Updating bootloader EEPROM\n image: %s\nconfig: %s\n" % (latest, config))
+    args = ['sudo', 'rpi-eeprom-update', '-d', '-f', tmp_update]
+    resp = shell_cmd(args)
+    sys.stdout.write(resp)
+
+def read_current_config():
+    """
+    Reads the configuration used by the current bootloader.
+    """
+    result = shell_cmd(['vcgencmd', 'bootloader_config'])
+    if result is None:
+        exit_error('Failed to read the current bootloader configuration')
+    return result
 
 class BootloaderImage(object):
     def __init__(self, filename, output):
+        """
+        Instantiates a Bootloader image writer with a source eeprom (filename)
+        and optionally an output filename.
+        """
         self._filename = filename
+        try:
             self._bytes = bytearray(open(filename, 'rb').read())
+        except IOError as err:
+            exit_error("Failed to read \'%s\'\n%s\n" % (filename, str(err)))
         self._out = None
         if output is not None:
             self._out = open(output, 'wb')
 
         if len(self._bytes) != IMAGE_SIZE:
-            raise Exception("%s: Expected size %d bytes actual size %d bytes" %
+            exit_error("%s: Expected size %d bytes actual size %d bytes" %
                        (filename, IMAGE_SIZE, len(self._bytes)))
 
     def find_config(self):
@@ -51,7 +135,7 @@ class BootloaderImage(object):
             offset += 8 + length # length + type
             offset = (offset + 7) & ~7
 
-        raise Exception('Bootloader config not found')
+        raise Exception('EEPROM parse error: Bootloader config not found')
 
     def write(self, new_config):
         hdr_offset, length = self.find_config()
@@ -64,7 +148,8 @@ class BootloaderImage(object):
             raise Exception('EEPROM image size exceeded')
 
         struct.pack_into('>L', self._bytes, hdr_offset + 4, new_len)
-        struct.pack_into(("%ds" % len(new_config_bytes)), self._bytes, hdr_offset + 4 + FILE_HDR_LEN, new_config_bytes)
+        struct.pack_into(("%ds" % len(new_config_bytes)), self._bytes,
+                         hdr_offset + 4 + FILE_HDR_LEN, new_config_bytes)
 
         # If the new config is smaller than the old config then set any old
         # data which is now unused to all ones (erase value)
@@ -97,34 +182,69 @@ class BootloaderImage(object):
                 sys.stdout.write(config_bytes)
 
 def main():
-    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, \
+    """
-    description='Bootloader EEPROM configuration tool for the Raspberry Pi 4. \
+    Utility for reading and writing the configuration file in the
-\n\nThere are 3 operating modes: \
+    Raspberry Pi 4 bootloader EEPROM image.
-\n\n1. Output the bootloader configuration stored in an EEPROM image file to \
+    """
-the screen (STDOUT): specify only the name of an EEPROM image file using the \
+    description = """\
-\'eeprom\' option. \
+Bootloader EEPROM configuration tool for the Raspberry Pi 4.
-\n\n2. Output the bootloader configuration stored in an EEPROM image file to a \
+Operating modes:
-file: specify the EEPROM image file using the \'eeprom\' option, and the output \
+
-file using the \'--out\' option.\
+1. Output the current bootloader to configuration STDOUT if no arguments are
-\n\n3. Insert a new bootloader configuration into an EEPROM image file: specify \
+   specified OR the given output file if --out is specified.
-the source EEPROM image file using the \'eeprom\' option and the bootloader \
+
-configuration file using the \'--config\' option. A new file which is a \
+   rpi-eeprom-config (--out boot.conf)
-combination of the EEPROM image file, together with the new bootloader \
+
-configuration file will be created - specify its name using the \'--out\' option. \
+2. Extracts the configuration file from the given \'eeprom\' file and outputs
-The new bootloader configuration will replace any configuration present in the \
+   the result to STDOUT or the output file if --output is specified.
-source EEPROM image.\
+
-\n\nBootloader EEPROM images are contained in the \'rpi-eeprom-images\' package,\
+   rpi-eeprom-config pieeprom.bin (--out boot.conf)
- which installs them to the /lib/firmware/raspberrypi/bootloader directory.')
+
-    parser.add_argument('--config', help='Name of bootloader configuration file')
+3. Writes a new EEPROM image replacing the configuration file with the contents
-    parser.add_argument('--out', help='Name of output file')
+   of the file specified by --config.
-    parser.add_argument('eeprom', help='Name of EEPROM file to use as input')
+
+   rpi-eeprom-config --config boot.conf --out newimage.bin pieeprom.bin
+
+   The new image file can be installed via rpi-eeprom-update
+   rpi-eeprom-update -d -f newimage.bin
+
+4. Applies a given config file to the latest available EEPROM image and
+   invokes rpi-eeprom-update to update the bootloader when the system is
+   rebooted.
+
+   rpi-eeprom-config --apply boot.conf
+
+   The latest available image is determined by querying \'rpi-eeprom-update -l\'
+   and depends on the rpi-eeprom-update configuration.
+
+Bootloader EEPROM images are contained in the \'rpi-eeprom-images\' package,
+which installs them to the /lib/firmware/raspberrypi/bootloader directory.'
+"""
+    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
+                                     description=description)
+
+    parser.add_argument('-a', '--apply', required=False,
+                        help='Updates the bootloader to the given config plus latest available EEPROM release.')
+    parser.add_argument('-c', '--config', help='Name of bootloader configuration file', required=False)
+    parser.add_argument('-o', '--out', help='Name of output file', required=False)
+    parser.add_argument('eeprom', nargs='?', help='Name of EEPROM file to use as input')
     args = parser.parse_args()
 
+    if args.apply is not None:
+        if not os.path.exists(args.apply):
+            exit_error("config file '%s' not found" % args.apply)
+        apply_update(args.apply)
+    elif args.eeprom is not None:
         image = BootloaderImage(args.eeprom, args.out)
         if args.config is not None:
+            if not os.path.exists(args.config):
+                exit_error("config file '%s' not found" % args.config)
             image.write(args.config)
         else:
             image.read()
+    elif args.config is None and args.eeprom is None:
+        sys.stdout.write(read_current_config())
 
 if __name__ == '__main__':
+    atexit.register(exit_handler)
     main()