Tighten restrictions and errors for tag APIs

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)]

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)
 

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]

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)
+        )

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
 
 

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()

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()
 

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)

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()
 
 

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) == []