diff --git a/pygbl/ebl.py b/pygbl/ebl.py index f28734a..f5a09af 100644 --- a/pygbl/ebl.py +++ b/pygbl/ebl.py @@ -286,22 +286,16 @@ def validate(self) -> None: f"Image CRC-32 is invalid: expected {VALID_CRC32:#010x}, got {crc:#010x}" ) - def get_tags(self, tag_type: type[T]) -> list[T]: - return [t for t in self.tags if type(t) is tag_type] + def get_tags(self, tag_type: type[T], *, allow_missing: bool = False) -> list[T]: + tags: list[T] = [t for t in self.tags if type(t) is tag_type] - def find_first_tag(self, tag_type: type[T]) -> T | None: - return next((t for t in self.tags if type(t) is tag_type), None) - - def get_first_tag(self, tag_type: type[T]) -> T: - tag = self.find_first_tag(tag_type) - - if tag is None: + if not tags and not allow_missing: raise KeyError(f"No {tag_type.__name__} tag exists") - return tag + return tags - def has_tag(self, tag_type: type[Any]) -> bool: - return self.find_first_tag(tag_type) is not None + def get_first_tag(self, tag_type: type[T]) -> T: + return self.get_tags(tag_type)[0] def regenerate_crc(self) -> EBLImage: tags = [t for t in self.tags if not isinstance(t, EBLEnd)] diff --git a/pygbl/gbl3.py b/pygbl/gbl3.py index 1830974..5fc86b1 100644 --- a/pygbl/gbl3.py +++ b/pygbl/gbl3.py @@ -729,28 +729,22 @@ def validate(self) -> None: f"Image CRC-32 is invalid: expected {VALID_CRC32:#010x}, got {crc:#010x}" ) - def get_tags(self, tag_type: type[T]) -> list[T]: - return [t for t in self.tags if type(t) is tag_type] + def get_tags(self, tag_type: type[T], *, allow_missing: bool = False) -> list[T]: + tags: list[T] = [t for t in self.tags if type(t) is tag_type] - def find_first_tag(self, tag_type: type[T]) -> T | None: - return next((t for t in self.tags if type(t) is tag_type), None) - - def get_first_tag(self, tag_type: type[T]) -> T: - tag = self.find_first_tag(tag_type) - - if tag is None: + if not tags and not allow_missing: raise KeyError(f"No {tag_type.__name__} tag exists") - return tag + return tags - def has_tag(self, tag_type: type[Any]) -> bool: - return self.find_first_tag(tag_type) is not None + def get_first_tag(self, tag_type: type[T]) -> T: + return self.get_tags(tag_type)[0] def get_metadata(self) -> bytes | None: """The metadata tag contents, if the image has one. Opaque to the bootloader.""" - tag = self.find_first_tag(GBL3Metadata) + metadata = self.get_tags(GBL3Metadata, allow_missing=True) - return None if tag is None else tag.metadata + return metadata[0].metadata if metadata else None def regenerate_crc(self) -> GBL3Image: tags = [t for t in self.tags if not isinstance(t, GBL3End)] @@ -793,11 +787,13 @@ def sign(self, private_key: ec.EllipticCurvePrivateKey) -> GBL3Image: return type(self)(tags=[*tags, GBL3Signature(r=r, s=s)]).regenerate_crc() def verify_signature(self, public_key: ec.EllipticCurvePublicKey) -> bool: - signature = self.find_first_tag(GBL3Signature) + signatures = self.get_tags(GBL3Signature, allow_missing=True) - if signature is None: + if not signatures: raise KeyError("Image is not signed") + signature = signatures[0] + return verify_digest( public_key, self.signing_digest(), signature.r, signature.s ) @@ -848,11 +844,12 @@ def decrypt(self, key: bytes) -> GBL3Image: A signature covers the ciphertext, so it cannot carry over to the plaintext. """ - init = self.find_first_tag(GBL3EncryptionInitAesCcm) + inits = self.get_tags(GBL3EncryptionInitAesCcm, allow_missing=True) - if init is None: + if not inits: raise KeyError("Image is not encrypted") + init = inits[0] ciphertext = b"".join(t.ciphertext for t in self.get_tags(GBL3EncryptedData)) plaintext = aes_ctr_crypt(key, init.nonce, ciphertext) tags, _ = parse_tag_stream(plaintext, stop_at_end=False) @@ -901,11 +898,13 @@ def decompress(self) -> GBL3Image: return type(self)(tags=tags).regenerate_crc() def is_combined_bootloader_app(self) -> bool: - app_info = self.find_first_tag(GBL3ApplicationInfo) + app_infos = self.get_tags(GBL3ApplicationInfo, allow_missing=True) - if app_info is None or not self.has_tag(GBL3Bootloader): + if not app_infos or not self.get_tags(GBL3Bootloader, allow_missing=True): return False + app_info = app_infos[0] + return bool( GBL3ApplicationType.BOOTLOADER in app_info.type and app_info.type & ~GBL3ApplicationType.BOOTLOADER @@ -941,10 +940,13 @@ def split_bootloader_app(self) -> list[GBL3Image]: def combine_bootloader_app(self, other: GBL3Image) -> GBL3Image: """Combine a bootloader image and an application image into a single image.""" - if self.has_tag(GBL3Bootloader) and other.has_tag(GBL3Bootloader): + ours = self.get_tags(GBL3Bootloader, allow_missing=True) + theirs = other.get_tags(GBL3Bootloader, allow_missing=True) + + if ours and theirs: raise ValueError("Both images contain bootloaders") - if not self.has_tag(GBL3Bootloader) and not other.has_tag(GBL3Bootloader): + if not ours and not theirs: raise ValueError("Neither image contains a bootloader") header = self.get_first_tag(GBL3Header) @@ -955,7 +957,7 @@ def combine_bootloader_app(self, other: GBL3Image) -> GBL3Image: if header.type != GBL3Type.NONE: raise ValueError("Cannot combine signed or encrypted images") - bootloader_image = self if self.has_tag(GBL3Bootloader) else other + bootloader_image = self if ours else other application = other if bootloader_image is self else self app_info = application.get_first_tag(GBL3ApplicationInfo) diff --git a/pygbl/gbl4.py b/pygbl/gbl4.py index 1ce15aa..060c9ad 100644 --- a/pygbl/gbl4.py +++ b/pygbl/gbl4.py @@ -534,7 +534,7 @@ def from_bytes(cls, data: bytes) -> Self: def serialize(self) -> bytes: return serialize_tag(self.root) - def get_tags(self, tag_type: type[T]) -> list[T]: + def get_tags(self, tag_type: type[T], *, allow_missing: bool = False) -> list[T]: """Every tag of a type, at any depth.""" found: list[T] = [] @@ -548,20 +548,10 @@ def walk(tags: list[GBL4TagBase]) -> None: walk([self.root]) - return found - - def find_first_tag(self, tag_type: type[T]) -> T | None: - tags = self.get_tags(tag_type) - - return tags[0] if tags else None - - def get_first_tag(self, tag_type: type[T]) -> T: - tag = self.find_first_tag(tag_type) - - if tag is None: + if not found and not allow_missing: raise KeyError(f"No {tag_type.__name__} tag exists") - return tag + return found - def has_tag(self, tag_type: type[Any]) -> bool: - return self.find_first_tag(tag_type) is not None + def get_first_tag(self, tag_type: type[T]) -> T: + return self.get_tags(tag_type)[0] diff --git a/tests/test_commander_cli.py b/tests/test_commander_cli.py index 5926574..2fd6ea0 100644 --- a/tests/test_commander_cli.py +++ b/tests/test_commander_cli.py @@ -194,7 +194,9 @@ def test_real_images_respect_the_bootloader_limits(path: pathlib.Path) -> None: """Every LZMA tag in the wild fits the decoder's 10 KiB probs and 8 KiB dictionary.""" image = pygbl.parse_firmware_image(path.read_bytes()) - if not isinstance(image, GBL3Image) or not image.get_tags(GBL3ProgLZMA): + if not isinstance(image, GBL3Image) or not image.get_tags( + GBL3ProgLZMA, allow_missing=True + ): pytest.skip("no LZMA tags") for tag in image.get_tags(GBL3ProgLZMA): @@ -254,8 +256,12 @@ def test_lzma_props_match(case: tuple[str, GBL3Compression | None]) -> None: """The props bytes are what the device reads to size its decoder buffers.""" ours, theirs = build(*case) - assert [t.data[:LZMA_PROPS_SIZE] for t in ours.get_tags(GBL3ProgLZMA)] == [ - t.data[:LZMA_PROPS_SIZE] for t in theirs.get_tags(GBL3ProgLZMA) + assert [ + t.data[:LZMA_PROPS_SIZE] + for t in ours.get_tags(GBL3ProgLZMA, allow_missing=True) + ] == [ + t.data[:LZMA_PROPS_SIZE] + for t in theirs.get_tags(GBL3ProgLZMA, allow_missing=True) ] @@ -267,4 +273,7 @@ def test_both_encoders_write_lzma_end_markers( ours, theirs = build(*case) for image in (ours, theirs): - assert all(ends_with_end_marker(t.data) for t in image.get_tags(GBL3ProgLZMA)) + assert all( + ends_with_end_marker(t.data) + for t in image.get_tags(GBL3ProgLZMA, allow_missing=True) + ) diff --git a/tests/test_compression.py b/tests/test_compression.py index b6c58a3..dbee1ab 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -137,7 +137,7 @@ def test_compress_replaces_program_data_tags( compressed = build_image().compress(algorithm) assert len(compressed.get_tags(tag_type)) == 2 - assert compressed.get_tags(GBL3EraseProg) == [] + assert compressed.get_tags(GBL3EraseProg, allow_missing=True) == [] assert [t.address for t in compressed.get_tags(tag_type)] == [0x8006000, 0x8106000] @@ -177,7 +177,7 @@ def test_compress_handles_program_data1() -> None: compressed = variant.compress(GBL3Compression.LZMA) - assert compressed.get_tags(GBL3Prog) == [] + assert compressed.get_tags(GBL3Prog, allow_missing=True) == [] assert len(compressed.get_tags(GBL3ProgLZMA)) == 2 diff --git a/tests/test_crypto.py b/tests/test_crypto.py index 84a6573..2305bdc 100644 --- a/tests/test_crypto.py +++ b/tests/test_crypto.py @@ -256,6 +256,6 @@ def test_decrypt_drops_signature_and_flags() -> None: sealed = image.encrypt(KEY, nonce=NONCE).sign(PRIVATE_KEY) decrypted = sealed.decrypt(KEY) - assert decrypted.find_first_tag(GBL3Signature) is None + assert decrypted.get_tags(GBL3Signature, allow_missing=True) == [] assert decrypted.get_first_tag(GBL3Header).type == GBL3Type.NONE assert decrypted.serialize() == image.serialize() diff --git a/tests/test_elf.py b/tests/test_elf.py index 39be04c..eface80 100644 --- a/tests/test_elf.py +++ b/tests/test_elf.py @@ -86,7 +86,7 @@ def rebuild(name: str) -> GBL3Image: theirs = commander_gbl(name) builder = ( build_bootloader_gbl3 - if theirs.has_tag(GBL3Bootloader) + if theirs.get_tags(GBL3Bootloader, allow_missing=True) else build_application_gbl3 ) @@ -205,7 +205,7 @@ def test_application_gbl_without_metadata() -> None: with elf_path(APPLICATION).open("rb") as f: image = build_application_gbl3(f) - assert image.find_first_tag(GBL3Metadata) is None + assert image.get_tags(GBL3Metadata, allow_missing=True) == [] assert image.get_first_tag(GBL3Header).type == GBL3Type.NONE image.validate() diff --git a/tests/test_gbl.py b/tests/test_gbl.py index 5e502d7..b293bc5 100644 --- a/tests/test_gbl.py +++ b/tests/test_gbl.py @@ -219,9 +219,10 @@ def test_tag_lookup_helpers() -> None: assert image.get_tags(GBL3EraseProg) == [first, second] assert image.get_first_tag(GBL3EraseProg) == first - assert image.find_first_tag(GBL3Metadata) is None - assert image.has_tag(GBL3EraseProg) - assert not image.has_tag(GBL3Metadata) + assert image.get_tags(GBL3Metadata, allow_missing=True) == [] + + with pytest.raises(KeyError, match="No GBL3Metadata tag exists"): + image.get_tags(GBL3Metadata) with pytest.raises(KeyError, match="No GBL3Metadata tag exists"): image.get_first_tag(GBL3Metadata) diff --git a/tests/test_gbl4.py b/tests/test_gbl4.py index 03e7a77..ed37a0a 100644 --- a/tests/test_gbl4.py +++ b/tests/test_gbl4.py @@ -63,7 +63,7 @@ def test_dispatches_to_gbl4(path: pathlib.Path) -> None: @pytest.mark.parametrize("path", IMAGES, ids=lambda p: p.name) def test_no_unknown_tags(path: pathlib.Path) -> None: """Every tag in a real image should be one we model.""" - assert load(path).get_tags(GBL4UnknownTag) == [] + assert load(path).get_tags(GBL4UnknownTag, allow_missing=True) == [] @pytest.mark.parametrize("path", IMAGES, ids=lambda p: p.name) @@ -130,7 +130,7 @@ def test_hash_values_are_variable_length(path: pathlib.Path) -> None: @pytest.mark.parametrize("path", IMAGES, ids=lambda p: p.name) def test_padding_is_preserved(path: pathlib.Path) -> None: """`PAD` has the container bit set but holds raw bytes, so it must stay opaque.""" - for pad in load(path).get_tags(GBL4Pad): + for pad in load(path).get_tags(GBL4Pad, allow_missing=True): assert pad.data == pad.serialize_payload() diff --git a/tests/test_rare_tags.py b/tests/test_rare_tags.py index bf48272..1c8ca68 100644 --- a/tests/test_rare_tags.py +++ b/tests/test_rare_tags.py @@ -171,4 +171,4 @@ def test_prog_survives_a_roundtrip() -> None: assert reparsed == image assert isinstance(reparsed.get_first_tag(GBL3Prog), GBL3Prog) - assert reparsed.find_first_tag(GBL3EraseProg) is None + assert reparsed.get_tags(GBL3EraseProg, allow_missing=True) == []