diff --git a/scripts/mbedtls_dev/crypto_knowledge.py b/scripts/mbedtls_dev/crypto_knowledge.py
index 4b4e2df..01a8c5d 100644
--- a/scripts/mbedtls_dev/crypto_knowledge.py
+++ b/scripts/mbedtls_dev/crypto_knowledge.py
@@ -19,7 +19,7 @@
 # limitations under the License.
 
 import re
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Optional, Tuple, Dict
 
 from mbedtls_dev.asymmetric_key_data import ASYMMETRIC_KEY_DATA
 
@@ -138,3 +138,18 @@
             return des3[:length]
         return b''.join([self.DATA_BLOCK] * (length // len(self.DATA_BLOCK)) +
                         [self.DATA_BLOCK[:length % len(self.DATA_BLOCK)]])
+
+    KEY_TYPE_FOR_SIGNATURE = {
+        'PSA_KEY_USAGE_SIGN_HASH': '.*KEY_PAIR',
+        'PSA_KEY_USAGE_VERIFY_HASH': '.*KEY.*'
+    } #type: Dict[str, str]
+    """Use a regexp to determine key types for which signature is possible
+       when using the actual usage flag.
+    """
+    def is_valid_for_signature(self, usage: str) -> bool:
+        """Determine if the key type is compatible with the specified
+           signitute type.
+
+        """
+        # This is just temporaly solution for the implicit usage flags.
+        return re.match(self.KEY_TYPE_FOR_SIGNATURE[usage], self.name) is not None
diff --git a/scripts/mbedtls_dev/psa_storage.py b/scripts/mbedtls_dev/psa_storage.py
index 5ff1738..88992a6 100644
--- a/scripts/mbedtls_dev/psa_storage.py
+++ b/scripts/mbedtls_dev/psa_storage.py
@@ -107,14 +107,6 @@
     } #type: Dict[str, str]
     """Mapping of usage flags to the flags that they imply."""
 
-    IMPLICIT_USAGE_FLAGS_KEY_RESTRICTION = {
-        'PSA_KEY_USAGE_SIGN_HASH': '.*KEY_PAIR',
-        'PSA_KEY_USAGE_VERIFY_HASH': '.*KEY.*'
-    } #type: Dict[str, str]
-    """Use a regexp to determine key types for which signature is possible
-       when using the actual usage flag.
-    """
-
     def __init__(self, *,
                  version: Optional[int] = None,
                  id: Optional[int] = None, #pylint: disable=redefined-builtin
diff --git a/tests/scripts/generate_psa_tests.py b/tests/scripts/generate_psa_tests.py
index fe865ae..4e2842c 100755
--- a/tests/scripts/generate_psa_tests.py
+++ b/tests/scripts/generate_psa_tests.py
@@ -494,32 +494,30 @@
             self,
             implyer_usage: str,
             alg: str,
-            key_type: str,
-            params: Optional[Iterable[str]] = None
+            key_type: crypto_knowledge.KeyType
     ) -> StorageKey:
         # pylint: disable=too-many-locals
         """Generate test keys for the specified implicit usage flag,
            algorithm and key type combination.
         """
-        kt = crypto_knowledge.KeyType(key_type, params)
-        bits = kt.sizes_to_test()[0]
+        bits = key_type.sizes_to_test()[0]
         implicit_usage = StorageKey.IMPLICIT_USAGE_FLAGS[implyer_usage]
         usage_flags = 'PSA_KEY_USAGE_EXPORT'
         material_usage_flags = usage_flags + ' | ' + implyer_usage
         expected_usage_flags = material_usage_flags + ' | ' + implicit_usage
         alg2 = 0
-        key_material = kt.key_material(bits)
+        key_material = key_type.key_material(bits)
         usage_expression = re.sub(r'PSA_KEY_USAGE_', r'', implyer_usage)
         alg_expression = re.sub(r'PSA_ALG_', r'', alg)
         alg_expression = re.sub(r',', r', ', re.sub(r' +', r'', alg_expression))
         key_type_expression = re.sub(r'\bPSA_(?:KEY_TYPE|ECC_FAMILY)_',
                                      r'',
-                                     kt.expression)
+                                     key_type.expression)
         description = 'implied by {}: {} {} {}-bit'.format(
             usage_expression, alg_expression, key_type_expression, bits)
         return StorageKey(version=self.version,
                           id=1, lifetime=0x00000001,
-                          type=kt.expression, bits=bits,
+                          type=key_type.expression, bits=bits,
                           usage=material_usage_flags,
                           expected_usage=expected_usage_flags,
                           alg=alg, alg2=alg2,
@@ -581,14 +579,14 @@
         # flag to generate a valid key for exercising. The key is generated
         # without usage extension to check the extension compatiblity.
         alg_with_keys = self.gather_key_types_for_sign_alg()
-        key_filter = StorageKey.IMPLICIT_USAGE_FLAGS_KEY_RESTRICTION
 
         for usage in sorted(StorageKey.IMPLICIT_USAGE_FLAGS, key=str):
             for alg in sorted(alg_with_keys):
                 for key_type in sorted(alg_with_keys[alg]):
                     # The key types must be filtered to fit the specific usage flag.
-                    if re.match(key_filter[usage], key_type):
-                        yield self.keys_for_implicit_usage(usage, alg, key_type)
+                    kt = crypto_knowledge.KeyType(key_type)
+                    if kt.is_valid_for_signature(usage):
+                        yield self.keys_for_implicit_usage(usage, alg, kt)
 
     def generate_all_keys(self) -> List[StorageKey]:
         keys = super().generate_all_keys()