path: root/scripts/tizen/sd_fusing.py

#!/usr/bin/env python3

from functools import reduce

import argparse
import atexit
import errno
import logging
import os
import re
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile

__version__ = "1.1.1"

Format = False
Device = ""
File = ""
Yes = False
SuperDelivered = False

LOGGING_NOTICE = int((logging.INFO + logging.WARNING) / 2)

class DebugFormatter(logging.Formatter):
    def format(self, record):
        if record.levelno == logging.DEBUG:
            record.debuginfo = "[{}:{}] ".format(os.path.basename(record.pathname), record.lineno)
        else:
            record.debuginfo = ''
        return logging.Formatter.format(self, record)

class ColorFormatter(DebugFormatter):
    _levelToColor = {
        logging.CRITICAL: "\x1b[35;1m",
        logging.ERROR: "\x1b[33;1m",
        logging.WARNING: "\x1b[33;1m",
        LOGGING_NOTICE: "\x1b[0m",
        logging.INFO: "\x1b[0m",
        logging.DEBUG: "\x1b[30;1m",
        logging.NOTSET: "\x1b[30;1m"
    }
    def format(self, record):
        record.levelcolor = self._levelToColor[record.levelno]
        record.msg = record.msg
        return super().format(record)

class ColorStreamHandler(logging.StreamHandler):
    def __init__(self, stream=None, format=None, datefmt=None, style='%', cformat=None):
        logging.StreamHandler.__init__(self, stream)
        if os.isatty(self.stream.fileno()):
            self.formatter = ColorFormatter(cformat, datefmt, style)
            self.terminator = "\x1b[0m\n"
        else:
            self.formatter = DebugFormatter(format, datefmt, style)

class Partition:
    def __init__(self, name, size, start=None, ptype=None, fstype="raw", bootable=False, **kwargs):
        self.name = name
        self.size = size
        self.size_sectors = kwargs.get("size_sectors", None)
        self.start = start
        self.start_sector = kwargs.get("start_sector", None)
        self.ptype = ptype
        self.bootable = bootable
        if type(self.size_sectors) == int and self.size_sectors >= 0:
            if type(self.size) == int and self.size >= 0:
                logging.warning(f"partition:{name} overriding size to the value obtained from size_sectors")
            # size is used to calculate free space, so adjust it here
            self.size = (self.size_sectors * 512 - 1) / (1024*1024) + 1
        if type(self.start_sector) == int and self.start_sector >= 0:
            if type(self.start) == int and self.start >= 0:
                logging.warning(f"partition:{name} overriding start to the value obtained from start_sector")
            self.size = None

    def __str__(self):
        output = []
        if self.start_sector:
            output.append(f"start={self.start_sector}")
        elif self.start:
            output.append(f"start={self.start}MiB")
        if type(self.size_sectors) == int and self.size_sectors >= 0:
            output.append(f"size={self.size_sectors}")
        elif type(self.size) == int and self.size >= 0:
            output.append(f"size={self.size}MiB")
        if self.name:
            output.append(f"name={self.name}")
        output.append(f"type={self.ptype}")
        if self.bootable:
                       output.append("bootable")
        return ", ".join(output) + "\n"

class Label:
    def __init__(self, part_table, ltype):
        self.ltype = ltype
        if ltype == 'gpt':
            ptype = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
        elif ltype == 'dos':
            ptype = '83'
        self.part_table = []
        for part in part_table:
            part["ptype"] = part.get("ptype", ptype)
            self.part_table.append(Partition(**part))
    def __str__(self):
        output = f"label: {self.ltype}\n"
        if self.ltype == 'gpt':
            output += f"first-lba: 34\n"
        for part in self.part_table:
            output += str(part)
        return output

class SdFusingTarget:
    params = (('reboot-param.bin', 'norm'),
              ('reboot-param.info', 'norm'),
              ('upgrade-status.info', '0'))

    def __init__(self, device, ltype):
        # TODO: make a copy of a sublcass part_table
        self.with_super = False
        self.device = device
        total_size = device_size(device)

        if hasattr(self, 'user_partition'):
            self.user_size = total_size - self.reserved_space - \
                reduce(lambda x, y: x + (y["size"] or 0), self.part_table, 0)
            if self.user_size < 100:
                logging.error(f"Not enough space for user data ({self.user_size}). Use larger storage.")
                raise OSError(errno.ENOSPC, os.strerror(errno.ENOSPC), device)
            # self.user_partition counts from 0
            self.part_table[self.user_partition]["size"] = self.user_size

        self.label = Label(self.part_table, ltype)
        if not hasattr(self, 'bootcode'):
            self.bootcode = None
        self.binaries = self._get_binaries('binaries')

    def apply_partition_sizes(self, partition_sizes):
        if partition_sizes is None or len(partition_sizes) == 0:
            return 0
        resized_total = 0
        for name, size in partition_sizes.items():
            resized_count = 0
            for part in self.part_table:
                if part['name'] == name:
                    psize = part['size']
                    part['size'] = size
                    logging.debug(f"overriding partition:{name}, old-size:{psize} MiB new-size:{size} MiB")
                    resized_count = resized_count + 1
            if resized_count == 0:
                logging.error(f"partition:{name} not found when attempting to apply_partition_sizes")
            resized_total = resized_total + resized_count
        return resized_total

    def _get_binaries(self, key):
        binaries = {}
        for i, p in enumerate(self.part_table):
            b = p.get(key, None)
            if b is None:
                continue
            if isinstance(b, str):
                binaries[b] = i + 1
            elif isinstance(b, list):
                for f in b:
                    binaries[f] = i + 1
        return binaries

    def get_partition_index_list(self, binary):
        if hasattr(self, 'update'):
            logging.error("You have requested to update the {} partition set. "
                          "This target does not support A/B partition sets."
                          .format(self.update.upper()))
            sys.exit(1)
        return [self.binaries.get(binary, None)]

    def ensure_parttable(self):
        logging.notice(f"Verifying that partition table on {Device} matches target specification")
        for partnum, part in enumerate(self.part_table, 1):
            bo = subprocess.check_output(["blkid", "-o", "export", Device + str(partnum)]).decode('utf-8')
            if "PARTLABEL=" in bo and f"PARTLABEL={part['name']}" not in bo:
                logging.error(f'On-device partition label mismatch with selected target: partlabel={part["name"]}, on-device:\n{bo}')
                sys.exit(1)

    def initialize_parameters(self):
        pass

    def write_parameters(self, params = None):
        pass

    def update_parameters(self):
        self.write_parameters()

class SdFusingTargetAB(SdFusingTarget):
    def __init__(self, device, ltype):
        super().__init__(device, ltype)
        self.binaries_b = self._get_binaries('binaries_b')

    def get_partition_index_list(self, binary):
        if self.update == 'b':
            return [self.binaries_b.get(binary, None)]
        elif self.update == 'ab':
            return [self.binaries.get(binary, None), self.binaries_b.get(binary, None)]

        return [self.binaries.get(binary, None)]

    def update_parameters(self):
        part_ab = 'a' if self.update in [None, '', 'a', 'ab'] else 'b'
        part_cloned = '1' if self.update == 'ab' else '0'
        params = [('partition-ab.info', part_ab),
                  ('partition-ab-cloned.info', part_cloned)]
        if not self.update in [None, '', 'a', 'ab']:
            params.append(('partition-a-status.info', 'ok'))
        if self.update in ['b', 'ab']:
            params.append(('partition-b-status.info', 'ok'))
        self.write_parameters(self.params + tuple(params))

class InitParams:
    def find_inform(self):
        n = None
        for i, p in enumerate(self.part_table):
            if p['name'] == 'inform':
                n = i + 1;
                break
        d = "/dev/" + get_partition_device(self.device, n)
        return d

    def initialize_parameters(self):
        logging.debug("Initializing parameters")
        d = self.find_inform()

        argv = ['tune2fs', '-O', '^metadata_csum', d]
        logging.debug(" ".join(argv))
        subprocess.run(argv,
                       stdin=subprocess.DEVNULL,
                       stdout=None, stderr=None)

    def write_parameters(self, params = None):
        d = self.find_inform()
        logging.debug(f"Writing parameters to {d}")
        with tempfile.TemporaryDirectory() as mnt:
            argv = ['mount', '-t', 'ext4', d, mnt]
            logging.debug(" ".join(argv))
            proc = subprocess.run(argv,
                                  stdin=subprocess.DEVNULL,
                                  stdout=None, stderr=None)
            if proc.returncode != 0:
                logging.error(f"Failed to mount {d} in {mnt} (Has the device been initialized with --format?)")
                return
            parameters = self.params if params is None else params
            for param, value in parameters:
                with open(os.path.join(mnt, param), 'w') as f:
                    logging.debug(f"Writing parameter {param}={value}")
                    f.write(value + '\n')
            argv = ['umount', d]
            logging.debug(" ".join(argv))
            subprocess.run(argv,
                           stdin=subprocess.DEVNULL,
                           stdout=None, stderr=None)

class Rpi3(InitParams, SdFusingTarget):
    long_name = "Raspberry Pi 3"
    part_table = [
        {"size": 64,   "name": "boot", "start": 4, "ptype": "0xe", "bootable": True,
         "binaries":   "boot.img"},
        {"size": 3072, "name": "rootfs",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "ptype":  "5",    "name": "extended", "start": 4484},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "modules",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal",
         "binaries":   "hal.img"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.reserved_space = 12
        self.user_partition = 4
        super().__init__(device, "dos")

class Rpi4Super(InitParams, SdFusingTargetAB):
    long_name = "Raspberry Pi 4 w/ super partition"
    part_table = [
        {"size": 64,   "name": "boot_a","start": 4,
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 6657, "name": "super",
         "binaries":   "super.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": 36,   "fstype": "raw",  "name": "none"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 64,   "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"}
    ]

    def __init__(self, device, args):
        self.reserved_space = 8
        self.user_partition = 4
        self.update = args.update
        super().__init__(device, "gpt")
        self.with_super = True
        self.super_alignment = 1048576

class Rpi4(InitParams, SdFusingTargetAB):
    long_name = "Raspberry Pi 4"
    part_table = [
        {"size": 64,   "name": "boot_a", "start": 4,
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 3072, "name": "rootfs_a",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": 36,   "name": "none"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal_a",
         "binaries":   "hal.img"},
        {"size": 64,   "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 3072, "name": "rootfs_b",
         "binaries_b": "rootfs.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 256,  "name": "hal_b",
         "binaries_b": "hal.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.reserved_space = 5
        self.user_partition = 4
        self.update = args.update
        super().__init__(device, "gpt")

class Rpi4AoT(InitParams, SdFusingTargetAB):
    long_name = "Raspberry Pi 4 for AoT"
    part_table = [
        {"size": 64,   "name": "boot_a", "start": 4,
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 3072, "name": "rootfs_a",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": 36,   "name": "none"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal_a",
         "binaries":   "hal.img"},
        {"size": 64,   "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 3072, "name": "rootfs_b",
         "binaries_b": "rootfs.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 256,  "name": "hal_b",
         "binaries_b": "hal.img"},
        {"size": 1536, "name": "aot-system_a",
         "binaries":   "system.img"},
        {"size": 1536, "name": "aot-system_b",
         "binaries_b": "system.img"},
        {"size": 256,  "name": "aot-vendor_a",
         "binaries":   "vendor.img"},
        {"size": 256,  "name": "aot-vendor_b",
         "binaries_b": "vendor.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.reserved_space = 5
        self.user_partition = 4
        self.update = args.update
        super().__init__(device, "gpt")

class RV64(InitParams, SdFusingTarget):
    long_name = "QEMU RISC-V 64-bit"
    part_table = [
        {"size": 2,    "name": "SPL", "start": 4,
         "ptype":      "2E54B353-1271-4842-806F-E436D6AF6985",
         "binaries":   ""},
        {"size": 4,    "name": "u-boot",
         "ptype":      "5B193300-FC78-40CD-8002-E86C45580B47",
         "binaries":  ["u-boot.img", "u-boot.itb"],},
        {"size": 292,  "name": "boot_a",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 36,   "name": "none"},
        {"size": 3072, "name": "rootfs_a",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal_a",
         "binaries":   "hal.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.user_partition = 6
        self.reserved_space = 5
        self.apply_partition_sizes(args.partition_sizes)
        super().__init__(device, 'gpt')

class VF2(InitParams, SdFusingTargetAB):
    long_name = "VisionFive2"
    part_table = [
        {"size": 2,    "name": "SPL", "start": 4,
         "ptype":      "2E54B353-1271-4842-806F-E436D6AF6985",
         "binaries":  ["u-boot-spl.bin.normal.out"],},
        {"size": 4,    "name": "u-boot",
         "ptype":      "5B193300-FC78-40CD-8002-E86C45580B47",
         "binaries":  ["u-boot.img", "u-boot.itb"],},
        {"size": 128,  "name": "boot_a",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 36,   "name": "none"},
        {"size": 3072, "name": "rootfs_a",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal_a",
         "binaries":   "hal.img"},
        {"size": 128,  "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 3072, "name": "rootfs_b",
         "binaries_b": "rootfs.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 256,  "name": "hal_b",
         "binaries_b": "hal.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.user_partition = 6
        self.reserved_space = 5
        self.update = args.update
        self.apply_partition_sizes(args.partition_sizes)
        super().__init__(device, 'gpt')

class VF2Super(InitParams, SdFusingTargetAB):
    long_name = "VisionFive2 w/ super partition"
    part_table = [
        {"size": 2,    "name": "SPL", "start": 4,
         "ptype":      "2E54B353-1271-4842-806F-E436D6AF6985",
         "binaries":  ["u-boot-spl.bin.normal.out"],},
        {"size": 4,    "name": "u-boot",
         "ptype":      "5B193300-FC78-40CD-8002-E86C45580B47",
         "binaries":  ["u-boot.img", "u-boot.itb"],},
        {"size": 128,  "name": "boot_a",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 36,   "fstype": "raw",  "name": "none"},
        {"size": 6656, "name": "super",
         "binaries":   "super.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 128,  "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    def __init__(self, device, args):
        self.user_partition = 6
        self.reserved_space = 5
        self.update = args.update
        super().__init__(device, 'gpt')
        self.with_super = True
        self.super_alignment = 1048576

class LicheePi4A(InitParams, SdFusingTargetAB):
    long_name = "LicheePi4A"
    part_table = [
        {"size": None, "name": "spl+uboot",
         "start_sector": 34, "size_sectors": 4062,
         "ptype":      "8DA63339-0007-60C0-C436-083AC8230908",
         "binaries":  ["u-boot-with-spl.bin"],},
        {"size": 128,  "name": "boot_a",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 3072, "name": "rootfs_a",
         "binaries":   "rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 256,  "name": "hal_a",
         "binaries":   "hal.img"},
        {"size": 128,  "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 3072, "name": "rootfs_b",
         "binaries_b": "rootfs.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 256,  "name": "hal_b",
         "binaries_b": "hal.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"},
    ]

    # bootcode written to the protective MBR, aka RV64 'J 0x4400' (sector 34)
    bootcode = b'\x6f\x40\x00\x40'

    def __init__(self, device, args):
        self.user_partition = 4
        self.reserved_space = 5
        self.update = args.update
        self.apply_partition_sizes(args.partition_sizes)
        super().__init__(device, 'gpt')

class LicheePi4ASuper(InitParams, SdFusingTargetAB):
    long_name = "LicheePi4A w/ super partition"
    part_table = [
        {"size": None, "name": "spl+uboot",
         "start_sector": 34, "size_sectors": 4062,
         "ptype":      "8DA63339-0007-60C0-C436-083AC8230908",
         "binaries":  ["u-boot-with-spl.bin"],},
        {"size": 128,  "name": "boot_a",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "boot.img"},
        {"size": 6656, "name": "super",
         "binaries":   "super.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "system-data.img"},
        {"size": None, "name": "user",
         "binaries":   "user.img"},
        {"size": 32,   "name": "module_a",
         "binaries":   "modules.img"},
        {"size": 32,   "name": "ramdisk_a",
         "binaries":   "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_a",
         "binaries":   "ramdisk-recovery.img"},
        {"size": 8,    "name": "inform", "fstype": "ext4"},
        {"size": 128,  "name": "boot_b",
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries_b": "boot.img"},
        {"size": 32,   "name": "module_b",
         "binaries_b": "modules.img"},
        {"size": 32,   "name": "ramdisk_b",
         "binaries_b": "ramdisk.img"},
        {"size": 32,   "name": "ramdisk-recovery_b",
         "binaries_b": "ramdisk-recovery.img"},
        {"size": 4,    "name": "reserved0"},
        {"size": 64,   "name": "reserved1"},
        {"size": 125,  "name": "reserved2"}
    ]

    # bootcode written to the protective MBR, aka RV64 'J 0x4400' (sector 34)
    bootcode = b'\x6f\x40\x00\x40'

    def __init__(self, device, args):
        self.reserved_space = 8
        self.user_partition = 4
        self.update = args.update
        super().__init__(device, 'gpt')
        self.with_super = True
        self.super_alignment = 1048576


class X86emu(SdFusingTarget):
    part_table = [
        {"size": 512,  "fstype": "vfat",  "name": "EFI", "start": 4,
         "ptype":      "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
         "binaries":   "",},
        {"size": 512,  "name": "boot",
         "binaries":   "emulator-boot.img",},
        {"size": 2048, "fstype": "ext4", "name": "rootfs",
         "binaries":   "emulator-rootfs.img"},
        {"size": 1344, "name": "system-data",
         "binaries":   "emulator-sysdata.img"},
        {"size": 1024, "name": "emulator-swap",
         "ptype":      "0657FD6D-A4AB-43C4-84E5-0933C84B4F4F"},
    ]

    def __init__(self, device, args):
        super().__init__(device, 'gpt')
        for p in self.label.part_table:
            if p.name == "rootfs":
                p.ptype = args._rootfs_uuid
                break

class X86emu32(X86emu):
    long_name = "QEMU x86 32-bit"

    def __init__(self, device, args):
        setattr(args, "_rootfs_uuid", "44479540-F297-41B2-9AF7-D131D5F0458A")
        super().__init__(device, args)

class X86emu64(X86emu):
    long_name = "QEMU x86 64-bit"

    def __init__(self, device, args):
        setattr(args, "_rootfs_uuid", "44479540-F297-41B2-9AF7-D131D5F0458A")
        super().__init__(device, args)

TARGETS = {
    'rpi3': Rpi3,
    'rpi4': Rpi4,
    'rpi4s': Rpi4Super,
    'rpi4aot': Rpi4AoT,
    'vf2': VF2,
    'vf2s': VF2Super,
    'rv64': RV64,
    'lpi4a': LicheePi4A,
    'lpi4as': LicheePi4ASuper,
    'x86emu32': X86emu32,
    'x86emu64': X86emu64,
}

def device_size(device):
    argv = ["sfdisk", "-s", device]
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdout=subprocess.PIPE)
    size = int(proc.stdout.decode('utf-8').strip()) >> 10
    logging.debug(f"{device} size {size}MiB")
    return size

def check_sfdisk():
    min_major = 2
    min_minor = 26
    min_minor_del = 28

    proc = subprocess.run(['sfdisk', '-v'],
                          stdout=subprocess.PIPE)
    version = proc.stdout.decode('utf-8').strip()
    logging.debug(f"Found {version}")
    version_tokens = [int(x) for x in re.findall('[0-9]+', version)]

    if len(version_tokens) == 3:
        major, minor, patch = version_tokens[0:3]
        version_str = f"{major}.{minor}.{patch}"
    elif len(version_tokens) == 2:
        major, minor = version_tokens[0:2]
        version_str = f"{major}.{minor}"
    else:
        logging.warning("Did not read version of sfdisk correctly.")
        return False,False

    support_delete = False

    if major < min_major or major == min_major and minor < min_minor:
        logging.error(f"Your sfdisk {version_str} is too old. Please switch to at least {min_major}.{min_minor}")
        return False,False
    elif major == min_major and minor >= min_minor_del:
        support_delete = True

    return True, support_delete

def mkpart(args, target):
    global Device
    new, support_delete = check_sfdisk()

    if not new:
        logging.error('sfdisk too old')
        sys.exit(1)

    with open('/proc/self/mounts') as mounts:
        device_kname = '/dev/' + get_device_kname(Device)
        device_re = re.compile(device_kname + '[^ ]*')
        logging.debug(f"Checking for mounted partitions on {device_kname}")
        for m in mounts:
            match = device_re.match(m)
            if match:
                logging.warning('Found mounted device: ' + match[0])
                argv = ['umount', match[0]]
                logging.debug(" ".join(argv))
                proc = subprocess.run(argv)
                if proc.returncode != 0:
                    logging.error(f"Failed to unmount {match[0]}")
                    sys.exit(1)

    if support_delete:
        logging.info("Removing old partitions")
        argv = ['sfdisk', '--delete', Device]
        logging.debug(" ".join(argv))
        proc = subprocess.run(argv)
        if proc.returncode != 0:
            logging.error(f"Failed to remove the old partitions from {Device}")
    else:
        logging.info("Removing old partition table")
        argv = ['dd', 'if=/dev/zero', 'of=' + Device,
                'bs=512', 'count=32', 'conv=notrunc']
        logging.debug(" ".join(argv))
        proc = subprocess.run(argv)
        if proc.returncode != 0:
            logging.error(f"Failed to clear the old partition table on {Device}")
            sys.exit(1)

    logging.debug("New partition table:\n" + str(target.label))
    argv = ['sfdisk', '--wipe-partitions', 'always', Device]
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdout=None,
                          stderr=None,
                          input=str(target.label).encode())
    if proc.returncode != 0:
        logging.error(f"Failed to create partition a new table on {Device}")
        logging.error(f"New partition table:\n" + str(target.label))
        sys.exit(1)

    # Run `udevadm settle` to ensure that partition change made by `sfdisk` is completely reflected in userspace.
    logging.info("Waiting for the udev event queue to empty...")
    argv = ['udevadm', 'settle']
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdout=None,
                          stderr=None)
    if proc.returncode != 0:
        logging.warning("udevadm settle exited without clearing the udev event queue.")
    else:
        logging.info("The udev event queue is empty.")

    if target.bootcode:
        logging.debug("Writing bootcode\n")
        with open(Device, "wb") as f:
            f.write(target.bootcode)
            f.close

    for i, part in enumerate(target.part_table):
        d = "/dev/" + get_partition_device(target.device, i+1)
        if not 'fstype' in part:
            logging.debug(f"Filesystem not defined for {d}, skipping")
            continue
        logging.debug(f"Formatting {d} as {part['fstype']}")
        if part['fstype'] == 'vfat':
            argv = ['mkfs.vfat', '-F', '16', '-n', part['name'], d]
            logging.debug(" ".join(argv))
            proc = subprocess.run(argv,
                                  stdin=subprocess.DEVNULL,
                                  stdout=None, stderr=None)
            if proc.returncode != 0:
                logging.error(f"Failed to create FAT filesystem on {d}")
                sys.exit(1)
        elif part['fstype'] == 'ext4':
            argv = ['mkfs.ext4', '-q', '-L', part['name'], d]
            logging.debug(" ".join(argv))
            proc = subprocess.run(argv,
                                  stdin=subprocess.DEVNULL,
                                  stdout=None, stderr=None)
            if proc.returncode != 0:
                logging.error(f"Failed to create ext4 filesystem on {d}")
                sys.exit(1)
        elif part['fstype'] == 'swap':
            argv = ['mkswap', '-L', part['name'], d]
            logging.debug(" ".join(argv))
            proc = subprocess.run(argv,
                                  stdin=subprocess.DEVNULL,
                                  stdout=None, stderr=None)
            if proc.returncode != 0:
                logging.error(f"Failed to format swap partition {d}")
                sys.exit(1)
        elif part['fstype'] == 'raw':
            pass
    target.initialize_parameters()

def check_args(args):
    global Format
    global Yes
    global SuperDelivered

    logging.info(f"Device: {args.device}")

    if args.binaries and len(args.binaries) > 0:
        logging.info("Fusing binar{}: {}".format("y" if len(args.binaries) == 1 else "ies",
                     ", ".join(args.binaries)))

    if args.YES:
        Yes = True

    if args.create:
        Format = True
        Yes = True

    if args.format:
        if Yes:
            Format = True
        else:
            response = input(f"{args.device} will be formatted. Continue? [y/N] ")
            if response.lower() in ('y', 'yes'):
                Format = True
            else:
                Format = False

    if args.super_delivered:
        SuperDelivered = True

def check_device(args):
    global Format
    global Device
    Device = args.device

    if args.create:
        if os.path.exists(Device):
            logging.error(f"Failed to create '{Device}', the file alread exists")
            sys.exit(1)
        else:
            argv = ["dd", "if=/dev/zero", f"of={Device}",
                    "conv=sparse", "bs=1M", f"count={args.size}"]
            logging.debug(" ".join(argv))
            rc = subprocess.run(argv)
            if rc.returncode != 0:
                logging.error("Failed to create the backing file")
                sys.exit(1)

    if os.path.isfile(Device):
        global File
        File = Device

        argv = ["losetup", "--show", "--partscan", "--find", f"{File}"]
        logging.debug(" ".join(argv))
        proc = subprocess.run(argv,
                              stdout=subprocess.PIPE)
        Device = proc.stdout.decode('utf-8').strip()
        if proc.returncode != 0:
            logging.error(f"Failed to attach {File} to a loopback device")
            sys.exit(1)
        logging.debug(f"Loop device found: {Device}")
        atexit.register(lambda: subprocess.run(["losetup", "-d", Device]))

    try:
        s = os.stat(Device)
        if not stat.S_ISBLK(s.st_mode):
            raise TypeError
    except FileNotFoundError:
        logging.error(f"No such device: {Device}")
        sys.exit(1)
    except TypeError:
        logging.error(f"{Device} is not a block device")
        sys.exit(1)

def check_partition_format(args, target):
    global Format
    global Device

    if not Format:
        logging.info(f"Skip formatting of {Device}".format(Device))
        target.ensure_parttable()
        return
    logging.info(f"Start formatting of {Device}")
    mkpart(args, target)
    logging.info(f"{Device} formatted")

def check_ddversion():
    proc = subprocess.run(["dd", "--version"],
                            stdout=subprocess.PIPE)
    version = proc.stdout.decode('utf-8').split('\n')[0].strip()
    logging.debug(f"Found {version}")
    major, minor = (int(x) for x in re.findall('[0-9]+', version))

    if major < 8 or major == 8 and minor < 24:
        return False

    return True

def get_partition_device(device, idx):
    argv = ['lsblk', device, '-o', 'TYPE,KNAME']
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdout=subprocess.PIPE)
    if proc.returncode != 0:
        logging.error("lsblk has failed")
        sys.exit(1)
    part_re = re.compile(f"^part\s+(.*[^0-9]{idx})$")
    for l in proc.stdout.decode('utf-8').splitlines():
        match = part_re.match(l)
        if match:
            return match[1]
    logging.error("device entry not found")
    sys.exit(1)

def get_device_kname(device):
    argv = ['lsblk', device, '-o', 'TYPE,KNAME']
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdout=subprocess.PIPE)
    if proc.returncode != 0:
        logging.error("lsblk has failed")
        sys.exit(1)
    for l in proc.stdout.decode('utf-8').splitlines():
        match = re.search(f"^(disk|loop)\s+(.*)", l)
        if match:
            return match[2]
    logging.error("kname entry not found")
    sys.exit(1)

def do_fuse_file(f, name, target):
    indexes = target.get_partition_index_list(name)
    if len(indexes) == 0:
        logging.info(f"No partition defined for {name}, skipping.")
        return
    for idx in indexes:
        if idx is None:
            logging.info(f"No partition defined for {name}, skipping.")
            continue
        pdevice = "/dev/" + get_partition_device(Device, idx)
        argv = ['dd', 'bs=4M',
                'oflag=direct',
                'iflag=fullblock',
                'conv=nocreat',
                'status=progress',
                f"of={pdevice}"]
        logging.debug(" ".join(argv))
        proc_dd = subprocess.Popen(argv,
                                   bufsize=(4 << 20),
                                   stdin=subprocess.PIPE,
                                   stdout=None, stderr=None)
        logging.notice(f"Writing {name} to {pdevice}")
        buf = f.read(4 << 20)
        while len(buf) > 0:
            proc_dd.stdin.write(buf)
            buf = f.read(4 << 20)
        proc_dd.communicate()
        logging.info("Done")
        #TODO: verification

#TODO: functions with the target argument should probably
#      be part of some class

def do_fuse_image_super(tmpd, target):
    metadata_slots = 2
    metadata_size = 65536

    hal_path = os.path.join(tmpd, 'hal.img')
    rootfs_path = os.path.join(tmpd, 'rootfs.img')
    super_path = os.path.join(tmpd, 'super.img')

    try:
        hal_size = os.stat(hal_path).st_size
        rootfs_size = os.stat(rootfs_path).st_size
    except FileNotFoundError as e:
        fn = os.path.split(e.filename)[-1]
        logging.warning(f"{fn} is missing, skipping super partition image")
        return

    group_size = 2 * (hal_size + rootfs_size)
    super_size = 2 * group_size

    # calculate additional space needed for metadata.
    # There are 2 metadata slots having 65536 B each. 131072 B in total
    additional_space = (metadata_slots * metadata_size) / target.super_alignment

    if additional_space > 1:
        # if metadata takes more than super alignment, add 1 MiB to super
        super_size += 1024*1024
    else:
        # if metadata takes less than super alignment, add alignment size to super
        super_size += target.super_alignment

    argv = ["lpmake", "-F",
            f"-o={super_path}",
            f"--device-size={super_size}",
            f"--metadata-size={metadata_size}",
            f"--metadata-slots={metadata_slots}",
            "-g", f"tizen_a:{group_size}",
            "-p", f"rootfs_a:none:{rootfs_size}:tizen_a",
            "-p", f"hal_a:none:{hal_size}:tizen_a",
            "-g", f"tizen_b:{group_size}",
            "-p", f"rootfs_b:none:{rootfs_size}:tizen_b",
            "-p", f"hal_b:none:{hal_size}:tizen_b",
            "-i", f"rootfs_a={rootfs_path}",
            "-i", f"rootfs_b={rootfs_path}",
            "-i", f"hal_a={hal_path}",
            "-i", f"hal_b={hal_path}"]
    logging.debug(" ".join(argv))
    proc = subprocess.run(argv,
                          stdin=subprocess.DEVNULL,
                          stdout=None, stderr=None)

    if proc.returncode != 0:
        logging.error("Failed to create super.img")
    do_fuse_image(super_path, target)

def do_fuse_image_tarball(tarball, tmpd, target):
    with tarfile.open(tarball) as tf:
        for entry in tf:
            if target.with_super and not SuperDelivered:
                if entry.name in('hal.img', 'rootfs.img'):
                    tf.extract(entry, path=tmpd)
                    continue
            f = tf.extractfile(entry)
            do_fuse_file(f, entry.name, target)

def do_fuse_image(img, target):
    with open(img, 'rb') as f:
        do_fuse_file(f, os.path.basename(img), target)

def fuse_image(args, target):
    global Yes

    if args.binaries is None or len(args.binaries) == 0:
        return

    if not Yes and not Format:
        print(f"The following images will be written to {args.device} and the "
              "existing data will be lost.\n")
        for b in args.binaries:
            print("  " + b)
        response = input("\nContinue? [y/N] ")
        if not response.lower() in ('y', 'yes'):
            return

    with tempfile.TemporaryDirectory() as tmpd:
        for b in args.binaries:
            if re.search('\.(tar|tar\.gz|tgz)$', b):
                do_fuse_image_tarball(b, tmpd, target)
            else:
                fn = os.path.split(b)[-1]
                if target.with_super and fn in ('rootfs.img', 'hal.img') and not SuperDelivered:
                    shutil.copy(b, os.path.join(tmpd, fn))
                else:
                    do_fuse_image(b, target)

        if target.with_super and not SuperDelivered:
            do_fuse_image_super(tmpd, target)
    target.update_parameters()

def logger_notice(self, msg, *args, **kws):
    if self.isEnabledFor(LOGGING_NOTICE):
        self._log(LOGGING_NOTICE, msg, args, **kws)
logging.Logger.notice = logger_notice

def logging_notice(msg, *args, **kws):
    if len(logging.root.handlers) == 0:
        basicConfig()
    logging.root.notice(msg, *args, **kws)
logging.notice = logging_notice


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="For {}, version {}".format(
        ", ".join([v.long_name for k,v in TARGETS.items()]),
        __version__
    ))
    parser.add_argument("-b", "--binary", action="extend", dest="binaries",
                        nargs='+',
                        help="binary to flash, may be used multiple times")
    parser.add_argument("--create", action="store_true",
                        help="create the backing file and format the loopback device")
    parser.add_argument("--debug", action="store_const", const="debug",
                        default="notice", dest="log_level",
                        help="set log level to DEBUG")
    parser.add_argument("-d", "--device",
                        help="device node or loopback backing file")
    parser.add_argument("--format", action="store_true",
                        help="create new partition table on the target device")
    parser.add_argument("--log-level", dest="log_level", default="notice",
                        help="Verbosity, possible values: debug, info, notice, warning, "
                        "error, critical (default: notice)")
    parser.add_argument("--partition-size", type=str, action="extend", dest="partition_sizes",
                        nargs='*',
                        help="override default partition size (in MiB) (used with --format), "
                        "may be used multiple times, for example: --partition-size hal_a=256")
    parser.add_argument("--size", type=int, default=8192,
                        help="size of the backing file to create (in MiB)")
    parser.add_argument("-t", "--target", required=True,
                        help="Target device model. Use `--target list`"
                        " to show supported devices.")
    parser.add_argument("--update", choices=['a', 'b', 'ab'], default=None,
                        help="Choose partition set to update: a or b or ab.")
    parser.add_argument("--version", action="version",
                        version=f"%(prog)s {__version__}")
    parser.add_argument("--YES", action="store_true",
                        help="agree to destroy data on the DEVICE")
    parser.add_argument("--super_delivered", action="store_true",
                        help="indicate that super.img is already in tarball and doesn't have to be created during fusing")
    args = parser.parse_args()

    if args.target == 'list':
        print("\nSupported devices:\n")
        for k,v in TARGETS.items():
            print(f"  {k:6}  {v.long_name}")
        sys.exit(0)

    if args.device is None:
        parser.error('-d/--device argument is required for normal operation')

    if args.partition_sizes is not None:
        partition_sizes = {}
        for eqstr in args.partition_sizes:
            ptstr = eqstr.split('=')
            if len(ptstr) == 2:
                name = ptstr[0]
                size = int(ptstr[1])
                partition_sizes[name] = size
            else:
                parser.error('--partition-size must follow the name=size pattern')
        args.partition_sizes = partition_sizes

    logging.addLevelName(LOGGING_NOTICE, "NOTICE")
    conh = ColorStreamHandler(format='%(asctime)s.%(msecs)d %(debuginfo)s%(levelname)-8s %(message)s',
                              cformat='%(asctime)s.%(msecs)d %(debuginfo)s%(levelcolor)s%(message)s',
                              datefmt='%Y-%m-%dT%H:%M:%S')
    log_handlers = [conh]
    logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s',
                        handlers=log_handlers,
                        level=args.log_level.upper())

    logging.debug(" ".join(sys.argv))
    check_args(args)
    check_device(args)

    target = TARGETS[args.target](Device, args)

    check_partition_format(args, target)
    fuse_image(args, target)
    subprocess.run(['sync'],
                   stdin=subprocess.DEVNULL,
                   stdout=None, stderr=None )