From 390543381cce8f511c7abe91a310c3f5c2b5d751 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Thu, 29 Apr 2021 20:38:47 +0200
Subject: [PATCH] Add knowledge of the compatibility of key types and
 algorithms

Determine key types that are compatible with an algorithm based on
their names.

Key derivation and PAKE are not yet supported.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
---
 scripts/mbedtls_dev/crypto_knowledge.py | 62 +++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/scripts/mbedtls_dev/crypto_knowledge.py b/scripts/mbedtls_dev/crypto_knowledge.py
index 05db887a09..49ff0d9371 100644
--- a/scripts/mbedtls_dev/crypto_knowledge.py
+++ b/scripts/mbedtls_dev/crypto_knowledge.py
@@ -25,6 +25,7 @@ from typing import Dict, Iterable, Optional, Pattern, Tuple
 from mbedtls_dev.asymmetric_key_data import ASYMMETRIC_KEY_DATA
 
 
+BLOCK_CIPHERS = frozenset(['AES', 'ARIA', 'CAMELLIA', 'DES'])
 BLOCK_MAC_MODES = frozenset(['CBC_MAC', 'CMAC'])
 BLOCK_CIPHER_MODES = frozenset([
     'CTR', 'CFB', 'OFB', 'XTS', 'CCM_STAR_NO_TAG',
@@ -32,6 +33,25 @@ BLOCK_CIPHER_MODES = frozenset([
 ])
 BLOCK_AEAD_MODES = frozenset(['CCM', 'GCM'])
 
+class EllipticCurveCategory(enum.Enum):
+    """Categorization of elliptic curve families.
+
+    The category of a curve determines what algorithms are defined over it.
+    """
+
+    SHORT_WEIERSTRASS = 0
+    MONTGOMERY = 1
+    TWISTED_EDWARDS = 2
+
+    @staticmethod
+    def from_family(family: str) -> 'EllipticCurveCategory':
+        if family == 'PSA_ECC_FAMILY_MONTGOMERY':
+            return EllipticCurveCategory.MONTGOMERY
+        if family == 'PSA_ECC_FAMILY_TWISTED_EDWARDS':
+            return EllipticCurveCategory.TWISTED_EDWARDS
+        # Default to SW, which most curves belong to.
+        return EllipticCurveCategory.SHORT_WEIERSTRASS
+
 
 class KeyType:
     """Knowledge about a PSA key type."""
@@ -72,6 +92,11 @@ class KeyType:
         if self.params is not None:
             self.expression += '(' + ', '.join(self.params) + ')'
 
+        m = re.match(r'PSA_KEY_TYPE_(\w+)', self.name)
+        assert m
+        self.head = re.sub(r'_(?:PUBLIC_KEY|KEY_PAIR)\Z', r'', m.group(1))
+        """The key type macro name, with common prefixes and suffixes stripped."""
+
         self.private_type = re.sub(r'_PUBLIC_KEY\Z', r'_KEY_PAIR', self.name)
         """The key type macro name for the corresponding key pair type.
 
@@ -162,6 +187,43 @@ class KeyType:
         # This is just temporaly solution for the implicit usage flags.
         return re.match(self.KEY_TYPE_FOR_SIGNATURE[usage], self.name) is not None
 
+    def can_do(self, alg: 'Algorithm') -> bool:
+        """Whether this key type can be used for operations with the given algorithm.
+
+        This function does not currently handle key derivation or PAKE.
+        """
+        #pylint: disable=too-many-return-statements
+        if alg.is_wildcard:
+            return False
+        if self.head == 'HMAC' and alg.head == 'HMAC':
+            return True
+        if self.head in BLOCK_CIPHERS and \
+           alg.head in frozenset.union(BLOCK_MAC_MODES,
+                                       BLOCK_CIPHER_MODES,
+                                       BLOCK_AEAD_MODES):
+            return True
+        if self.head == 'CHACHA20' and alg.head == 'CHACHA20_POLY1305':
+            return True
+        if self.head in {'ARC4', 'CHACHA20'} and \
+           alg.head == 'STREAM_CIPHER':
+            return True
+        if self.head == 'RSA' and alg.head.startswith('RSA_'):
+            return True
+        if self.head == 'ECC':
+            assert self.params is not None
+            eccc = EllipticCurveCategory.from_family(self.params[0])
+            if alg.head == 'ECDH' and \
+               eccc in {EllipticCurveCategory.SHORT_WEIERSTRASS,
+                        EllipticCurveCategory.MONTGOMERY}:
+                return True
+            if alg.head == 'ECDSA' and \
+               eccc == EllipticCurveCategory.SHORT_WEIERSTRASS:
+                return True
+            if alg.head in {'PURE_EDDSA', 'EDDSA_PREHASH'} and \
+               eccc == EllipticCurveCategory.TWISTED_EDWARDS:
+                return True
+        return False
+
 
 class AlgorithmCategory(enum.Enum):
     """PSA algorithm categories."""