A key agreement algorithm can contain a key derivation

PSA_ALG_KEY_AGREEMENT(..., kdf) is a valid key derivation algorithm
when kdf is one.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/scripts/mbedtls_dev/crypto_knowledge.py b/scripts/mbedtls_dev/crypto_knowledge.py
index adda02f..73890f8 100644
--- a/scripts/mbedtls_dev/crypto_knowledge.py
+++ b/scripts/mbedtls_dev/crypto_knowledge.py
@@ -366,3 +366,23 @@
         self.head = self.determine_head(self.base_expression)
         self.category = self.determine_category(self.base_expression, self.head)
         self.is_wildcard = self.determine_wildcard(self.expression)
+
+    def is_key_agreement_with_derivation(self) -> bool:
+        """Whether this is a combined key agreement and key derivation algorithm."""
+        if self.category != AlgorithmCategory.KEY_AGREEMENT:
+            return False
+        m = re.match(r'PSA_ALG_KEY_AGREEMENT\(\w+,\s*(.*)\)\Z', self.expression)
+        if not m:
+            return False
+        kdf_alg = m.group(1)
+        # Assume kdf_alg is either a valid KDF or 0.
+        return not re.match(r'(?:0[Xx])?0+\s*\Z', kdf_alg)
+
+    def can_do(self, category: AlgorithmCategory) -> bool:
+        """Whether this algorithm fits the specified operation category."""
+        if category == self.category:
+            return True
+        if category == AlgorithmCategory.KEY_DERIVATION and \
+           self.is_key_agreement_with_derivation():
+            return True
+        return False
diff --git a/tests/scripts/generate_psa_tests.py b/tests/scripts/generate_psa_tests.py
index 6c05b89..81abce7 100755
--- a/tests/scripts/generate_psa_tests.py
+++ b/tests/scripts/generate_psa_tests.py
@@ -375,7 +375,7 @@
             category: crypto_knowledge.AlgorithmCategory,
     ) -> Iterator[test_case.TestCase]:
         """Generate failure test cases for keyless operations with the specified algorithm."""
-        if category == alg.category:
+        if alg.can_do(category):
             # Compatible operation, unsupported algorithm
             for dep in automatic_dependencies(alg.base_expression):
                 yield self.make_test_case(alg, category,
@@ -394,7 +394,7 @@
         for kt in self.key_types:
             key_is_compatible = kt.can_do(alg)
             # To do: public key for a private key operation
-            if key_is_compatible and category == alg.category:
+            if key_is_compatible and alg.can_do(category):
                 # Compatible key and operation, unsupported algorithm
                 for dep in automatic_dependencies(alg.base_expression):
                     yield self.make_test_case(alg, category,
@@ -405,7 +405,7 @@
                 yield self.make_test_case(alg, category,
                                           self.Reason.INVALID,
                                           kt=kt)
-            elif category == alg.category:
+            elif alg.can_do(category):
                 # Incompatible key, compatible operation, supported algorithm
                 yield self.make_test_case(alg, category,
                                           self.Reason.INCOMPATIBLE,