text: #!/usr/bin/env python3
  # Copyright 2020 The Fuchsia Authors. All rights reserved.
  # Use of this source code is governed by a BSD-style license that can be
  # found in the LICENSE file.
  
  from __future__ import print_function
  
  import argparse
  import functools
  import json
  import logging
  import os
  import platform
  import stat
  import subprocess
  import sys
  import time
  import typing
  
  from typing import List
  
  import paths
  
  if sys.hexversion < 0x030700F0:
    logging.critical(
        'This script requires Python >= 3.7 to run (you have %s), please upgrade!'
        % (platform.python_version()),
        file=sys.stderr)
    sys.exit(1)
  
  WORKSTATION_INSTALLER_GPT_GUID = '4dce98ce-e77e-45c1-a863-caf92f1330c1'
  
  def make_unique_name(name, type):
    return f'{name}_{type}'
  
  class ManifestImage:
    """ Represents an entry in the 'images.json' manifest that will be installed to disk. """
    def __init__(self, name: str, guids: List[str], type: str):
      """
      Args:
        name: 'name' of the partition in the image manifest.
        guids: List of GPT GUIDs that this partition should be written to the disk with.
        type: 'type' of the partition in the image manifest.
      """
      self.name = name
      self.guids = guids
      self.type = type
  
    def unique_name(self):
      return make_unique_name(self.name, self.type)
  
  
  # This is the list of Fuchsia build images we write to the final image,
  # and the partition types they will have (passed to cgpt)
  IMAGES = [
      # Standard x64 partitions
      # This is the zedboot image, which is actually booted.
      ManifestImage('zedboot-efi', ['efi'], 'blk'),
      # This is the EFI system partition that will be installed to the target.
      ManifestImage('efi', [WORKSTATION_INSTALLER_GPT_GUID], 'blk'),
      ManifestImage('zircon-a', [WORKSTATION_INSTALLER_GPT_GUID], 'zbi'),
      ManifestImage('zircon-r', [WORKSTATION_INSTALLER_GPT_GUID], 'zbi'),
  
      # ChromeOS partitions
      # The zircon-r.signed partition is used as both zedboot on the installation
      # disk and also the installed zircon-r partition.
      ManifestImage('zircon-r.signed', ['kernel', WORKSTATION_INSTALLER_GPT_GUID], 'zbi.signed'),
      ManifestImage('zircon-a.signed', [WORKSTATION_INSTALLER_GPT_GUID], 'zbi.signed'),
  
      # Common partitions - installed everywhere.
      ManifestImage('storage-sparse', [WORKSTATION_INSTALLER_GPT_GUID], 'blk'),
  ]
  
  
  def ParseSize(size):
    """Parse a size.
  
    Args:
      size: '<number><suffix>', where suffix is 'K', 'M', or 'G'.
  
    Returns:
      A size in bytes equivalent to the human-readable size given.
    Raises:
      A ValueError if <suffix> is unrecognised or <number> is not a base-10
      number.
    """
    units = ['K', 'M', 'G']
    if size.isdigit():
      return int(size)
    else:
      unit = size[-1].upper()
      size_bytes = int(size[:-1]) * 1024
  
    if unit not in units:
      raise ValueError('unrecognised unit suffix "{}" for size {}'.format(
          unit, size))
  
    while units.pop(0) != unit:
      size_bytes *= 1024
    return size_bytes
  
  
  def PrettySize(size):
    """Returns a size in bytes as a human-readable string."""
    units = 'BKMGT'
  
    unit = 0
    # By the time we get to 3072, the error caused by
    # shifting units is <2%, so we don't care.
    while size > 3072 and unit < len(units) - 1:
      size /= 1024
      unit += 1
  
    return '{:1.1f}{}'.format(size, units[unit])
  
  
  class Partition:
    """Represents a single partition to be written to the disk.
  
    Attributes:
      label: label of the partition on the output image.
      path: path of the file that is going to be written to the partition on
        the host.
      real_size: size of the file on the host, in bytes.
      size: size of the partition, in bytes. This may not match real_size due
        to sector size alignment or EFI partition rules.
      type: type of the partition, passed to `cgpt`.
    """
    FAT32_MIN_SIZE = (63 * 1024 * 1024)
  
    def __init__(self, path, part_type, label, size: typing.Optional[int] = None):
      self.path = path
      self.type = part_type
      self.label = label
  
      # Calculate sector-aligned size of this file.
      if path:
        stat_result = os.stat(path)
        if not stat.S_ISREG(stat_result.st_mode):
          raise ValueError('{} is not a regular file.'.format(path))
        rounded_size = stat_result.st_size
        size = stat_result.st_size
      elif size:
        rounded_size = size
      else:
        raise ValueError('Expected a source file or size')
      if rounded_size % Image.SECTOR_SIZE != 0:
        rounded_size += Image.SECTOR_SIZE
        rounded_size -= rounded_size % Image.SECTOR_SIZE
      if self.type == 'efi':
        # Gigaboot won't be able to load zedboot.bin from an EFI partition that's
        # too small, so we ensure the partition is at least big enough for it to
        # work.
        rounded_size = max(Partition.FAT32_MIN_SIZE, rounded_size)
      self.real_size = size
      self.size = rounded_size
  
  
  class Image:
    """Represents a single disk image to be written.
  
      Attributes:
          filename: output filename of the image.
          is_usb: True if writing to a USB, False if creating an image on the
            host.
          file: filehandle to the output image. Held open while we work to prevent
            auto-mounting of USB.
          block_size: number of bytes to write at a time to the disk.
          file_size: total size of the image, in bytes
          partitions: list of |Partition| objects to write to disk.
    """
    CGPT_BIN = os.path.join(paths.PREBUILT_DIR, 'tools/cgpt',
                            paths.PREBUILT_PLATFORM, 'cgpt')
    SECTOR_SIZE = 512
    GPT_SECTORS = 2048
    CROS_RESERVED_SECTORS = 2048
  
    def __init__(self, filename, is_usb, block_size):
      self.filename = filename
      self.is_usb = is_usb
      self.file = open(filename, mode='wb')
      self.block_size = block_size
  
      # Allocate space for the primary and backup GPTs
      self.file_size = 2 * Image.GPT_SECTORS * Image.SECTOR_SIZE
      self.partitions = []
  
      # Set up the ChromeOS reserved partition.
      reserved_part = Partition(None, 'reserved', 'reserved',
                                self.CROS_RESERVED_SECTORS * self.SECTOR_SIZE)
      self.AddPartition(reserved_part)
  
    def _Cgpt(self, args):
      args = [Image.CGPT_BIN] + args
      return subprocess.run(args, capture_output=True)
  
    def _CgptAdd(self, part, offset):
      """Add a partition to the GPT represnted by thsis |Image|.
  
      Args:
        part: partition to add
        offset: offset to add partition at. Must be a multiple of SECTOR_SIZE.
      Returns:
        True if add succeded, False if it failed.
      """
      if offset % Image.SECTOR_SIZE != 0:
        raise ValueError('Offset must be a multiple of SECTOR_SIZE!')
      if part.size % Image.SECTOR_SIZE != 0:
        raise ValueError('Size must be a multiple of SECTOR_SIZE!')
      size = part.size // Image.SECTOR_SIZE
      offset //= Image.SECTOR_SIZE
      cgpt_args = ['add', '-s', str(size), '-t', part.type, '-b', str(offset),
                   '-l', part.label, self.filename]
  
      if part.type == 'kernel':
        # Mark CrOS kernel partition as bootable.
        # We assume that the disk will only have 1 kernel partition.
        # -T 1 = Bootloader will try to boot this partition once.
        # -S 0 = This partition has been successfully booted.
        # -P 2 = Partition priority is 2. Partition with the highest priority gets
        # booted.
        cgpt_args += ['-T', '1', '-S', '1', '-P', '2']
  
      ret = self._Cgpt(cgpt_args)
      if ret.returncode != 0:
        logging.critical('\n'
              '======= CGPT ADD FAILED! =======\n'
              'Maybe your disk is too small?\n')
        logging.error(ret.stdout)
        logging.error(ret.stderr)
        return False
      return True
  
    def AddPartition(self, partition):
      """Add a partition to the outputted disk image.
  
      This function does not write any data - call Finalise() to write the
      disk image once all partitions have been added.
  
      Args:
        partition: partition to add
      """
      self.partitions.append(partition)
      self.file_size += partition.size
  
    def WritePart(self, part, offset):
      """Writes data to a partition on the output device.
  
      Args:
        part: partiton to write
        offset: offset in bytes to write to
      """
      if part.path is None:
        return
      self.file.seek(offset, 0)
  
      written = 0
      logging.info(
          '   Writing image {} to partition {}... '.format(
              part.path.split('/')[-1], part.label)
          )
      with open(part.path, 'rb') as fh:
        start = time.perf_counter()
        for block in iter(functools.partial(fh.read, self.block_size), b''):
          written += len(block)
          self.file.write(block)
        # flush and fsync to get accurate timing results
        self.file.flush()
        os.fsync(self.file.fileno())
        finish = time.perf_counter()
      per_second = written / (finish - start)
      logging.info('     Wrote {} in {:1.2f}s, {}/s'.format(
          PrettySize(written), finish - start, PrettySize(per_second)))
  
    def Finalise(self):
      """Write all the partitions this image represents to disk/file."""
      if not self.is_usb:
        # first, make sure the file is big enough.
        logging.info('Create image of size={} bytes'.format(self.file_size))
        self.file.truncate(self.file_size)
        self.file.flush()
        os.fsync(self.file.fileno())
  
      logging.info('Creating new GPT partition table...')
      self._Cgpt(['create', self.filename])
      self._Cgpt(['boot', '-p', self.filename])
      logging.info('Done.')
  
      logging.info('Creating and writing partitions...')
      current_offset = Image.SECTOR_SIZE * Image.GPT_SECTORS
      for part in self.partitions:
        if not self._CgptAdd(part, current_offset):
          logging.warning('Write failed, aborting.')
          self.file.close()
          return
        self.WritePart(part, current_offset)
        current_offset += part.size
      logging.info('Done.')
      self.file.close()
  
  
  def GetPartitions():
    """Get all partitions to be written to the output image.
  
    The list of partitions is currently determined by the IMAGES dict
    at the top of this file.
  
    Returns:
      a list of |Partition| objects to be written to the disk.
    """
    images = {}
    try:
      with open(paths.FUCHSIA_BUILD_DIR + '/images.json') as f:
        images_list = json.load(f)
        for image in images_list:
          images[make_unique_name(image['name'], image['type'])] = image
    except IOError as err:
      logging.critical('Failed to find image manifest. Have you run `fx build`?', exc_info=err)
      return []
  
    ret = []
    is_bootable = False
    for image in IMAGES:
      if image.unique_name() not in images:
        logging.debug("Skipping image that wasn't built: {}".format(image.unique_name()))
        continue
  
      for part_type in image.guids:
        full_path = os.path.join(paths.FUCHSIA_BUILD_DIR, images[image.unique_name()]['path'])
        ret.append(Partition(full_path, part_type, image.name))
        # Assume that any non-installer partition is a bootable partition.
        if part_type != WORKSTATION_INSTALLER_GPT_GUID:
          is_bootable = True
  
    if not is_bootable:
      logging.critical('ERROR: mkinstaller would generate an unbootable image.' +
            'Are you building for a supported platform?')
      return []
  
    return ret
  
  
  def GetUsbDisks():
    """Get a list of all USB disks on the system.
  
    Returns:
      A list where each entry is of the format '/path/to/disk - <disk name>'
    """
    res = subprocess.run(['fx', 'list-usb-disks'], capture_output=True)
    res.check_returncode()
  
    disks = res.stdout.decode('utf-8').split('\n')
    return disks
  
  
  def IsUsbDisk(path):
    """Is the given path a USB disk?
  
    Args:
      path: a path that may represent a USB disk. Does not have to exist.
  
    Returns:
      True if the path represents a USB disk, False otherwise.
    """
    return path in map(lambda a: a.split()[0], GetUsbDisks())
  
  
  def EjectDisk(path):
    """Eject the given USB disk from the system."""
    system = platform.system()
    if system == 'Linux':
      subprocess.run(['eject', path])
    elif system == 'Darwin':
      subprocess.run(['diskutil', 'eject', path])
    logging.info('Ejected USB disk')
  
  
  def Main(args):
    path = args.FILE
    if args.create:
      if not args.force and os.path.exists(path):
        logging.critical(
            'File {} already exists, not creating an image. Use --force if you want to proceed.'
            .format(path),
            file=sys.stderr)
        return 1
    else:
      if not os.path.exists(path):
        logging.critical(
            ('Path {} does not exist, use --create to create a disk image.\n'
             'Detected USB devices:\n'
             '{}').format(path, '\n'.join(GetUsbDisks())),
            end='',
            file=sys.stderr)
        return 1
      if not IsUsbDisk(path):
        logging.critical(
            ('Path {} is not a USB device. Use -f to force.\n'
             'Detected USB devices:\n'
             '{}').format(path, '\n'.join(GetUsbDisks())),
            end='',
            file=sys.stderr)
        return 1
  
      if not os.access(path, os.W_OK) and sys.stdin.isatty():
        logging.warning('Changing ownership of {} to {}'.format(path,
                                                      os.environ.get('USER')))
        subprocess.run(
            ['sudo', 'chown', os.environ.get('USER'), path],
            stdin=sys.stdin,
            stdout=sys.stdout,
            stderr=sys.stderr)
      elif not os.access(path, os.W_OK):
        logging.critical('Cannot write to {}. Please check file permissions!'.format(path))
        return 1
  
    parts = GetPartitions()
    if not parts:
      return 1
  
    output = Image(args.FILE, not args.create, ParseSize(args.block_size))
    for p in parts:
      output.AddPartition(p)
  
    output.Finalise()
    if not args.create:
      EjectDisk(path)
    return 0
  
  
  if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Create a Fuchsia installer image.', prog='fx mkinstaller')
    parser.add_argument(
        '-c',
        '--create',
        action='store_true',
        help='Create a disk image instead of writing to an existing disk.')
    parser.add_argument(
        '-f',
        '--force',
        action='store_true',
        help='Force writing to an image that already exists or a disk that might not be a USB.'
    )
    parser.add_argument(
        '-b',
        '--block-size',
        type=str,
        default='2M',
        help='Block size (optionally suffixed by K, M, G) to write. Default is 2M'
    )
    parser.add_argument(
      '-v',
      '--verbose',
      action='store_true',
      help='Be verbose while creating disk images.'
    )
    parser.add_argument('FILE', help='Path to USB device or installer image')
    argv = parser.parse_args()
    level = logging.WARNING
    if argv.verbose:
      level = logging.DEBUG
    logging.basicConfig(format='mkinstaller: %(levelname)s: %(message)s', level=level)
    sys.exit(Main(argv))
  