Make Mac class prevent further updates or final after finalisation. Keep cached values for hashlib API.

tests/test_mac.py

@@ -7,6 +7,9 @@ from pyaegis import aegis128l, aegis128x2, aegis128x4, aegis256, aegis256x2, aeg
 
 from .util import random_split_bytes
 
+# All AEGIS algorithm modules
+ALL_ALGORITHMS = [aegis128l, aegis128x2, aegis128x4, aegis256, aegis256x2, aegis256x4]
+
 
 def load_mac_test_vectors():
     """Load MAC test vectors from JSON file."""
@@ -100,3 +103,167 @@ def test_mac_class(vector):
         assert computed_tag256 == expected_tag256, (
             f"256-bit MAC mismatch for {vector['name']}"
         )
+
+
+@pytest.mark.parametrize("vector", load_mac_test_vectors(), ids=get_test_id)
+def test_mac_class_with_digest(vector):
+    """Test MAC computation using digest() and hexdigest() instead of final()."""
+    alg = get_algorithm_module(vector["name"])
+
+    key = bytes.fromhex(vector["key"])
+    nonce = bytes.fromhex(vector["nonce"])
+    data = bytes.fromhex(vector["data"])
+
+    # Test 128-bit MAC if present
+    if "tag128" in vector:
+        expected_tag128 = bytes.fromhex(vector["tag128"])
+
+        # Test with digest()
+        mac_state = alg.Mac(key, nonce, maclen=16)
+        for chunk in random_split_bytes(data):
+            mac_state.update(chunk)
+        computed_tag128 = mac_state.digest()
+        assert computed_tag128 == expected_tag128, (
+            f"128-bit MAC mismatch for {vector['name']} using digest()"
+        )
+
+        # Test that digest() can be called multiple times
+        computed_tag128_again = mac_state.digest()
+        assert computed_tag128 == computed_tag128_again, (
+            "digest() should return the same value on repeated calls"
+        )
+
+        # Test hexdigest()
+        mac_state2 = alg.Mac(key, nonce, maclen=16)
+        for chunk in random_split_bytes(data):
+            mac_state2.update(chunk)
+        hex_tag = mac_state2.hexdigest()
+        assert hex_tag == expected_tag128.hex(), (
+            f"128-bit MAC hexdigest mismatch for {vector['name']}"
+        )
+
+        # Test that hexdigest() can be called multiple times
+        hex_tag_again = mac_state2.hexdigest()
+        assert hex_tag == hex_tag_again, (
+            "hexdigest() should return the same value on repeated calls"
+        )
+
+    # Test 256-bit MAC if present
+    if "tag256" in vector:
+        expected_tag256 = bytes.fromhex(vector["tag256"])
+
+        # Test with digest()
+        mac_state = alg.Mac(key, nonce, maclen=32)
+        for chunk in random_split_bytes(data):
+            mac_state.update(chunk)
+        computed_tag256 = mac_state.digest()
+        assert computed_tag256 == expected_tag256, (
+            f"256-bit MAC mismatch for {vector['name']} using digest()"
+        )
+
+
+@pytest.mark.parametrize("vector", load_mac_test_vectors(), ids=get_test_id)
+def test_mac_clone(vector):
+    """Test that cloning a Mac state works correctly."""
+    alg = get_algorithm_module(vector["name"])
+
+    key = bytes.fromhex(vector["key"])
+    nonce = bytes.fromhex(vector["nonce"])
+    data = bytes.fromhex(vector["data"])
+
+    # Test 128-bit MAC if present
+    if "tag128" in vector:
+        expected_tag128 = bytes.fromhex(vector["tag128"])
+
+        mac_state = alg.Mac(key, nonce, maclen=16)
+        for chunk in random_split_bytes(data):
+            mac_state.update(chunk)
+
+        # Clone the state
+        cloned_state = mac_state.clone()
+
+        # Both should produce the same tag
+        tag1 = mac_state.final()
+        tag2 = cloned_state.final()
+
+        assert tag1 == expected_tag128
+        assert tag2 == expected_tag128
+        assert tag1 == tag2
+
+
+@pytest.mark.parametrize("vector", load_mac_test_vectors(), ids=get_test_id)
+def test_mac_reset(vector):
+    """Test that resetting a Mac state works correctly."""
+    alg = get_algorithm_module(vector["name"])
+
+    key = bytes.fromhex(vector["key"])
+    nonce = bytes.fromhex(vector["nonce"])
+    data = bytes.fromhex(vector["data"])
+
+    # Test 128-bit MAC if present
+    if "tag128" in vector:
+        expected_tag128 = bytes.fromhex(vector["tag128"])
+
+        mac_state = alg.Mac(key, nonce, maclen=16)
+        for chunk in random_split_bytes(data):
+            mac_state.update(chunk)
+        tag1 = mac_state.final()
+        assert tag1 == expected_tag128
+
+        # Reset and compute again
+        mac_state.reset()
+        for chunk in random_split_bytes(data):
+            mac_state.update(chunk)
+        tag2 = mac_state.final()
+
+        assert tag2 == expected_tag128
+        assert tag1 == tag2
+
+
+@pytest.mark.parametrize("alg", ALL_ALGORITHMS, ids=lambda x: x.__name__.split(".")[-1])
+def test_mac_reset_after_digest(alg):
+    """Test that reset() clears the cached digest and allows reuse."""
+    key = alg.random_key()
+    nonce = alg.random_nonce()
+
+    mac_state = alg.Mac(key, nonce)
+    mac_state.update(b"Hello, world!")
+    tag1 = mac_state.digest()
+
+    # After digest(), update should fail
+    with pytest.raises(RuntimeError):
+        mac_state.update(b"More data")
+
+    # Reset should clear the cached digest
+    mac_state.reset()
+
+    # Now we should be able to update again
+    mac_state.update(b"Different data")
+    tag2 = mac_state.digest()
+
+    # Tags should be different since we used different data
+    assert tag1 != tag2
+
+
+@pytest.mark.parametrize("alg", ALL_ALGORITHMS, ids=lambda x: x.__name__.split(".")[-1])
+def test_mac_clone_preserves_cached_digest(alg):
+    """Test that cloning preserves the cached digest state."""
+    key = alg.random_key()
+    nonce = alg.random_nonce()
+
+    mac_state = alg.Mac(key, nonce)
+    mac_state.update(b"Hello, world!")
+    tag1 = mac_state.digest()
+
+    # Clone after digest
+    cloned_state = mac_state.clone()
+
+    # Both should return the same cached tag
+    tag2 = cloned_state.digest()
+    assert tag1 == tag2
+
+    # Both should be unable to update
+    with pytest.raises(RuntimeError):
+        mac_state.update(b"More data")
+    with pytest.raises(RuntimeError):
+        cloned_state.update(b"More data")