| # Copyright 2022 The ChromiumOS Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Parse and operate based on disk layout files. |
| |
| For information on the JSON format, see: |
| http://dev.chromium.org/chromium-os/developer-guide/disk-layout-format |
| |
| The --adjust_part flag takes arguments like: |
| <label>:<op><size> |
| Where: |
| <label> is a label name as found in the disk layout file |
| <op> is one of the three: + - = |
| <size> is a number followed by an optional size qualifier: |
| B, KiB, MiB, GiB, TiB: bytes, kibi-, mebi-, gibi-, tebi- (base 1024) |
| B, K, M, G, T: short hand for above |
| B, KB, MB, GB, TB: bytes, kilo-, mega-, giga-, tera- (base 1000) |
| |
| This will set the ROOT-A partition size to 1 gibibytes (1024 * 1024 * 1024 * 1): |
| --adjust_part ROOT-A:=1GiB |
| This will grow the ROOT-A partition size by 500 mebibytes (1024 * 1024 * 500): |
| --adjust_part ROOT-A:+500MiB |
| This will shrink the ROOT-A partition size by 10 mebibytes (1024 * 1024 * 10): |
| --adjust_part ROOT-A:-20MiB |
| """ |
| |
| import argparse |
| import inspect |
| import os |
| import sys |
| from typing import List, Union |
| |
| from chromite.lib import disk_layout |
| |
| |
| def WritePartitionScript( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| sfilename: Union[str, os.PathLike], |
| vfilename: Union[str, os.PathLike], |
| ): |
| """Writes a shell script with functions for the base and requested layouts. |
| |
| Args: |
| options: Flags passed to the script |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| sfilename: Filename to write the finished script to |
| vfilename: Filename to write the partition variables json data to |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).WritePartitionScript(image_type, sfilename, vfilename) |
| |
| |
| def GetBlockSize( |
| _options: List[str], layout_filename: Union[str, os.PathLike] |
| ) -> str: |
| """Returns the partition table block size. |
| |
| Args: |
| options: Flags passed to the script. |
| layout_filename: Path to partition configuration file. |
| |
| Returns: |
| Block size of all partitions in the layout. |
| """ |
| |
| return disk_layout.DiskLayout(layout_filename).GetBlockSize() |
| |
| |
| def GetFilesystemBlockSize( |
| _options: List[str], layout_filename: Union[str, os.PathLike] |
| ) -> str: |
| """Returns the filesystem block size. |
| |
| This is used for all partitions in the table that have filesystems. |
| |
| Args: |
| options: Flags passed to the script. |
| layout_filename: Path to partition configuration file. |
| |
| Returns: |
| Block size of all filesystems in the layout. |
| """ |
| |
| return disk_layout.DiskLayout(layout_filename).GetFilesystemBlockSize() |
| |
| |
| def GetPartitionSize( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the partition size of a given partition for a given layout type. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Size of selected partition in bytes. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetPartitionSize(image_type, int(num)) |
| |
| |
| def GetFormat( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the format of a given partition for a given layout type. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Format of the selected partition's filesystem. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetFormat(image_type, int(num)) |
| |
| |
| def GetFilesystemFormat( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the filesystem format of a partition for a given layout type. |
| |
| Args: |
| options: Flags passed to the script |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| num: Number of the partition you want to read from |
| |
| Returns: |
| Format of the selected partition's filesystem |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetFilesystemFormat(image_type, int(num)) |
| |
| |
| def GetFilesystemSize( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> int: |
| """Returns the filesystem size of a given partition for a given layout type. |
| |
| If no filesystem size is specified, returns the partition size. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Size of selected partition filesystem in bytes. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetFilesystemSize(image_type, int(num)) |
| |
| |
| def GetImageTypes( |
| _options: List[str], layout_filename: Union[str, os.PathLike] |
| ) -> str: |
| """Returns a list of all the image types in the layout. |
| |
| Args: |
| options: Flags passed to the script. |
| layout_filename: Path to partition configuration file. |
| |
| Returns: |
| List of all image types. |
| """ |
| |
| image_types = disk_layout.DiskLayout(layout_filename).GetImageTypes() |
| return " ".join(image_types) |
| |
| |
| def GetFilesystemOptions( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the filesystem options of a given partition and layout type. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| The selected partition's filesystem options. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetFilesystemOptions(image_type, int(num)) |
| |
| |
| def GetLabel( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the label for a given partition. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Label of selected partition, or 'UNTITLED' if none specified. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetLabel(image_type, int(num)) |
| |
| |
| def GetNumber( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| label, |
| ) -> int: |
| """Returns the partition number of a given label. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| label: Number of the partition you want to read from. |
| |
| Returns: |
| The number of the partition corresponding to the label. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetNumber(image_type, label) |
| |
| |
| def GetReservedEraseBlocks( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> int: |
| """Returns the number of erase blocks reserved in the partition. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Number of reserved erase blocks. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetReservedEraseBlocks(image_type, int(num)) |
| |
| |
| def GetType( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the type of a given partition for a given layout. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| Type of the specified partition. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetType(image_type, int(num)) |
| |
| |
| def GetPartitions( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| ) -> str: |
| """Returns the partition numbers for the image_type. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| |
| Returns: |
| A space delimited string of partition numbers. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetPartitions(image_type) |
| |
| |
| def GetUUID( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| num: str, |
| ) -> str: |
| """Returns the filesystem UUID of a given partition for a given layout type. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| num: Number of the partition you want to read from. |
| |
| Returns: |
| UUID of specified partition. Defaults to random if not set. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).GetUUID(image_type, int(num)) |
| |
| |
| def DoDebugOutput( |
| options: List[str], |
| layout_filename: Union[str, os.PathLike], |
| image_type: str, |
| ): |
| """Prints out a human readable disk layout in on-disk order. |
| |
| Args: |
| options: Flags passed to the script. |
| layout_filename: Path to partition configuration file. |
| image_type: Type of image e.g. ALL/LIST/base/test/dev/factory_install. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).DoDebugOutput(image_type) |
| |
| |
| def GetVerityHashAlgo( |
| _options: List[str], layout_filename: Union[str, "os.PathLike[str]"] |
| ) -> str: |
| """Return hashing algo to be used with verity. |
| |
| Args: |
| options: Flags passed to the script. |
| layout_filename: Path to partition configuration file. |
| """ |
| return disk_layout.DiskLayout(layout_filename).GetVerityHashAlgo() |
| |
| |
| def Validate( |
| options: List[str], |
| image_type: str, |
| layout_filename: Union[str, os.PathLike], |
| ): |
| """Validates a layout file, used before reading sizes to check for errors. |
| |
| Args: |
| options: Flags passed to the script. |
| image_type: Type of image eg base/test/dev/factory_install. |
| layout_filename: Path to partition configuration file. |
| """ |
| |
| return disk_layout.DiskLayout( |
| layout_filename, options.adjust_part.split() |
| ).Validate(image_type) |
| |
| |
| class ArgsAction(argparse.Action): # pylint: disable=no-init |
| """Helper to add all arguments to an args array. |
| |
| ArgumentParser does not let you specify the same dest for multiple args. |
| We take care of appending to the 'args' array ourselves here. |
| """ |
| |
| def __call__(self, parser, namespace, values, option_string=None) -> None: |
| args = getattr(namespace, "args", []) |
| args.append(values) |
| setattr(namespace, "args", args) |
| |
| |
| class HelpAllAction(argparse.Action): |
| """Display all subcommands help in one go.""" |
| |
| def __init__(self, *args, **kwargs) -> None: |
| if "nargs" in kwargs: |
| raise ValueError("nargs not allowed") |
| kwargs["nargs"] = 0 |
| argparse.Action.__init__(self, *args, **kwargs) |
| |
| def __call__(self, parser, namespace, values, option_string=None) -> None: |
| print("%s\nCommands:" % (parser.description,), end="") |
| subparser = getattr(namespace, "help_all") |
| for key, subparser in namespace.help_all.choices.items(): |
| # Should we include the desc of each arg too ? |
| print( |
| "\n %s %s\n %s" |
| % ( |
| key, |
| subparser.get_default("help_all"), |
| subparser.description, |
| ) |
| ) |
| sys.exit(0) |
| |
| |
| def GetParser(): |
| """Return a parser for the CLI. |
| |
| We use the function docstring to build the cli argument, help text |
| and their arguments. |
| """ |
| parser = argparse.ArgumentParser( |
| description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter, |
| ) |
| parser.add_argument( |
| "--adjust_part", |
| metavar="SPEC", |
| default="", |
| help="adjust partition sizes", |
| ) |
| |
| action_map = { |
| "write": WritePartitionScript, |
| "readblocksize": GetBlockSize, |
| "readfsblocksize": GetFilesystemBlockSize, |
| "readpartsize": GetPartitionSize, |
| "readformat": GetFormat, |
| "readfsformat": GetFilesystemFormat, |
| "readfssize": GetFilesystemSize, |
| "readimagetypes": GetImageTypes, |
| "readfsoptions": GetFilesystemOptions, |
| "readlabel": GetLabel, |
| "readnumber": GetNumber, |
| "readreservederaseblocks": GetReservedEraseBlocks, |
| "readtype": GetType, |
| "readpartitionnums": GetPartitions, |
| "readuuid": GetUUID, |
| "readverityhashalgo": GetVerityHashAlgo, |
| "debug": DoDebugOutput, |
| "validate": Validate, |
| } |
| |
| subparsers = parser.add_subparsers( |
| title="Commands", dest="command", required=True |
| ) |
| |
| for name, func in sorted(action_map.items()): |
| # Turn the func's docstring into something we can show the user. |
| desc, doc = func.__doc__.split("\n", 1) |
| # Extract the help for each argument. |
| args_help = {} |
| for line in doc.splitlines(): |
| if ":" in line: |
| arg, text = line.split(":", 1) |
| args_help[arg.strip()] = text.strip() |
| |
| argspec = inspect.getfullargspec(func) |
| # Skip the first argument as that'll be the options field. |
| args = argspec.args[1:] |
| |
| subparser = subparsers.add_parser(name, description=desc, help=desc) |
| subparser.set_defaults( |
| callback=func, help_all=" ".join("<%s>" % x for x in args) |
| ) |
| for arg in args: |
| subparser.add_argument(arg, action=ArgsAction, help=args_help[arg]) |
| |
| parser.add_argument( |
| "--help-all", |
| action=HelpAllAction, |
| default=subparsers, |
| help="show all commands and their help in one screen", |
| ) |
| |
| return parser |
| |
| |
| def main(argv) -> None: |
| parser = GetParser() |
| opts = parser.parse_args(argv) |
| |
| ret = opts.callback(opts, *opts.args) |
| if ret is not None: |
| print(ret) |
| |
| |
| if __name__ == "__main__": |
| sys.exit(main(sys.argv[1:])) |