diff --git a/tests_ci/dummy_image.py b/tests_ci/dummy_image.py new file mode 100755 index 0000000..cfc71f2 --- /dev/null +++ b/tests_ci/dummy_image.py @@ -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() diff --git a/tests_ci/test_db.toml b/tests_ci/test_db.toml new file mode 100644 index 0000000..57a949a --- /dev/null +++ b/tests_ci/test_db.toml @@ -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"