1#!/usr/bin/env python3
2#
3# Copyright (C) 2023 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16"""Generate a GPT disk
17
18Example:
19
20The following command creates a 16mb disk file gpt.bin with two partitions:
21
22  1) "boot_a" of size 4096KB, with data initialized from "file_a",
23  2) "boot_b" of size 8192KB, with data initialized to zeroes.
24
25    gen_gpt_disk.py ./gpt.bin 16M
26        --partition "boot_a,4096K,<path to file_a>" \
27        --partition "boot_b,8192K,"
28
29All GUIDs will be non-deterministic for reproducibility.
30"""
31
32import argparse
33import os
34import pathlib
35import shutil
36import sys
37import subprocess
38import tempfile
39
40GPT_BLOCK_SIZE = 512
41
42# Some GPT libraries may expect a valid GUID here, these are just pre-generated
43# valid GUIDs for reproducibility.
44DISK_GUID = "b116114b-b6ae-4186-8231-b35104cb4c9e"
45PARTITION_GUIDS = [
46    "36bc51fd-c3d6-4109-a2ac-35bddd757e2a",
47    "42aaac2e-37e3-43ba-9930-42dfa96e6334",
48    "bdadfeca-879c-43e9-8f0d-8ef7da29b5e7",
49    "8d5b90b2-0d65-476b-8691-94f74fcac271",
50    "dafe682f-3f2b-4472-a1b8-a0f217701909",
51    "7097fd78-f559-461e-bb9b-db176a8169d8",
52    "c03dd176-117b-4e65-8205-37ebec007c1a",
53    "6d9159db-9b9e-4906-8fa4-31f5ffaef50e",
54    "4cdfda9f-23aa-4b27-9fea-24bb08238135",
55    "2a0a3df9-e8ef-4627-b4ce-ef1e847924f4",
56    "3c9b64f9-c659-4e5d-9d25-72028c46e6b8",
57    "95f142f9-d1f3-41ad-96dc-b969ee8242a1",
58    "5181811b-e836-4e66-bfe2-91aa86da515f",
59    "f2dbad77-38ac-43de-b4c7-2f7432d2339e",
60    "fc172d2c-f735-49a5-8be0-a475d0fc5be9",
61    "5ad48195-8e26-4496-a7e2-669028db2dce",
62    "f49a6965-8168-4c0c-8102-7b64a8176c25",
63    "be495262-5d9b-4fcb-9240-d9e84b195abe",
64    "c5ab3a8d-f898-420f-9039-a7445760fb0f",
65    "bc79912b-44bf-4ed8-8320-796ba47714d1",
66]
67
68
69def parse_args():
70    parser = argparse.ArgumentParser()
71    parser.add_argument("out", help="output file")
72    parser.add_argument(
73        "disk_size", type=str, help="disk size of the image. i.e. 64k, 1M etc"
74    )
75    parser.add_argument(
76        "--partition",
77        type=str,
78        action="append",
79        help="specifies a partition. Format should be"
80        "--partition=<part name>,<size>,<file name>",
81    )
82    return parser.parse_args()
83
84
85def parse_size_str(size_str: str) -> int:
86    """Parse size string into integer, taking into account unit.
87
88    Example:
89        2M, 2m -> 2*1024*1024
90        2K, 2k -> 2*1024
91        512 -> 512
92    """
93    size_str = size_str.lower()
94    if size_str.endswith("m"):
95        return int(size_str[:-1]) * 1024 * 1024
96    if size_str.endswith("k"):
97        return int(size_str[:-1]) * 1024
98    return int(size_str)
99
100
101def _append(src_file: str, offset: int, size: int, dst_file: str):
102    """Append a portion of a file to the end of another file
103
104    Args:
105        src_file: path to the source file
106        offset: offset in the source file
107        size: number of bytes to append
108        dst_file: destination file
109
110    Returns:
111        number of bytes actually copied.
112    """
113    with open(src_file, "rb") as src_f:
114        src_f.seek(offset, 0)
115        data = src_f.read(size)
116        if len(data) != size:
117            raise ValueError(f"Want {size} bytes from {src_file}, but got {len(data)}.")
118        with open(dst_file, "ab") as dst_f:
119            dst_f.write(data)
120    return size
121
122
123def main() -> int:
124    args = parse_args()
125    with tempfile.TemporaryDirectory() as temp_dir:
126        temp_disk = pathlib.Path(temp_dir) / "temp_disk"
127
128        disk_size = parse_size_str(args.disk_size)
129        temp_disk.write_bytes(disk_size * b"\x00")
130
131        part_start = 34 * GPT_BLOCK_SIZE  # MBR + GPT header + entries
132        partition_info = []
133        sgdisk_command = [
134            "sgdisk",
135            str(temp_disk),
136            "--clear",
137            "--set-alignment",
138            "1",
139            "--disk-guid",
140            DISK_GUID,
141        ]
142
143        for i, part in enumerate(args.partition, start=1):
144            name, size, file = part.split(",")
145            if not size:
146                raise ValueError("Must provide a size")
147            size = parse_size_str(size)
148
149            sgdisk_command += [
150                "--new",
151                f"{i}:{part_start // GPT_BLOCK_SIZE}:{(part_start + size) // GPT_BLOCK_SIZE - 1}",
152                "--partition-guid",
153                f"{i}:{PARTITION_GUIDS[i]}",
154                "--change-name",
155                f"{i}:{name}",
156            ]
157
158            partition_info.append((name, part_start, size, file))
159            part_start += size
160
161        res = subprocess.run(sgdisk_command, check=True, text=True)
162
163        # Create the final disk with partition content
164        dest_file = pathlib.Path(args.out)
165        dest_file.write_bytes(0 * b"\x00")
166        append_offset = 0
167        for name, start, size, file in partition_info:
168            end = start + size
169            # Fill gap
170            append_offset += _append(
171                str(temp_disk), append_offset, start - append_offset, args.out
172            )
173
174            # Copy over file
175            if file:
176                file_size = os.path.getsize(file)
177                if file_size > size:
178                    raise ValueError(f"{file} too large for {name}")
179                append_offset += _append(file, 0, file_size, args.out)
180
181            # If partition size greater than file size, copy the remaining
182            # partition content from `temp_disk` (zeroes)
183            append_offset += _append(
184                str(temp_disk), append_offset, end - append_offset, args.out
185            )
186
187        # Copy the remaining disk from `temp_disk` (zeroes + back up header/entries)
188        append_offset += _append(
189            str(temp_disk), append_offset, disk_size - append_offset, args.out
190        )
191
192    return 0
193
194
195if __name__ == "__main__":
196    sys.exit(main())
197