1#!/usr/bin/env python3
2#
3# Copyright (C) 2024 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 test data files for libgbl tests"""
17
18import os
19import pathlib
20import random
21import shutil
22import subprocess
23import tempfile
24from typing import List
25
26SCRIPT_DIR = pathlib.Path(os.path.dirname(os.path.realpath(__file__)))
27GPT_TOOL = pathlib.Path(SCRIPT_DIR.parents[1]) / "tools" / "gen_gpt_disk.py"
28AVB_DIR = pathlib.Path(SCRIPT_DIR.parents[4]) / "external" / "avb"
29AVB_TOOL = AVB_DIR / "avbtool.py"
30AVB_TEST_DATA_DIR = AVB_DIR / "test" / "data"
31SZ_KB = 1024
32
33# RNG seed values. Keep the same seed value for a given file to ensure
34# reproducibility as much as possible; this will prevent adding a bunch of
35# unnecessary test binaries to the git history.
36RNG_SEED_SPARSE_TEST_RAW = 1
37RNG_SEED_ZIRCON = {"a": 2, "b": 3, "r": 4}
38
39
40# A helper for writing bytes to a file at a given offset.
41def write_file(file, offset, data):
42    file.seek(offset, 0)
43    file.write(data)
44
45
46# Generates sparse image for flashing test
47def gen_sparse_test_file():
48    out_file_raw = SCRIPT_DIR / "sparse_test_raw.bin"
49    random.seed(RNG_SEED_SPARSE_TEST_RAW)
50    with open(out_file_raw, "wb") as f:
51        # 4k filled with 0x78563412
52        write_file(f, 0, b"\x12\x34\x56\x78" * 1024)
53        # 8k file hole (will become dont-care with the "-s" option)
54        # 12k raw data
55        write_file(f, 12 * SZ_KB, random.randbytes(12 * SZ_KB))
56        # 8k filled with 0x78563412
57        write_file(f, 24 * SZ_KB, b"\x12\x34\x56\x78" * 1024 * 2)
58        # 12k raw data
59        write_file(f, 32 * SZ_KB, random.randbytes(12 * SZ_KB))
60        # 4k filled with 0x78563412
61        write_file(f, 44 * SZ_KB, b"\x12\x34\x56\x78" * 1024)
62        # 8k filled with 0xEFCDAB90
63        write_file(f, 48 * SZ_KB, b"\x90\xab\xcd\xef" * 1024 * 2)
64
65    # For now this requires that img2simg exists on $PATH.
66    # It can be built from an Android checkout via `m img2simg`; the resulting
67    # binary will be at out/host/linux-x86/bin/img2simg.
68    subprocess.run(["img2simg", "-s", out_file_raw, SCRIPT_DIR / "sparse_test.bin"])
69    subprocess.run(
70        [
71            "img2simg",
72            "-s",
73            out_file_raw,
74            SCRIPT_DIR / "sparse_test_blk1024.bin",
75            "1024",
76        ]
77    )
78
79
80# Generates GPT disk, kernel data for Zircon tests
81def gen_zircon_gpt():
82    gen_gpt_args = []
83    for suffix in ["a", "b", "r"]:
84        random.seed(RNG_SEED_ZIRCON[suffix])
85        zircon = random.randbytes(16 * SZ_KB)
86        out_file = SCRIPT_DIR / f"zircon_{suffix}.bin"
87        out_file.write_bytes(zircon)
88        gen_gpt_args.append(f"--partition=zircon_{suffix},16K,{str(out_file)}")
89
90    subprocess.run(
91        [GPT_TOOL, SCRIPT_DIR / "zircon_gpt.bin", "128K"] + gen_gpt_args, check=True
92    )
93
94
95# Generates test data for A/B slot Manager writeback test
96def gen_writeback_test_bin():
97    subprocess.run(
98        [
99            GPT_TOOL,
100            SCRIPT_DIR / "writeback_test_disk.bin",
101            "64K",
102            "--partition=test_partition,4k,/dev/zero",
103        ],
104        check=True,
105    )
106
107
108def gen_vbmeta():
109    """Creates the vbmeta keys and signs some images."""
110    # Use the test vbmeta keys from libavb.
111    for name in ["testkey_rsa4096.pem", "testkey_rsa4096_pub.pem"]:
112        shutil.copyfile(AVB_TEST_DATA_DIR / name, SCRIPT_DIR / name)
113
114    # Convert the public key to raw bytes for use in verification.
115    subprocess.run(
116        [
117            AVB_TOOL,
118            "extract_public_key",
119            "--key",
120            SCRIPT_DIR / "testkey_rsa4096_pub.pem",
121            "--output",
122            SCRIPT_DIR / "testkey_rsa4096_pub.bin",
123        ],
124        check=True,
125    )
126
127    with tempfile.TemporaryDirectory() as temp_dir:
128        temp_dir = pathlib.Path(temp_dir)
129
130        # Create the hash descriptor. We only need this temporarily until we add
131        # it into the final vbmeta image.
132        hash_descriptor_path = temp_dir / "hash_descriptor.bin"
133        subprocess.run(
134            [
135                AVB_TOOL,
136                "add_hash_footer",
137                "--dynamic_partition_size",
138                "--do_not_append_vbmeta_image",
139                "--partition_name",
140                "zircon_a",
141                "--image",
142                SCRIPT_DIR / "zircon_a.bin",
143                "--output_vbmeta_image",
144                hash_descriptor_path,
145                "--salt",
146                "2000",
147            ],
148            check=True,
149        )
150
151        # Create the final signed vbmeta including the hash descriptor.
152        subprocess.run(
153            [
154                AVB_TOOL,
155                "make_vbmeta_image",
156                "--key",
157                SCRIPT_DIR / "testkey_rsa4096.pem",
158                "--algorithm",
159                "SHA512_RSA4096",
160                "--include_descriptors_from_image",
161                hash_descriptor_path,
162                "--output",
163                SCRIPT_DIR / "zircon_a.vbmeta",
164            ],
165            check=True,
166        )
167
168
169if __name__ == "__main__":
170    gen_writeback_test_bin()
171    gen_sparse_test_file()
172    gen_zircon_gpt()
173    gen_vbmeta()
174