Merge pull request #18 from GiacomoPope/pk_from_sk

allow creation of pk from sk

Author: GiacomoPope

src/dilithium_py/ml_dsa/ml_dsa.py

@@ -420,6 +420,40 @@ def verify(self, pk_bytes, m, sig_bytes, ctx=b""):
 
         return self._verify_internal(pk_bytes, m_prime, sig_bytes)
 
+    """
+    The following additional function follows an outline from:
+        https://github.com/aws/aws-lc/pull/2142
+    which computes pk_bytes when only the sk_bytes are known.
+    """
+
+    def pk_from_sk(self, sk_bytes: bytes) -> bytes:
+        """
+        Given the packed representation of a ML-DSA secret key,
+        compute the corresponding packed public key bytes.
+        """
+        # First unpack the secret key
+        rho, K, tr, s1, s2, t0 = self._unpack_sk(sk_bytes)
+
+        # Compute the matrix A from rho in NTT form
+        A_hat = self._expand_matrix_from_seed(rho)
+
+        # Convert s1 to NTT form
+        s1_hat = s1.to_ntt()
+
+        # Compute the polynomial t, we have the lower bits t0,
+        # but we need the higher bits t1 for the public key
+        t = (A_hat @ s1_hat).from_ntt() + s2
+        t1, _ = t.power_2_round(self.d)
+
+        # The packed public key is made from rho || t1
+        pk_bytes = self._pack_pk(rho, t1)
+
+        # Ensure the public key matches the hash within the secret key
+        if tr != self._h(pk_bytes, 64):
+            raise ValueError("maleformed secret key")
+
+        return pk_bytes
+
     """
     The following external mu functions are not in FIPS 204, but are in
     Appendix D of the following IETF draft and are included for experimentation

src/dilithium_py/ml_dsa/pkcs.py

@@ -266,10 +266,10 @@ def sk_from_der(enc_key):
     if not expanded:
         pk, expanded = ml_dsa.key_derive(seed)
 
-    if not pk and seed:
-        # deriving verifying key from expanded key in ML-DSA is non-trivial
-        # do that only for encodings that include the seed
-        pk, _ = ml_dsa.key_derive(seed)
+    if not pk:
+        # If we reach here, we need to compute the public key
+        # directly from the secret key bytes
+        pk = ml_dsa.pk_from_sk(expanded)
 
     return ml_dsa, expanded, seed, pk
 

tests/test_ml_dsa.py

@@ -31,6 +31,12 @@ def generic_test_ml_dsa(self, ML_DSA, count=5):
             check_wrong_msg = ML_DSA.verify(pk, b"", sig, ctx=ctx)
             check_no_ctx = ML_DSA.verify(pk, msg, sig)
 
+            # Generate the public key directly from the secret key
+            recovered_pk = ML_DSA.pk_from_sk(sk)
+
+            # Check that recovering the pk works
+            self.assertEqual(pk, recovered_pk)
+
             # Check that signature works
             self.assertTrue(check_verify)
 

tests/test_pkcs.py

@@ -161,7 +161,7 @@ def test_import_from_expanded(self):
         self.assertIs(ml_dsa, self.ml_dsa)
         self.assertEqual(sk, self.sk)
         self.assertEqual(seed, None)
-        self.assertEqual(pk, None)
+        self.assertEqual(pk, self.pk)
 
 
 @unittest.skipUnless(ECDSA_PRESENT, "requires ecdsa package")
@@ -1006,7 +1006,7 @@ def test_sk_sanity_expanded_only(self):
         self.assertEqual(self.ml_dsa, ml_dsa)
         self.assertEqual(self.sk, expanded)
         self.assertEqual(None, seed)
-        self.assertEqual(None, pk)
+        self.assertEqual(self.pk, pk)
 
     def test_sk_trailing_junk(self):
         enc = (