fully working timestamping enricher

src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py

@@ -2,8 +2,10 @@ import os
 from loguru import logger
 
 from importlib.metadata import version
+import hashlib
 
 import requests
+
 from rfc3161_client import (
     TimestampRequestBuilder,
     TimeStampResponse,
@@ -20,6 +22,8 @@ from auto_archiver.core import Enricher
 from auto_archiver.core import Metadata, Media
 from auto_archiver.version import __version__
 
+
+
 class TimestampingEnricher(Enricher):
     """
     Uses several RFC3161 Time Stamp Authorities to generate a timestamp token that will be preserved. This can be used to prove that a certain file existed at a certain time, useful for legal purposes, for example, to prove that a certain file was not tampered with after a certain date.
@@ -98,7 +102,7 @@ class TimestampingEnricher(Enricher):
         else:
             logger.warning(f"No successful timestamps for {url=}")
 
-    def verify_signed(self, timestamp_response: TimeStampResponse, signature: bytes) ->  None:
+    def verify_signed(self, timestamp_response: TimeStampResponse, message: bytes) ->  None:
         """
         Verify a Signed Timestamp using the TSA provided by the Trusted Root.
         """
@@ -117,6 +121,16 @@ class TimestampingEnricher(Enricher):
         for i, cert in enumerate(timestamp_certs): # cannot use list comprehension, it's a set
             intermediate_certs.append(cert)
 
+
+        message_hash = None
+        hash_algorithm = timestamp_response.tst_info.message_imprint.hash_algorithm
+        if hash_algorithm == x509.ObjectIdentifier(value="2.16.840.1.101.3.4.2.3"):
+            message_hash = hashlib.sha512(message).digest()
+        elif hash_algorithm == x509.ObjectIdentifier(value="2.16.840.1.101.3.4.2.1"):
+            message_hash = hashlib.sha256(message).digest()
+        else:
+            raise ValueError(f"Unsupported hash algorithm: {hash_algorithm}")
+        
         for certificate in cert_authorities:
             builder = VerifierBuilder()
             builder.add_root_certificate(certificate)
@@ -125,8 +139,10 @@ class TimestampingEnricher(Enricher):
                 builder.add_intermediate_certificate(intermediate_cert)
 
             verifier = builder.build()
+
+            
             try:
-                verifier.verify(timestamp_response, signature)
+                verifier.verify(timestamp_response, message_hash)
                 return certificate
             except Rfc3161VerificationError as e:
                 logger.debug(f"Unable to verify Timestamp with CA {certificate.subject}: {e}")
@@ -163,7 +179,7 @@ class TimestampingEnricher(Enricher):
     def save_certificate(self, tsp_response: TimeStampResponse) -> list[Media]:
         # returns the leaf certificate URL, fails if not set
 
-        certificates = self.load_tst_certs(tsp_response)
+        certificates = self.tst_certs(tsp_response)
 
 
         cert_chain = []

tests/enrichers/test_timestamping_enricher.py

@@ -6,6 +6,8 @@ from rfc3161_client import (
     decode_timestamp_response,
 )
 
+from cryptography import x509
+
 @pytest.fixture
 def digicert():
     with open("tests/data/timestamp_token_digicert_com.crt", "rb") as f:
@@ -20,17 +22,12 @@ def test_sign_data(setup_module):
     result: TimeStampResponse = tsp.sign_data(tsa_url, data)
     assert isinstance(result, TimeStampResponse)
 
-    cert_chain = tsp.download_certificate(result)
-
-    assert len(cert_chain) == 2
-
-    try:
-        valid_root = tsp.verify_signed(result, data)
-        assert valid_root.subject == "CN=Entrust Root Certification Authority - G2, OU=(c) 2009 Entrust, Inc. - for authorized use only, OU=See www.entrust.net/legal-terms, O=Entrust, Inc., C="
-    except Exception as e:
-        pytest.fail(f"Verification failed: {e}")
+    root_cert: x509.Certificate = tsp.verify_signed(result, data)
+    assert root_cert.subject.rfc4514_string() == "CN=IdenTrust Commercial Root CA 1,O=IdenTrust,C=US"
 
     # test downloading the cert
+    cert_chain = tsp.save_certificate(result)
+    assert len(cert_chain) == 2
 
 def test_tsp_enricher_download_syndication(setup_module, digicert):
     tsp: TimestampingEnricher = setup_module("timestamping_enricher")