Add key usage policy extension support for key generation

Signed-off-by: gabor-mezei-arm <gabor.mezei@arm.com>
diff --git a/scripts/mbedtls_dev/psa_storage.py b/scripts/mbedtls_dev/psa_storage.py
index 45f0380..4cd3dfe 100644
--- a/scripts/mbedtls_dev/psa_storage.py
+++ b/scripts/mbedtls_dev/psa_storage.py
@@ -101,6 +101,12 @@
     LATEST_VERSION = 0
     """The latest version of the storage format."""
 
+    EXTENDABLE_USAGE_FLAGS = {
+        Expr('PSA_KEY_USAGE_SIGN_HASH'): Expr('PSA_KEY_USAGE_SIGN_MESSAGE'),
+        Expr('PSA_KEY_USAGE_VERIFY_HASH'): Expr('PSA_KEY_USAGE_VERIFY_MESSAGE')
+    } #type: Dict[Expr, Expr]
+    """The extendable usage flags with the corresponding extension flags."""
+
     def __init__(self, *,
                  version: Optional[int] = None,
                  id: Optional[int] = None, #pylint: disable=redefined-builtin
@@ -108,18 +114,27 @@
                  type: Exprable, #pylint: disable=redefined-builtin
                  bits: int,
                  usage: Exprable, alg: Exprable, alg2: Exprable,
-                 material: bytes #pylint: disable=used-before-assignment
+                 material: bytes, #pylint: disable=used-before-assignment
+                 usage_extension: bool = True
                 ) -> None:
         self.version = self.LATEST_VERSION if version is None else version
         self.id = id #pylint: disable=invalid-name #type: Optional[int]
         self.lifetime = as_expr(lifetime) #type: Expr
         self.type = as_expr(type) #type: Expr
         self.bits = bits #type: int
-        self.usage = as_expr(usage) #type: Expr
+        self.original_usage = as_expr(usage) #type: Expr
+        self.updated_usage = self.original_usage #type: Expr
         self.alg = as_expr(alg) #type: Expr
         self.alg2 = as_expr(alg2) #type: Expr
         self.material = material #type: bytes
 
+        if usage_extension:
+            for flag, extension in self.EXTENDABLE_USAGE_FLAGS.items():
+                if self.original_usage.value() & flag.value() and \
+                   self.original_usage.value() & extension.value() == 0:
+                    self.updated_usage = Expr(self.updated_usage.string +
+                                              ' | ' + extension.string)
+
     MAGIC = b'PSA\000KEY\000'
 
     @staticmethod
@@ -151,7 +166,7 @@
         if self.version == 0:
             attributes = self.pack('LHHLLL',
                                    self.lifetime, self.type, self.bits,
-                                   self.usage, self.alg, self.alg2)
+                                   self.updated_usage, self.alg, self.alg2)
             material = self.pack('L', len(self.material)) + self.material
         else:
             raise NotImplementedError