mirror of
https://github.com/chenxiaolong/avbroot.git
synced 2026-07-03 14:05:11 +02:00
Add creation of dummy images for use in tests
This commit is contained in:
Pascal Roeleven
Executable
+247
@@ -0,0 +1,247 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
from binascii import crc32
|
||||
import contextlib
|
||||
import gzip
|
||||
import os.path
|
||||
import sys
|
||||
import urllib.request
|
||||
import zipfile
|
||||
|
||||
import tomlkit
|
||||
|
||||
sys.path.append(os.path.join(sys.path[0], ".."))
|
||||
from avbroot import ota, util
|
||||
from avbroot.main import PATH_PAYLOAD
|
||||
|
||||
ZIP_HEADER_NO_FILENAME_SIZE = 30
|
||||
|
||||
PARTITIONS_TO_ZERO = {"product", "system", "vendor", "system_ext", "modem"}
|
||||
|
||||
|
||||
class Crc32Hasher:
|
||||
existing_crc32 = None
|
||||
|
||||
def get_checksum(self):
|
||||
return self.existing_crc32
|
||||
|
||||
def set_checksum(self, existing_crc32):
|
||||
self.existing_crc32 = existing_crc32
|
||||
|
||||
def update(self, *args, **kwargs):
|
||||
if self.existing_crc32:
|
||||
args = args + (self.existing_crc32,)
|
||||
|
||||
self.existing_crc32 = crc32(*args, **kwargs)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def output_open(filename, mode, compress=True):
|
||||
if compress and not filename.endswith(".gz"):
|
||||
filename += ".gz"
|
||||
|
||||
if compress:
|
||||
with gzip.GzipFile(filename, mode, compresslevel=1, mtime=0) as f:
|
||||
yield f
|
||||
else:
|
||||
with open(filename, mode) as f:
|
||||
yield f
|
||||
|
||||
|
||||
def output_writer(f_in, f_out, size, db=None):
|
||||
if not f_in:
|
||||
util.zero_n(f_out, size)
|
||||
else:
|
||||
crc_hash = None
|
||||
mapping = None
|
||||
|
||||
if db is not None and size != 0:
|
||||
mapping = True
|
||||
|
||||
if mapping:
|
||||
crc_hash = Crc32Hasher()
|
||||
|
||||
# Combine sections if adjacent
|
||||
if db and db[-1]["range"][1] == f_out.tell() - 1:
|
||||
entry = db.pop()
|
||||
crc_hash.set_checksum(entry["checksum"])
|
||||
section_start = entry["range"][0]
|
||||
else:
|
||||
section_start = f_out.tell()
|
||||
|
||||
section_end = f_out.tell() + (size - 1)
|
||||
|
||||
util.copyfileobj_n(f_in, f_out, size, hasher=crc_hash)
|
||||
|
||||
if mapping:
|
||||
db.append(
|
||||
{
|
||||
"range": [section_start, section_end],
|
||||
"checksum": crc_hash.get_checksum(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def download_file(f_out, url, section):
|
||||
crc_hash = Crc32Hasher()
|
||||
|
||||
req = urllib.request.Request(url)
|
||||
req.add_header("Range", f"bytes={section['range'][0]}-{section['range'][1]}")
|
||||
with urllib.request.urlopen(req) as data:
|
||||
util.copyfileobj_n(
|
||||
data, f_out, section["range"][1] - section["range"][0] + 1, hasher=crc_hash
|
||||
)
|
||||
|
||||
if crc_hash.get_checksum() != int(section["checksum"], 16):
|
||||
raise Exception(f"ERROR: Checksum of range {section['range']} doesn't match")
|
||||
|
||||
|
||||
def download_dummy_image(args):
|
||||
with open(args.db, "r") as f:
|
||||
database = f.read()
|
||||
|
||||
database = tomlkit.parse(database)
|
||||
|
||||
device_entry = database["device"].get(args.device_name)
|
||||
if not device_entry:
|
||||
raise Exception(f"ERROR: {args.device_name} cannot be found in the database")
|
||||
|
||||
# Assume we always download the beginning and end (which makes sense for a ZIP file)
|
||||
previous_offset = 0
|
||||
with output_open(args.output, "wb", args.compress) as f_out:
|
||||
for section in device_entry["sections"]:
|
||||
util.zero_n(f_out, section["range"][0] - previous_offset)
|
||||
download_file(f_out, device_entry["url"], section)
|
||||
previous_offset = section["range"][1] + 1
|
||||
|
||||
|
||||
def print_entry(db, original_filename):
|
||||
new_map = tomlkit.array()
|
||||
new_map.append(tomlkit.nl())
|
||||
|
||||
for section in db:
|
||||
entry = tomlkit.inline_table()
|
||||
|
||||
for key, value in section.items():
|
||||
entry[key] = f"{value:x}" if key == "checksum" else value
|
||||
|
||||
new_map.append(entry)
|
||||
|
||||
print("*** Add this entry to the database ***")
|
||||
print(
|
||||
tomlkit.dumps(
|
||||
{
|
||||
"device": {
|
||||
"new_device_name": {
|
||||
"filename": os.path.basename(original_filename),
|
||||
"url": "",
|
||||
"sections": new_map,
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
print("*** End ***")
|
||||
|
||||
|
||||
def convert_image(args):
|
||||
with zipfile.ZipFile(args.input, "r") as z:
|
||||
info = z.getinfo(PATH_PAYLOAD)
|
||||
|
||||
with z.open(info, "r") as f:
|
||||
_, manifest, blob_offset = ota.parse_payload(f)
|
||||
|
||||
file_offset = info.header_offset + ZIP_HEADER_NO_FILENAME_SIZE + len(info.filename)
|
||||
|
||||
partition_operations = {}
|
||||
|
||||
for p in manifest.partitions:
|
||||
if p.partition_name in PARTITIONS_TO_ZERO:
|
||||
old_data_offset = 0
|
||||
|
||||
# Ensure all operations are in order
|
||||
for op in p.operations:
|
||||
if old_data_offset > op.data_offset:
|
||||
raise Exception("Operations are expected to be ordered")
|
||||
old_data_offset = op.data_offset
|
||||
|
||||
partition_operations[p.partition_name] = p.operations
|
||||
|
||||
db = []
|
||||
|
||||
with open(args.input, "rb") as f_in, output_open(
|
||||
args.output, "wb", args.compress
|
||||
) as f_out:
|
||||
# Copy zip data until payload.bin + payload.bin header
|
||||
output_writer(f_in, f_out, file_offset + blob_offset, db)
|
||||
|
||||
# Sort partitions based on data_offset of first operation and iterate
|
||||
for _, p_operations in sorted(
|
||||
partition_operations.items(), key=lambda x: x[1][0].data_offset
|
||||
):
|
||||
# Copy all data until first partition offset
|
||||
output_writer(
|
||||
f_in,
|
||||
f_out,
|
||||
p_operations[0].data_offset - (f_in.tell() - file_offset - blob_offset),
|
||||
db,
|
||||
)
|
||||
|
||||
for op in p_operations:
|
||||
if f_out.tell() != (file_offset + blob_offset + op.data_offset):
|
||||
raise Exception("f_out should be equal to the offset")
|
||||
output_writer(None, f_out, op.data_length, db)
|
||||
f_in.seek(f_out.tell())
|
||||
|
||||
# Copy remaining data
|
||||
f_in.seek(0, 2)
|
||||
remaining_size = f_in.tell() - f_out.tell()
|
||||
f_in.seek(f_out.tell())
|
||||
output_writer(f_in, f_out, remaining_size, db)
|
||||
|
||||
print_entry(db, args.input)
|
||||
|
||||
|
||||
def parse_args(args=None):
|
||||
parser = argparse.ArgumentParser()
|
||||
subparsers = parser.add_subparsers(
|
||||
dest="subcommand", required=True, help="Subcommands"
|
||||
)
|
||||
|
||||
base = argparse.ArgumentParser(add_help=False)
|
||||
base.add_argument(
|
||||
"--compress",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Compress output file",
|
||||
)
|
||||
base.add_argument("--output", required=True, help="Path to new OTA zip")
|
||||
|
||||
convert = subparsers.add_parser(
|
||||
"convert", help="Convert full OTA zip to dummy zip", parents=[base]
|
||||
)
|
||||
convert.add_argument("--input", required=True, help="Path to original OTA zip")
|
||||
|
||||
download = subparsers.add_parser(
|
||||
"download", help="Assemble dummy zip while downloading", parents=[base]
|
||||
)
|
||||
download.add_argument("--db", required=True, help="Path to database file")
|
||||
download.add_argument("device_name", help="Name of device (as in database)")
|
||||
|
||||
return parser.parse_args(args=args)
|
||||
|
||||
|
||||
def main(args=None):
|
||||
args = parse_args(args=args)
|
||||
|
||||
if args.subcommand == "convert":
|
||||
convert_image(args)
|
||||
elif args.subcommand == "download":
|
||||
download_dummy_image(args)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,65 @@
|
||||
[magisk]
|
||||
filename = "magisk.apk"
|
||||
url = "https://github.com/topjohnwu/Magisk/releases/download/v25.2/Magisk-v25.2.apk"
|
||||
sha256 = "0bdc32918b6ea502dca769b1c7089200da51ea1def170824c2812925b426d509"
|
||||
|
||||
[device.bluejay]
|
||||
filename = "bluejay-ota-tq1a.230205.002-1ec64b58.zip"
|
||||
url = "https://dl.google.com/dl/android/aosp/bluejay-ota-tq1a.230205.002-1ec64b58.zip"
|
||||
sections = [
|
||||
{range = [0, 25781534], checksum = "2db28a01"},
|
||||
{range = [50595563, 50613918], checksum = "6f0635b4"},
|
||||
{range = [1870236118, 1871590329], checksum = "272a03c1"},
|
||||
{range = [2051986898, 2081231853], checksum = "39940962"}
|
||||
]
|
||||
|
||||
[device.bluejay.hashes]
|
||||
"boot.img" = "8131084ad95d8ffcb917bc239ffdeef6"
|
||||
"vbmeta.img" = "38bc2245f6d41877ebf92d6612d3ca04"
|
||||
"vendor_boot.img" = "4611f85f4eae65e0eca09605068d14c8"
|
||||
|
||||
[device.bramble]
|
||||
filename = "bramble-ota-tq1a.230205.002-49f2dc71.zip"
|
||||
url = "https://dl.google.com/dl/android/aosp/bramble-ota-tq1a.230205.002-49f2dc71.zip"
|
||||
sections = [
|
||||
{range = [0, 12759576], checksum = "b0a1b579"},
|
||||
{range = [1155394531, 1155409742], checksum = "b018de12"},
|
||||
{range = [1636569527, 1637452126], checksum = "b8747c8"},
|
||||
{range = [1870498525, 1896655733], checksum = "f3d4f27b"}
|
||||
]
|
||||
|
||||
[device.bramble.hashes]
|
||||
"boot.img" = "098500dfba880f9a61a379d342a84cc9"
|
||||
"vbmeta.img" = "17230bbb41023bd3b03aaf94203e9330"
|
||||
"vendor_boot.img" = "34e3559f2460d8ef52ccb5352e09dfec"
|
||||
|
||||
[device.cheetah]
|
||||
filename = "cheetah-ota-tq1a.230205.002-ad27a6d7.zip"
|
||||
url = "https://dl.google.com/dl/android/aosp/cheetah-ota-tq1a.230205.002-ad27a6d7.zip"
|
||||
sections = [
|
||||
{range = [0, 23399500], checksum = "34f19830"},
|
||||
{range = [125881702, 125898553], checksum = "9f68fc6d"},
|
||||
{range = [1518207067, 1518293650], checksum = "a003f2d1"},
|
||||
{range = [1883282667, 1883297278], checksum = "9537e70d"},
|
||||
{range = [2010705366, 2012070705], checksum = "94fdddd2"},
|
||||
{range = [2278970391, 2307353441], checksum = "6ba7d418"}
|
||||
]
|
||||
|
||||
[device.cheetah.hashes]
|
||||
"init_boot.img" = "7c200e0dfd1af446d68a3f009cc86f8b"
|
||||
"vbmeta.img" = "b12f7e00cc7b9152f7aaa00ab2c4bae0"
|
||||
"vendor_boot.img" = "e39682b1656b8c59ac2717f330a08bd9"
|
||||
|
||||
[device.sunfish]
|
||||
filename = "sunfish-ota-tq1a.230205.002-536bce5f.zip"
|
||||
url = "https://dl.google.com/dl/android/aosp/sunfish-ota-tq1a.230205.002-536bce5f.zip"
|
||||
sections = [
|
||||
{range = [0, 34269195], checksum = "93177bf6"},
|
||||
{range = [1146499715, 1146513902], checksum = "83903c62"},
|
||||
{range = [1610565619, 1611042042], checksum = "114b64ce"},
|
||||
{range = [1807072291, 1809631036], checksum = "4d33dae2"}
|
||||
]
|
||||
|
||||
[device.sunfish.hashes]
|
||||
"boot.img" = "ddd9bcf0bdc4c727f33110f90103158a"
|
||||
"vbmeta.img" = "7388afabde4bb490ddef491dd2a8fe27"
|
||||
Reference in New Issue
Block a user