src/firmware/tools/gpt.py - fuchsia - Git at Google

 #!/usr/bin/env fuchsia-vendored-python
 #
 # Copyright 2022 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.
 """Script to create Fuchsia GPT images.

 This script wraps the `sgdisk` utility to provide a simpler interface tailored
 specifically for Fuchsia GPT images.

 When creating a new GPT, the layout is specified by a JSON file that looks like:

 {
     "size_mib": disk size in MiB
     "block_size": disk R/W block size in bytes
     "alignment_kib": partition start alignment in KiB; should generally be a
                      multiple of the disk erase block size
     "partitions": [
         {
             "name": partition name
             "size_kib": size in KiB, or "fill" to use the rest of the disk
             "unique_guid": unique GUID; omit to assign randomly
             "type_guid": type GUID, or one of ("fvm", "vbmeta", "zircon") to use
                          the standard Fuchsia types
         }
     ]
 }

 Each partition starts immediately following the previous; this tool does not
 support creating GPTs with out-of-order partitions or gaps between partitions.
 """

 import argparse
 import contextlib
 import dataclasses
 import json
 import logging
 import os
 import pathlib
 import re
 import subprocess
 import sys
 import tempfile
 import textwrap
 from typing import Dict, Iterator, List, Optional

 # A regexp to capture GUIDs as reported by `sgdisk`.
 # Example: 4B7A23A5-E9D1-456A-BACB-78B23D855324.
 GUID_RE = r"[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}"

 # Maps partition type names to known type GUIDs.
 # https://cs.opensource.google/fuchsia/fuchsia/+/main:zircon/system/public/zircon/hw/gpt.h.
 _TYPE_GUID_FROM_NAME = {
     "fvm": "49FD7CB8-DF15-4E73-B9D9-992070127F0F",
     "vbmeta": "421A8BFC-85D9-4D85-ACDA-B64EEC0133E9",
     "zircon": "9B37FFF6-2E58-466A-983A-F7926D0B04E0",
 }


 @dataclasses.dataclass
 class Partition:
     """GPT partition info."""

     # 1-based index in the GPT entry array.
     index: int
     name: str

     # Partition first block, last block (inclusive), and size in blocks.
     # Logical block addressing (LBA) is used here because it's what GPT uses
     # internally and is indivisible so will always be integers for exact
     # calculations (unlike KiB/MiB).
     first_lba: int
     last_lba: int
     size_lba: int

     type_guid: str
     unique_guid: str
     attr_flags: str


 @contextlib.contextmanager
 def gen_disk_image(source_gpt: str, image_size_mib: int) -> Iterator[str]:
     """Generates empty disk image using the given source GPT."""
     print("Generating disk image")

     with tempfile.TemporaryDirectory() as temp_dir:
         image_path = os.path.join(temp_dir, "disk_image")

         with open(image_path, "wb") as file:
             file.truncate(image_size_mib * 1024 * 1024)

         subprocess.run(
             ["sgdisk", f"--load-backup={source_gpt}", image_path],
             capture_output=True,
             check=True,
         )
         yield image_path
         print("Deleting disk image")


 # This function is a bit brittle as it depends heavily on `sgdisk` output format
 # which is not guaranteed to be stable. If it becomes unreliable we may either
 # need to bring a fixed version of sgdisk into //prebuilts, or just implement
 # GPT handling ourselves.
 def get_partitions(disk_file: str) -> List[Partition]:
     """Returns a list of all Partitions."""
     # Find all the partition indices first.
     # Output looks like:
     # -----------------------------------------------------------------
     # Disk test_disk.img: 131072 sectors, 64.0 MiB
     # Sector size (logical): 512 bytes
     # Disk identifier (GUID): F1D74D6E-1DB4-4E46-8D2A-87AC93543ACA
     # Partition table holds up to 128 entries
     # Main partition table begins at sector 2 and ends at sector 33
     # First usable sector is 34, last usable sector is 131038
     # Partitions will be aligned on 2048-sector boundaries
     # Total free space is 126909 sectors (62.0 MiB)
     #
     # Number  Start (sector)    End (sector)  Size       Code  Name
     #    1            2048            4095   1024.0 KiB  8300
     #    2            4096            6143   1024.0 KiB  8300
     # -----------------------------------------------------------------
     # Here we just care about the index, which will be the first entry on lines
     # that start with a number.
     output = subprocess.run(
         ["sgdisk", disk_file, "--print"],
         text=True,
         capture_output=True,
         check=True,
     ).stdout
     indices = re.findall(r"^\s*(\d+)\s+", output, re.MULTILINE)

     partitions = []
     for index in indices:
         # Now fill in the full partition information. Output looks like:
         # -----------------------------------------------------------------
         # Partition GUID code: 0FC63DAF-8483-4772-8E79-3D69D8477DE4 (Linux filesystem)
         # Partition unique GUID: 4B7A23A5-E9D1-456A-BACB-78B23D855324
         # First sector: 2048 (at 1024.0 KiB)
         # Last sector: 4095 (at 2.0 MiB)
         # Partition size: 2048 sectors (1024.0 KiB)
         # Attribute flags: 0000000000000000
         # Partition name: 'foo'
         # -----------------------------------------------------------------
         # We could potentially query all partitions in a single command, but
         # the extra delay is negligible and this is simpler.
         output = subprocess.run(
             ["sgdisk", disk_file, "--info", index],
             text=True,
             capture_output=True,
             check=True,
         ).stdout

         def get_param(pattern):
             return re.search(f"^{pattern}.*?$", output, re.MULTILINE).group(1)

         partitions.append(
             Partition(
                 index=int(index),
                 name=get_param(r"Partition name: '(.*?)'"),
                 first_lba=int(get_param(r"First sector: (\d+)")),
                 last_lba=int(get_param(r"Last sector: (\d+)")),
                 size_lba=int(get_param(r"Partition size: (\d+) sectors")),
                 type_guid=get_param(f"Partition GUID code: ({GUID_RE})"),
                 unique_guid=get_param(f"Partition unique GUID: ({GUID_RE})"),
                 attr_flags=get_param(r"Attribute flags: (\S*)"),
             )
         )

     return partitions


 def get_partition(
     disk_file: str,
     *,
     part_name: Optional[str] = None,
     index: Optional[int] = None,
 ) -> Partition:
     """Returns a single Partition keyed by name and/or index.

     Raises ValueError if the requested partition couldn't be found.
     """
     if part_name is None and index is None:
         raise ValueError("Must specify either a name or index")

     def filter(part):
         return (part_name is None or part.name == part_name) and (
             index is None or part.index == index
         )

     partitions = get_partitions(disk_file)
     try:
         return next(p for p in partitions if filter(p))
     except StopIteration:
         raise ValueError(f"Failed to find partition {part_name}")


 def get_part_num(disk_file: str, part_name: str) -> int:
     """Returns the numeric index of a given partition.

     Raises ValueError if the given partition name doesn't exist.
     """
     return get_partition(disk_file, part_name=part_name).index


 def mib_to_lba(mib: int, blk_size: int) -> int:
     """Convert size in MB to lba"""
     return int(mib * 1024 * 1024 / blk_size)


 def save_gpt(
     disk_path: pathlib.Path,
     block_size: int,
     *,
     mbr_path: Optional[pathlib.Path] = None,
     primary_path: Optional[pathlib.Path] = None,
     backup_path: Optional[pathlib.Path] = None,
 ):
     """Extracts and saves the GPT sections from a given disk image.

     Args:
         disk_path: path to the disk image.
         block_size: block size in bytes.
         mbr_path: the path to write the MBR binary, or None.
         primary_path: the path to write the primary GPT binary, or None.
         backup_path: the path to write the backup GPT binary, or None.
     """
     # The MBR is always the first block.
     mbr_size = block_size

     # GPT header is always one block, followed by 16KiB of space for partition
     # entries. Technically it can be more than 16KiB but none of our layouts
     # need this so we don't worry about it.
     gpt_size = block_size + (16 * 1024)

     with open(disk_path, "rb") as disk_file:
         mbr = disk_file.read(mbr_size)
         primary = disk_file.read(gpt_size)
         disk_file.seek(-gpt_size, 2)
         backup = disk_file.read()

     for path, contents in (
         (mbr_path, mbr),
         (primary_path, primary),
         (backup_path, backup),
     ):
         if path:
             path.write_bytes(contents)


 def delete_part(disk_file: str, part_num: int):
     """Delete a partition from the GPT"""
     subprocess.run(
         ["sgdisk", "--delete", f"{part_num}", disk_file],
         capture_output=True,
         check=True,
     )


 def delete_part_by_name(disk_file: str, name: str):
     """Delete a partition by name"""
     part_num = get_part_num(disk_file, name)
     delete_part(disk_file, part_num)


 def new_part(disk_file: str, part_num: int, part: Partition):
     """Create a new partition"""
     # does not set attr_flags - this is fine for now as current Fuchsia devices
     # don't depend on it.
     subprocess.run(
         [
             "sgdisk",
             "-a",
             "1",
             "--new",
             f"{part_num}:{part.first_lba}:{part.last_lba}",
             disk_file,
         ],
         capture_output=True,
         check=True,
     )

     subprocess.run(
         ["sgdisk", "--change-name", f"{part_num}:{part.name}", disk_file],
         capture_output=True,
         check=True,
     )
     subprocess.run(
         ["sgdisk", "--typecode", f"{part_num}:{part.type_guid}", disk_file],
         capture_output=True,
         check=True,
     )

     # If we don't have a specific UUID to use, "R" creates a random one.
     guid = part.unique_guid or "R"
     subprocess.run(
         ["sgdisk", "--partition-guid", f"{part_num}:{guid}", disk_file],
         capture_output=True,
         check=True,
     )


 def gpt_rename_part(
     disk_image: str,
     old_part_name: str,
     new_part_name: str,
     new_part_type_guid: str = None,
 ):
     """Rename a partition in the GPT"""
     orig_part = get_partition(disk_image, part_name=old_part_name)
     new_partition = Partition(
         index=orig_part.index,
         name=new_part_name,
         first_lba=orig_part.first_lba,
         last_lba=orig_part.last_lba,
         size_lba=orig_part.size_lba,
         type_guid=new_part_type_guid or orig_part.type_guid,
         unique_guid="",  # generate new
         attr_flags=orig_part.attr_flags,
     )

     delete_part(disk_image, orig_part.index)
     new_part(disk_image, orig_part.index, new_partition)


 def create_gpt(spec: Dict, out_dir: pathlib.Path):
     """Creates a GPT from the given specification.

     See the file docstring for a description of the spec.

     Args:
         spec_path: the layout specification.
         out_dir: directory to write resulting images.

     Raises:
         ValueError if the given spec is invalid.
     """
     disk_size = spec["size_mib"] * 1024 * 1024
     alignment = spec["alignment_kib"] * 1024
     block_size = spec["block_size"]

     # TODO: sgdisk appears to hardcode a 512-byte block size when working with
     # raw image files and I can't find any option to modify this behavior. For
     # now all our disks also have 512-byte blocks so it's fine, just punt on
     # this until we need it.
     if block_size != 512:
         raise NotImplementedError(
             "Only 512-byte blocks are currently implemented"
         )

     if disk_size % alignment != 0:
         raise ValueError("Disk size must be a multiple of alignment")
     if alignment % block_size != 0:
         raise ValueError("Disk alignment must be a multiple of block size")

     with tempfile.TemporaryDirectory() as temp_dir:
         temp_dir = pathlib.Path(temp_dir)
         disk_path = temp_dir / "disk.img"

         with open(disk_path, "wb") as disk_file:
             disk_file.truncate(disk_size)

         # We generate a single big `sgdisk` command so that it only has to read
         # the file once, do all the operations in memory, and then write the
         # final result once.

         # Initialize the disk.
         command = ["sgdisk", disk_path, "--clear"]

         # Partition alignment is in block units.
         command += ["--set-alignment", str(alignment // block_size)]

         # Add the partitions.
         for i, part in enumerate(spec["partitions"], start=1):
             name = part["name"]

             size_kib = part["size_kib"]
             if size_kib == "fill":
                 # "fill" means use the rest of the partition, which sgdisk does
                 # if you pass "0" as the end sector.
                 end = 0
             else:
                 # Otherwise, use "+" notation to tell sgdisk to set the end
                 # relative to the partition start (i.e. partition size rather
                 # than end position).
                 end = f"+{size_kib}K"

             # If a unique GUID isn't provided, "R" tells sgdisk to generate one
             # randomly.
             unique_guid = part.get("unique_guid", "R")

             # Fetch the type GUID from the name map if it exists.
             type_guid = part["type_guid"]
             if type_guid in _TYPE_GUID_FROM_NAME:
                 type_guid = _TYPE_GUID_FROM_NAME[type_guid]

             command += ["--new", f"{i}:0:{end}"]
             command += ["--change-name", f"{i}:{name}"]
             command += ["--typecode", f"{i}:{type_guid}"]
             command += ["--partition-guid", f"{i}:{unique_guid}"]

         # Run the command to write the GPT.
         subprocess.run(command, check=True, capture_output=True, text=True)

         # Extract the pieces to individual files.
         out_dir.mkdir(exist_ok=True)
         save_gpt(
             disk_path,
             block_size,
             mbr_path=out_dir / "mbr.bin",
             primary_path=out_dir / "primary.bin",
             backup_path=out_dir / "backup.bin",
         )

         # Grab a high level summary and the final partition details for the
         # README file.
         summary = subprocess.run(
             ["sgdisk", disk_path, "--print"],
             check=True,
             capture_output=True,
             text=True,
         ).stdout.strip()

         partitions = get_partitions(disk_path)

     # Write a README to document this GPT so it's easy for a reader to see
     # where it came from and what it looks like.
     (out_dir / "README.md").write_text(
         textwrap.dedent(
             """\
             # {name} GPT binaries

             Generated by
             https://cs.opensource.google/fuchsia/fuchsia/+/main:src/firmware/tools/gpt.py.

             ## Summary

             ```
             {summary}
             ```

             ## Details

             ```
             {details}
             ```

             ## Input specification

             ```
             {spec}
             ```

             ## Command

             ```
             {command}
             ```

             """
         ).format(
             name=out_dir.name,
             summary=summary,
             details=json.dumps(
                 [dataclasses.asdict(p) for p in partitions], indent=2
             ),
             spec=json.dumps(spec, indent=2),
             command=" ".join(str(c) for c in command),
         )
     )


 def parse_args() -> argparse.Namespace:
     """Parses commandline args."""
     parser = argparse.ArgumentParser(
         description=__doc__,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )

     subparsers = parser.add_subparsers(
         title="Subcommands",
         dest="action",
         help="Action to perform",
         required=True,
     )

     create_parser = subparsers.add_parser(
         "create", help="create GPT images from a given specification"
     )
     create_parser.add_argument(
         "spec",
         type=pathlib.Path,
         help="layout specification JSON file; see --help for details",
     )
     create_parser.add_argument(
         "out_dir",
         type=pathlib.Path,
         help="directory to write the resulting files to, with names:"
         " [mbr.bin, primary.bin, backup.bin, README.md]",
     )

     return parser.parse_args()


 def main():
     """Main function entry point. Raises an exception on failure."""
     logging.basicConfig(level=logging.INFO)

     args = parse_args()

     if args.action == "create":
         create_gpt(json.loads(args.spec.read_text()), args.out_dir)

     else:
         raise ValueError(f"Unknown action {args.action}")


 if __name__ == "__main__":
     try:
         main()
         sys.exit(0)
     except KeyboardInterrupt:
         # Catch this to avoid printing exception context in this case.
         sys.exit(1)
     except subprocess.CalledProcessError as error:
         # If we failed with a subprocess that captured stdout/stderr,
         # dump it out to the user so they can see what went wrong.
         logging.error(
             "Command failed: `%s`", " ".join(str(c) for c in error.cmd)
         )
         if error.stdout or error.stderr:
             logging.error(
                 "\n" "----\n" "%s\n" "----",
                 (error.stdout or "") + (error.stderr or ""),
             )
         raise
	#!/usr/bin/env fuchsia-vendored-python
	#
	# Copyright 2022 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.
	"""Script to create Fuchsia GPT images.

	This script wraps the `sgdisk` utility to provide a simpler interface tailored
	specifically for Fuchsia GPT images.

	When creating a new GPT, the layout is specified by a JSON file that looks like:

	{
	"size_mib": disk size in MiB
	"block_size": disk R/W block size in bytes
	"alignment_kib": partition start alignment in KiB; should generally be a
	multiple of the disk erase block size
	"partitions": [
	{
	"name": partition name
	"size_kib": size in KiB, or "fill" to use the rest of the disk
	"unique_guid": unique GUID; omit to assign randomly
	"type_guid": type GUID, or one of ("fvm", "vbmeta", "zircon") to use
	the standard Fuchsia types
	}
	]
	}

	Each partition starts immediately following the previous; this tool does not
	support creating GPTs with out-of-order partitions or gaps between partitions.
	"""

	import argparse
	import contextlib
	import dataclasses
	import json
	import logging
	import os
	import pathlib
	import re
	import subprocess
	import sys
	import tempfile
	import textwrap
	from typing import Dict, Iterator, List, Optional

	# A regexp to capture GUIDs as reported by `sgdisk`.
	# Example: 4B7A23A5-E9D1-456A-BACB-78B23D855324.
	GUID_RE = r"[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}"

	# Maps partition type names to known type GUIDs.
	# https://cs.opensource.google/fuchsia/fuchsia/+/main:zircon/system/public/zircon/hw/gpt.h.
	_TYPE_GUID_FROM_NAME = {
	"fvm": "49FD7CB8-DF15-4E73-B9D9-992070127F0F",
	"vbmeta": "421A8BFC-85D9-4D85-ACDA-B64EEC0133E9",
	"zircon": "9B37FFF6-2E58-466A-983A-F7926D0B04E0",
	}


	@dataclasses.dataclass
	class Partition:
	"""GPT partition info."""

	# 1-based index in the GPT entry array.
	index: int
	name: str

	# Partition first block, last block (inclusive), and size in blocks.
	# Logical block addressing (LBA) is used here because it's what GPT uses
	# internally and is indivisible so will always be integers for exact
	# calculations (unlike KiB/MiB).
	first_lba: int
	last_lba: int
	size_lba: int

	type_guid: str
	unique_guid: str
	attr_flags: str


	@contextlib.contextmanager
	def gen_disk_image(source_gpt: str, image_size_mib: int) -> Iterator[str]:
	"""Generates empty disk image using the given source GPT."""
	print("Generating disk image")

	with tempfile.TemporaryDirectory() as temp_dir:
	image_path = os.path.join(temp_dir, "disk_image")

	with open(image_path, "wb") as file:
	file.truncate(image_size_mib * 1024 * 1024)

	subprocess.run(
	["sgdisk", f"--load-backup={source_gpt}", image_path],
	capture_output=True,
	check=True,
	)
	yield image_path
	print("Deleting disk image")


	# This function is a bit brittle as it depends heavily on `sgdisk` output format
	# which is not guaranteed to be stable. If it becomes unreliable we may either
	# need to bring a fixed version of sgdisk into //prebuilts, or just implement
	# GPT handling ourselves.
	def get_partitions(disk_file: str) -> List[Partition]:
	"""Returns a list of all Partitions."""
	# Find all the partition indices first.
	# Output looks like:
	# -----------------------------------------------------------------
	# Disk test_disk.img: 131072 sectors, 64.0 MiB
	# Sector size (logical): 512 bytes
	# Disk identifier (GUID): F1D74D6E-1DB4-4E46-8D2A-87AC93543ACA
	# Partition table holds up to 128 entries
	# Main partition table begins at sector 2 and ends at sector 33
	# First usable sector is 34, last usable sector is 131038
	# Partitions will be aligned on 2048-sector boundaries
	# Total free space is 126909 sectors (62.0 MiB)
	#
	# Number Start (sector) End (sector) Size Code Name
	# 1 2048 4095 1024.0 KiB 8300
	# 2 4096 6143 1024.0 KiB 8300
	# -----------------------------------------------------------------
	# Here we just care about the index, which will be the first entry on lines
	# that start with a number.
	output = subprocess.run(
	["sgdisk", disk_file, "--print"],
	text=True,
	capture_output=True,
	check=True,
	).stdout
	indices = re.findall(r"^\s*(\d+)\s+", output, re.MULTILINE)

	partitions = []
	for index in indices:
	# Now fill in the full partition information. Output looks like:
	# -----------------------------------------------------------------
	# Partition GUID code: 0FC63DAF-8483-4772-8E79-3D69D8477DE4 (Linux filesystem)
	# Partition unique GUID: 4B7A23A5-E9D1-456A-BACB-78B23D855324
	# First sector: 2048 (at 1024.0 KiB)
	# Last sector: 4095 (at 2.0 MiB)
	# Partition size: 2048 sectors (1024.0 KiB)
	# Attribute flags: 0000000000000000
	# Partition name: 'foo'
	# -----------------------------------------------------------------
	# We could potentially query all partitions in a single command, but
	# the extra delay is negligible and this is simpler.
	output = subprocess.run(
	["sgdisk", disk_file, "--info", index],
	text=True,
	capture_output=True,
	check=True,
	).stdout

	def get_param(pattern):
	return re.search(f"^{pattern}.*?$", output, re.MULTILINE).group(1)

	partitions.append(
	Partition(
	index=int(index),
	name=get_param(r"Partition name: '(.*?)'"),
	first_lba=int(get_param(r"First sector: (\d+)")),
	last_lba=int(get_param(r"Last sector: (\d+)")),
	size_lba=int(get_param(r"Partition size: (\d+) sectors")),
	type_guid=get_param(f"Partition GUID code: ({GUID_RE})"),
	unique_guid=get_param(f"Partition unique GUID: ({GUID_RE})"),
	attr_flags=get_param(r"Attribute flags: (\S*)"),
	)
	)

	return partitions


	def get_partition(
	disk_file: str,
	*,
	part_name: Optional[str] = None,
	index: Optional[int] = None,
	) -> Partition:
	"""Returns a single Partition keyed by name and/or index.

	Raises ValueError if the requested partition couldn't be found.
	"""
	if part_name is None and index is None:
	raise ValueError("Must specify either a name or index")

	def filter(part):
	return (part_name is None or part.name == part_name) and (
	index is None or part.index == index
	)

	partitions = get_partitions(disk_file)
	try:
	return next(p for p in partitions if filter(p))
	except StopIteration:
	raise ValueError(f"Failed to find partition {part_name}")


	def get_part_num(disk_file: str, part_name: str) -> int:
	"""Returns the numeric index of a given partition.

	Raises ValueError if the given partition name doesn't exist.
	"""
	return get_partition(disk_file, part_name=part_name).index


	def mib_to_lba(mib: int, blk_size: int) -> int:
	"""Convert size in MB to lba"""
	return int(mib * 1024 * 1024 / blk_size)


	def save_gpt(
	disk_path: pathlib.Path,
	block_size: int,
	*,
	mbr_path: Optional[pathlib.Path] = None,
	primary_path: Optional[pathlib.Path] = None,
	backup_path: Optional[pathlib.Path] = None,
	):
	"""Extracts and saves the GPT sections from a given disk image.

	Args:
	disk_path: path to the disk image.
	block_size: block size in bytes.
	mbr_path: the path to write the MBR binary, or None.
	primary_path: the path to write the primary GPT binary, or None.
	backup_path: the path to write the backup GPT binary, or None.
	"""
	# The MBR is always the first block.
	mbr_size = block_size

	# GPT header is always one block, followed by 16KiB of space for partition
	# entries. Technically it can be more than 16KiB but none of our layouts
	# need this so we don't worry about it.
	gpt_size = block_size + (16 * 1024)

	with open(disk_path, "rb") as disk_file:
	mbr = disk_file.read(mbr_size)
	primary = disk_file.read(gpt_size)
	disk_file.seek(-gpt_size, 2)
	backup = disk_file.read()

	for path, contents in (
	(mbr_path, mbr),
	(primary_path, primary),
	(backup_path, backup),
	):
	if path:
	path.write_bytes(contents)


	def delete_part(disk_file: str, part_num: int):
	"""Delete a partition from the GPT"""
	subprocess.run(
	["sgdisk", "--delete", f"{part_num}", disk_file],
	capture_output=True,
	check=True,
	)


	def delete_part_by_name(disk_file: str, name: str):
	"""Delete a partition by name"""
	part_num = get_part_num(disk_file, name)
	delete_part(disk_file, part_num)


	def new_part(disk_file: str, part_num: int, part: Partition):
	"""Create a new partition"""
	# does not set attr_flags - this is fine for now as current Fuchsia devices
	# don't depend on it.
	subprocess.run(
	[
	"sgdisk",
	"-a",
	"1",
	"--new",
	f"{part_num}:{part.first_lba}:{part.last_lba}",
	disk_file,
	],
	capture_output=True,
	check=True,
	)

	subprocess.run(
	["sgdisk", "--change-name", f"{part_num}:{part.name}", disk_file],
	capture_output=True,
	check=True,
	)
	subprocess.run(
	["sgdisk", "--typecode", f"{part_num}:{part.type_guid}", disk_file],
	capture_output=True,
	check=True,
	)

	# If we don't have a specific UUID to use, "R" creates a random one.
	guid = part.unique_guid or "R"
	subprocess.run(
	["sgdisk", "--partition-guid", f"{part_num}:{guid}", disk_file],
	capture_output=True,
	check=True,
	)


	def gpt_rename_part(
	disk_image: str,
	old_part_name: str,
	new_part_name: str,
	new_part_type_guid: str = None,
	):
	"""Rename a partition in the GPT"""
	orig_part = get_partition(disk_image, part_name=old_part_name)
	new_partition = Partition(
	index=orig_part.index,
	name=new_part_name,
	first_lba=orig_part.first_lba,
	last_lba=orig_part.last_lba,
	size_lba=orig_part.size_lba,
	type_guid=new_part_type_guid or orig_part.type_guid,
	unique_guid="", # generate new
	attr_flags=orig_part.attr_flags,
	)

	delete_part(disk_image, orig_part.index)
	new_part(disk_image, orig_part.index, new_partition)


	def create_gpt(spec: Dict, out_dir: pathlib.Path):
	"""Creates a GPT from the given specification.

	See the file docstring for a description of the spec.

	Args:
	spec_path: the layout specification.
	out_dir: directory to write resulting images.

	Raises:
	ValueError if the given spec is invalid.
	"""
	disk_size = spec["size_mib"] * 1024 * 1024
	alignment = spec["alignment_kib"] * 1024
	block_size = spec["block_size"]

	# TODO: sgdisk appears to hardcode a 512-byte block size when working with
	# raw image files and I can't find any option to modify this behavior. For
	# now all our disks also have 512-byte blocks so it's fine, just punt on
	# this until we need it.
	if block_size != 512:
	raise NotImplementedError(
	"Only 512-byte blocks are currently implemented"
	)

	if disk_size % alignment != 0:
	raise ValueError("Disk size must be a multiple of alignment")
	if alignment % block_size != 0:
	raise ValueError("Disk alignment must be a multiple of block size")

	with tempfile.TemporaryDirectory() as temp_dir:
	temp_dir = pathlib.Path(temp_dir)
	disk_path = temp_dir / "disk.img"

	with open(disk_path, "wb") as disk_file:
	disk_file.truncate(disk_size)

	# We generate a single big `sgdisk` command so that it only has to read
	# the file once, do all the operations in memory, and then write the
	# final result once.

	# Initialize the disk.
	command = ["sgdisk", disk_path, "--clear"]

	# Partition alignment is in block units.
	command += ["--set-alignment", str(alignment // block_size)]

	# Add the partitions.
	for i, part in enumerate(spec["partitions"], start=1):
	name = part["name"]

	size_kib = part["size_kib"]
	if size_kib == "fill":
	# "fill" means use the rest of the partition, which sgdisk does
	# if you pass "0" as the end sector.
	end = 0
	else:
	# Otherwise, use "+" notation to tell sgdisk to set the end
	# relative to the partition start (i.e. partition size rather
	# than end position).
	end = f"+{size_kib}K"

	# If a unique GUID isn't provided, "R" tells sgdisk to generate one
	# randomly.
	unique_guid = part.get("unique_guid", "R")

	# Fetch the type GUID from the name map if it exists.
	type_guid = part["type_guid"]
	if type_guid in _TYPE_GUID_FROM_NAME:
	type_guid = _TYPE_GUID_FROM_NAME[type_guid]

	command += ["--new", f"{i}:0:{end}"]
	command += ["--change-name", f"{i}:{name}"]
	command += ["--typecode", f"{i}:{type_guid}"]
	command += ["--partition-guid", f"{i}:{unique_guid}"]

	# Run the command to write the GPT.
	subprocess.run(command, check=True, capture_output=True, text=True)

	# Extract the pieces to individual files.
	out_dir.mkdir(exist_ok=True)
	save_gpt(
	disk_path,
	block_size,
	mbr_path=out_dir / "mbr.bin",
	primary_path=out_dir / "primary.bin",
	backup_path=out_dir / "backup.bin",
	)

	# Grab a high level summary and the final partition details for the
	# README file.
	summary = subprocess.run(
	["sgdisk", disk_path, "--print"],
	check=True,
	capture_output=True,
	text=True,
	).stdout.strip()

	partitions = get_partitions(disk_path)

	# Write a README to document this GPT so it's easy for a reader to see
	# where it came from and what it looks like.
	(out_dir / "README.md").write_text(
	textwrap.dedent(
	"""\
	# {name} GPT binaries

	Generated by
	https://cs.opensource.google/fuchsia/fuchsia/+/main:src/firmware/tools/gpt.py.

	## Summary

	```
	{summary}
	```

	## Details

	```
	{details}
	```

	## Input specification

	```
	{spec}
	```

	## Command

	```
	{command}
	```

	"""
	).format(
	name=out_dir.name,
	summary=summary,
	details=json.dumps(
	[dataclasses.asdict(p) for p in partitions], indent=2
	),
	spec=json.dumps(spec, indent=2),
	command=" ".join(str(c) for c in command),
	)
	)


	def parse_args() -> argparse.Namespace:
	"""Parses commandline args."""
	parser = argparse.ArgumentParser(
	description=__doc__,
	formatter_class=argparse.RawDescriptionHelpFormatter,
	)

	subparsers = parser.add_subparsers(
	title="Subcommands",
	dest="action",
	help="Action to perform",
	required=True,
	)

	create_parser = subparsers.add_parser(
	"create", help="create GPT images from a given specification"
	)
	create_parser.add_argument(
	"spec",
	type=pathlib.Path,
	help="layout specification JSON file; see --help for details",
	)
	create_parser.add_argument(
	"out_dir",
	type=pathlib.Path,
	help="directory to write the resulting files to, with names:"
	" [mbr.bin, primary.bin, backup.bin, README.md]",
	)

	return parser.parse_args()


	def main():
	"""Main function entry point. Raises an exception on failure."""
	logging.basicConfig(level=logging.INFO)

	args = parse_args()

	if args.action == "create":
	create_gpt(json.loads(args.spec.read_text()), args.out_dir)

	else:
	raise ValueError(f"Unknown action {args.action}")


	if __name__ == "__main__":
	try:
	main()
	sys.exit(0)
	except KeyboardInterrupt:
	# Catch this to avoid printing exception context in this case.
	sys.exit(1)
	except subprocess.CalledProcessError as error:
	# If we failed with a subprocess that captured stdout/stderr,
	# dump it out to the user so they can see what went wrong.
	logging.error(
	"Command failed: `%s`", " ".join(str(c) for c in error.cmd)
	)
	if error.stdout or error.stderr:
	logging.error(
	"\n" "----\n" "%s\n" "----",
	(error.stdout or "") + (error.stderr or ""),
	)
	raise