allow signing via digest, fixes #4

pyauth · Mar 14, 2024 · 22f766f · 22f766f
1 parent ed1a52b
commit 22f766f
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 14 deletions.
diff --git a/README.rst b/README.rst
@@ -27,10 +27,12 @@ Synopsis
     # Verify a presented timestamp token offline using the original message
     verified = TSPVerifier().verify(signed, message=message)
 
-    # Or verify using the message digest (digest algorithm may vary)
+    # Or sign and verify using the message digest (digest algorithm may vary)
     import hashlib
 
     digest = hashlib.sha512(message).digest()
+    
+    signer.sign(message_digest=digest)
     verified = TSPVerifier().verify(signed, message_digest=digest)
 
     print(verified.tst_info)  # Parsed TSTInfo (CMS SignedData) structure

diff --git a/test/test.py b/test/test.py
@@ -31,13 +31,7 @@ def setUp(self):
         self.signer = TSPSigner()
         self.verifier = TSPVerifier()
 
-    def test_basic_tsp_client_operations(self):
-        message = b"abc"
-        signed = self.signer.sign(message)
-        verified = self.verifier.verify(signed, message=message)
-        digest = hashlib.sha512(message).digest()
-        verified = self.verifier.verify(signed, message_digest=digest)
-
+    def check_results(self, verified, signed, message, digest):
         self.assertTrue(verified.tst_info)
         self.assertTrue(verified.signed_attrs)
 
@@ -53,6 +47,25 @@ def test_basic_tsp_client_operations(self):
         with self.assertRaises(NonceMismatchError):
             self.verifier.verify(signed, message_digest=digest, nonce=123)
 
+    def test_basic_tsp_client_operations(self):
+        message = b"abc"
+        digest = hashlib.sha512(message).digest()
+
+        # sign and verify by message
+        signed = self.signer.sign(message)
+        verified_by_message = self.verifier.verify(signed, message=message)
+        self.check_results(verified_by_message, signed, message, digest)
+
+        # verify by digest
+        verified_by_digest = self.verifier.verify(signed, message_digest=digest)
+        self.check_results(verified_by_digest, signed, message, digest)
+
+        # sign and verify by digest only
+        signed_by_digest = self.signer.sign(message_digest=digest)
+        verified_by_digest = self.verifier.verify(signed_by_digest, message_digest=digest)
+        self.check_results(verified_by_digest, signed_by_digest, message, digest)
+
+
     def test_set_custom_tsa(self):
         message = b"abc"
         for tsp_server in self.tsp_servers:

diff --git a/tsp_client/signer.py b/tsp_client/signer.py
@@ -8,7 +8,7 @@
 from asn1crypto import algos, tsp
 
 from .algorithms import DigestAlgorithm
-from .exceptions import TSPClientSigningError
+from .exceptions import TSPClientSigningError, InvalidInput
 from .verifier import TSPVerifier, VerifyResult
 
 
@@ -34,18 +34,26 @@ def _verify_timestamp(self, verify_result: VerifyResult):
         if verify_result.tst_info["gen_time"] > now + self.max_clock_drift:
             raise TSPClientSigningError("Timestamp returned by server is too far in the future")
 
-    def sign(self, message, *, signing_settings: SigningSettings = SigningSettings()) -> bytes:
+    def sign(self, message=None, *, message_digest=None, signing_settings: SigningSettings = SigningSettings()) -> bytes:
+        if not message and not message_digest:
+            raise InvalidInput("Expected at least one of message or message_digest to be set")
+        if message is not None and message_digest is not None:
+            raise InvalidInput("Expected only one of message and message_digest to be set")
+
         hasher = signing_settings.digest_algorithm.implementation()
-        hasher.update(message)
-        digest = hasher.digest()
+
+        if message:
+            hasher.update(message)
+            message_digest = hasher.digest()
+
         nonce = int.from_bytes(secrets.token_bytes(), byteorder=sys.byteorder)
         tsp_request = tsp.TimeStampReq(
             {
                 "version": 1,
                 "message_imprint": tsp.MessageImprint(
                     {
                         "hash_algorithm": algos.DigestAlgorithm({"algorithm": hasher.name}),
-                        "hashed_message": digest,
+                        "hashed_message": message_digest,
                     }
                 ),
                 "cert_req": True,
@@ -62,6 +70,6 @@ def sign(self, message, *, signing_settings: SigningSettings = SigningSettings()
                 f'{tsp_response["status"]["fail_info"].native}'
             )
         tst = tsp_response["time_stamp_token"].dump()
-        verify_result = self._verifier.verify(tst, nonce=nonce, message_digest=digest)
+        verify_result = self._verifier.verify(tst, nonce=nonce, message_digest=message_digest)
         self._verify_timestamp(verify_result)
         return tst