diff --git a/tests_ci/dummy_image.py b/tests_ci/dummy_image.py index 03c026b..4fccd9f 100755 --- a/tests_ci/dummy_image.py +++ b/tests_ci/dummy_image.py @@ -19,6 +19,21 @@ import ota_utils PARTITIONS_TO_PRESERVE = set(sum(PARTITION_PRIORITIES.values(), ())) +def strtobool(val): + """Convert a string representation of truth to true (1) or false (0). + True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values + are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if + 'val' is anything else. + """ + val = val.lower() + if val in ("y", "yes", "t", "true", "on", "1"): + return 1 + elif val in ("n", "no", "f", "false", "off", "0"): + return 0 + else: + raise ValueError("invalid truth value %r" % (val,)) + + class Crc32Hasher: existing_crc32 = None @@ -115,33 +130,53 @@ def download_dummy_image(args): previous_offset = section["range"][1] + 1 -def print_entry(db, original_filename): - new_map = tomlkit.array() - new_map.append(tomlkit.nl()) +def add_entry(db, db_entry, device_name, original_filename): + new_sections = tomlkit.array() + new_sections.append(tomlkit.nl()) - for section in db: + for section in db_entry: entry = tomlkit.inline_table() for key, value in section.items(): entry[key] = f"{value:x}" if key == "checksum" else value - new_map.append(entry) + new_sections.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 ***") + device_entry = { + device_name: { + "filename": os.path.basename(original_filename), + "url": "", + "sections": new_sections, + } + } + + with open(db, "r") as f: + toml_db = tomlkit.load(f) + + if toml_db.get("device").get(device_name) is not None: + overwrite_entry = None + + while overwrite_entry is None: + choice = input( + f"{device_name} is already in the database. Do you want to overwrite it? [y/n] " + ).lower() + + try: + overwrite_entry = strtobool(choice) + except ValueError: + pass + + if not overwrite_entry: + print("Entry not added to the database. It is printed below instead:") + print(tomlkit.dumps({"device": device_entry})) + return + + toml_db.get("device").update(device_entry) + + with open(db, "w") as f: + tomlkit.dump(toml_db, f) + + print(f"{device_name} is added to the database") def convert_image(args): @@ -167,13 +202,13 @@ def convert_image(args): partition_operations[p.partition_name] = p.operations - db = [] + db_entry = [] 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) + output_writer(f_in, f_out, file_offset + blob_offset, db_entry) # Sort partitions based on data_offset of first operation and iterate for _, p_operations in sorted( @@ -184,22 +219,22 @@ def convert_image(args): f_in, f_out, p_operations[0].data_offset - (f_in.tell() - file_offset - blob_offset), - db, + db_entry, ) 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) + output_writer(None, f_out, op.data_length, db_entry) f_in.seek(f_out.tell()) # Copy remaining data f_in.seek(0, os.SEEK_END) remaining_size = f_in.tell() - f_out.tell() f_in.seek(f_out.tell()) - output_writer(f_in, f_out, remaining_size, db) + output_writer(f_in, f_out, remaining_size, db_entry) - print_entry(db, args.input) + add_entry(args.db, db_entry, args.device_name, args.input) def parse_args(args=None): @@ -215,18 +250,18 @@ def parse_args(args=None): default=True, help="Compress output file", ) + base.add_argument("--db", required=True, help="Path to database file") base.add_argument("--output", required=True, help="Path to new OTA zip") + base.add_argument("device_name", help="Name of device (as in database)") 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( + 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)