import re
from typing import Dict, Iterable, Optional, Pattern, Tuple
from mbedtls_dev.asymmetric_key_data import ASYMMETRIC_KEY_DATA
class KeyType:
def __init__(self, name: str, params: Optional[Iterable[str]] = None):
self.name = name.strip()
if params is None:
if '(' in self.name:
m = re.match(r'(\w+)\s*\((.*)\)\Z', self.name)
assert m is not None
self.name = m.group(1)
params = m.group(2).split(',')
self.params = (None if params is None else
[param.strip() for param in params])
assert re.match(r'PSA_KEY_TYPE_\w+\Z', self.name)
self.expression = self.name
if self.params is not None:
self.expression += '(' + ', '.join(self.params) + ')'
self.private_type = re.sub(r'_PUBLIC_KEY\Z', r'_KEY_PAIR', self.name)
ECC_KEY_SIZES = {
'PSA_ECC_FAMILY_SECP_K1': (192, 224, 256),
'PSA_ECC_FAMILY_SECP_R1': (225, 256, 384, 521),
'PSA_ECC_FAMILY_SECP_R2': (160,),
'PSA_ECC_FAMILY_SECT_K1': (163, 233, 239, 283, 409, 571),
'PSA_ECC_FAMILY_SECT_R1': (163, 233, 283, 409, 571),
'PSA_ECC_FAMILY_SECT_R2': (163,),
'PSA_ECC_FAMILY_BRAINPOOL_P_R1': (160, 192, 224, 256, 320, 384, 512),
'PSA_ECC_FAMILY_MONTGOMERY': (255, 448),
'PSA_ECC_FAMILY_TWISTED_EDWARDS': (255, 448),
}
KEY_TYPE_SIZES = {
'PSA_KEY_TYPE_AES': (128, 192, 256), 'PSA_KEY_TYPE_ARIA': (128, 192, 256), 'PSA_KEY_TYPE_CAMELLIA': (128, 192, 256), 'PSA_KEY_TYPE_CHACHA20': (256,), 'PSA_KEY_TYPE_DERIVE': (120, 128), 'PSA_KEY_TYPE_DES': (64, 128, 192), 'PSA_KEY_TYPE_HMAC': (128, 160, 224, 256, 384, 512), 'PSA_KEY_TYPE_PASSWORD': (48, 168, 336), 'PSA_KEY_TYPE_PASSWORD_HASH': (128, 256), 'PSA_KEY_TYPE_PEPPER': (128, 256), 'PSA_KEY_TYPE_RAW_DATA': (8, 40, 128), 'PSA_KEY_TYPE_RSA_KEY_PAIR': (1024, 1536), }
def sizes_to_test(self) -> Tuple[int, ...]:
if self.private_type == 'PSA_KEY_TYPE_ECC_KEY_PAIR':
assert self.params is not None
return self.ECC_KEY_SIZES[self.params[0]]
return self.KEY_TYPE_SIZES[self.private_type]
DATA_BLOCK = b'Here\000is key\240data'
def key_material(self, bits: int) -> bytes:
if self.expression in ASYMMETRIC_KEY_DATA:
if bits not in ASYMMETRIC_KEY_DATA[self.expression]:
raise ValueError('No key data for {}-bit {}'
.format(bits, self.expression))
return ASYMMETRIC_KEY_DATA[self.expression][bits]
if bits % 8 != 0:
raise ValueError('Non-integer number of bytes: {} bits for {}'
.format(bits, self.expression))
length = bits // 8
if self.name == 'PSA_KEY_TYPE_DES':
des3 = b'dEs kEy\001dEs kEy\002dEs kEy\004'
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': re.compile('.*KEY_PAIR'),
'PSA_KEY_USAGE_VERIFY_HASH': re.compile('.*KEY.*')
}
def is_valid_for_signature(self, usage: str) -> bool:
return re.match(self.KEY_TYPE_FOR_SIGNATURE[usage], self.name) is not None