gbl/tools/gen_gpt_disk.py

#!/usr/bin/env python3
#
# Copyright (C) 2023 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generate a GPT disk

Example:

The following command creates a 16mb disk file gpt.bin with two partitions:

  1) "boot_a" of size 4096KB, with data initialized from "file_a",
  2) "boot_b" of size 8192KB, with data initialized to zeroes.

    gen_gpt_disk.py ./gpt.bin 16M
        --partition "boot_a,4096K,<path to file_a>" \
        --partition "boot_b,8192K,"
"""

import argparse
import os
import pathlib
import shutil
import sys
import subprocess
import tempfile

GPT_BLOCK_SIZE = 512


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('out', help="output file")
    parser.add_argument('disk_size',
                        type=str,
                        help="disk size of the image. i.e. 64k, 1M etc")
    parser.add_argument("--partition",
                        type=str,
                        action='append',
                        help="specifies a partition. Format should be"
                        "--partition=<part name>,<size>,<file name>")
    return parser.parse_args()


def parse_size_str(size_str: str) -> int:
    """Parse size string into integer, taking into account unit.

    Example:
        2M, 2m -> 2*1024*1024
        2K, 2k -> 2*1024
        512 -> 512
    """
    size_str = size_str.lower()
    if size_str.endswith("m"):
        return int(size_str[:-1]) * 1024 * 1024
    if size_str.endswith("k"):
        return int(size_str[:-1]) * 1024
    return int(size_str)


def _append(src_file: str, offset: int, size: int, dst_file: str):
    """Append a portion of a file to the end of another file

    Args:
        src_file: path to the source file
        offset: offset in the source file
        size: number of bytes to append
        dst_file: destination file

    Returns:
        number of bytes actually copied.
    """
    with open(src_file, 'rb') as src_f:
        src_f.seek(offset, 0)
        data = src_f.read(size)
        if len(data) != size:
            raise ValueError(
                f"Want {size} bytes from {src_file}, but got {len(data)}.")
        with open(dst_file, 'ab') as dst_f:
            dst_f.write(data)
    return size


def main() -> int:
    args = parse_args()
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_disk = pathlib.Path(temp_dir) / "temp_disk"

        disk_size = parse_size_str(args.disk_size)
        temp_disk.write_bytes(disk_size * b'\x00')

        part_start = 34 * GPT_BLOCK_SIZE  # MBR + GPT header + entries
        partition_info = []
        sgdisk_command = [
            "sgdisk",
            str(temp_disk), "--clear", "--set-alignment", "1"
        ]

        for i, part in enumerate(args.partition, start=1):
            name, size, file = part.split(',')
            if not size:
                raise ValueError("Must provide a size")
            size = parse_size_str(size)

            sgdisk_command += [
                "--new",
                f"{i}:{part_start // GPT_BLOCK_SIZE}:{(part_start + size) // GPT_BLOCK_SIZE - 1}"
            ]
            sgdisk_command += ["--change-name", f"{i}:{name}"]

            partition_info.append((name, part_start, size, file))
            part_start += size

        res = subprocess.run(sgdisk_command, check=True, text=True)

        # Create the final disk with partition content
        dest_file = pathlib.Path(args.out)
        dest_file.write_bytes(0 * b'\x00')
        append_offset = 0
        for name, start, size, file in partition_info:
            end = start + size
            # Fill gap
            append_offset += _append(str(temp_disk), append_offset,
                                     start - append_offset, args.out)

            # Copy over file
            if file:
                file_size = os.path.getsize(file)
                if file_size > size:
                    raise ValueError(f"{file} too large for {name}")
                append_offset += _append(file, 0, file_size, args.out)

            # If partition size greater than file size, copy the remaining
            # partition content from `temp_disk` (zeroes)
            append_offset += _append(str(temp_disk), append_offset,
                                     end - append_offset, args.out)

        # Copy the remaining disk from `temp_disk` (zeroes + back up header/entries)
        append_offset += _append(str(temp_disk), append_offset,
                                 disk_size - append_offset, args.out)

    return 0


if __name__ == '__main__':
    sys.exit(main())