def encrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
if len(self.phdr):
raise CoseException(
f"Protected header must be empty when using an AE algorithm: {alg}"
)
if alg is None:
raise CoseException(
"The algorithm parameter should at least be included in the unprotected header"
)
elif alg in {A128KW, A192KW, A256KW}:
key_ops = [WrapOp, EncryptOp]
kek = SymmetricKey(k=self._compute_kek(target_alg, ops='encrypt'),
optional_params={
KpAlg: alg,
KpKeyOps: key_ops
})
kek.verify(SymmetricKey, alg, [WrapOp, EncryptOp])
elif alg in {RsaesOaepSha512, RsaesOaepSha256, RsaesOaepSha1}:
kek = self.key
kek.verify(RSAKey, alg, [WrapOp, EncryptOp])
else:
raise CoseIllegalAlgorithm(f"Algorithm {alg} for {self.__name__}")
return alg.key_wrap(kek, self.payload)
def decrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
kek = SymmetricKey(k=self._compute_kek(target_alg, 'decrypt'),
optional_params={
KpAlg: alg,
KpKeyOps: [DecryptOp, UnwrapOp]
})
kek.verify(SymmetricKey, alg, [UnwrapOp, DecryptOp])
return alg.key_unwrap(kek, self.payload)
def decrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
key_ops = [DecryptOp, UnwrapOp]
if alg in {A128KW, A192KW, A256KW}:
kek = SymmetricKey(k=self._compute_kek(target_alg, 'decrypt'),
optional_params={
KpAlg: alg,
KpKeyOps: key_ops
})
kek.verify(SymmetricKey, alg, [UnwrapOp, DecryptOp])
elif alg in {RsaesOaepSha512, RsaesOaepSha256, RsaesOaepSha1}:
kek = self.key
kek.verify(RSAKey, alg, [UnwrapOp, DecryptOp])
else:
raise CoseException(
f"Unsupported algorithm for key unwrapping: {alg}")
return alg.key_unwrap(kek, self.payload)
def decrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
_ = target_alg
if alg in {EcdhEsA256KW, EcdhEsA192KW, EcdhEsA128KW}:
peer_key = self.get_attr(headers.EphemeralKey)
elif alg in {EcdhSsA256KW, EcdhSsA192KW, EcdhSsA128KW}:
peer_key = self.get_attr(headers.StaticKey)
else:
raise CoseIllegalAlgorithm(
f"Algorithm {alg} unsupported for {self.__name__}")
kek = SymmetricKey(k=self._compute_kek(alg.get_key_wrap_func(),
peer_key, self.key, alg),
optional_params={
KpAlg: alg,
KpKeyOps: [UnwrapOp, DecryptOp]
})
kek.verify(SymmetricKey, alg, [UnwrapOp, DecryptOp])
return alg.get_key_wrap_func().key_unwrap(kek, self.payload)
class KeyWrap(CoseRecipient):
@classmethod
def from_cose_obj(cls, cose_obj: list, *args, **kwargs) -> 'KeyWrap':
msg = super().from_cose_obj(cose_obj)
msg.context = kwargs.get('context')
# only AE algorithms supported thus the protected header must be empty
alg = msg.get_attr(headers.Algorithm)
if alg in {A128KW, A192KW, A256KW} and len(msg.phdr):
raise CoseMalformedMessage(
f"Recipient class KEY_WRAP with alg {alg} must have a zero-length protected header"
)
if msg.payload == b'':
raise CoseMalformedMessage(
f'Recipient class KEY_WRAP must carry the encrypted CEK in its payload'
)
msg.recipients = [
CoseRecipient.create_recipient(r, context='Rec_Recipient')
for r in msg.recipients
]
return msg
@property
def context(self):
return self._context
@context.setter
def context(self, context: str):
self._context = context
def encode(self, *args, **kwargs) -> list:
recipient = [
self.phdr_encoded, self.uhdr_encoded,
self.encrypt(kwargs.get('target_alg'))
]
if len(self.recipients):
recipient.append(
[r.encode(*args, **kwargs) for r in self.recipients])
return recipient
def _compute_kek(self, target_alg: '_EncAlg', ops: 'str') -> bytes:
if self.key is None:
# try to derive from this recipients' recipient list
if not len(self.recipients):
raise CoseException(f"No key found to {ops} the CEK")
else:
r_types = CoseRecipient.verify_recipients(self.recipients)
if ops == 'encrypt':
if DirectKeyAgreement in r_types:
self.key = self.recipients[0].compute_cek(target_alg)
elif KeyWrap in r_types or KeyAgreementWithKeyWrap in r_types:
key_bytes = os.urandom(self.get_attr(
headers.Algorithm))
for r in self.recipients:
r.payload = key_bytes
self.key = SymmetricKey(k=key_bytes)
else:
raise CoseException('Unsupported COSE recipient class')
else:
if DirectKeyAgreement in r_types or KeyWrap in r_types or KeyAgreementWithKeyWrap in r_types:
self.key = self.recipients[0].decrypt(
self.get_attr(headers.Algorithm))
else:
raise CoseException('Unsupported COSE recipient class')
if self.key is None:
raise CoseException("No key found to decrypt the CEK")
return self.key.k
def compute_cek(self, target_alg: '_EncAlg', ops: str) -> Optional['SK']:
if ops == "encrypt":
if self.payload == b'':
return None
else:
return SymmetricKey(k=self.payload,
optional_params={
KpAlg: target_alg,
KpKeyOps: [EncryptOp]
})
else:
return SymmetricKey(k=self.decrypt(target_alg),
optional_params={
KpAlg: target_alg,
KpKeyOps: [DecryptOp]
})
def encrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
if A128KW.identifier >= alg.identifier >= A256KW.identifier:
if len(self.phdr):
raise CoseException(
f"Protected header must be empty when using an AE algorithm: {alg}"
)
if alg is None:
raise CoseException(
"The algorithm parameter should at least be included in the unprotected header"
)
self.key = SymmetricKey(k=self._compute_kek(target_alg,
ops='encrypt'),
optional_params={
KpAlg: alg,
KpKeyOps: [WrapOp, EncryptOp]
})
self.key.verify(SymmetricKey, alg, [WrapOp, EncryptOp])
return alg.key_wrap(self.key, self.payload)
else:
raise CoseIllegalAlgorithm(f"Algorithm {alg} for {self.__name__}")
def decrypt(self, target_alg: '_EncAlg') -> bytes:
alg = self.get_attr(headers.Algorithm)
kek = SymmetricKey(k=self._compute_kek(target_alg, 'decrypt'),
optional_params={
KpAlg: alg,
KpKeyOps: [DecryptOp, UnwrapOp]
})
kek.verify(SymmetricKey, alg, [UnwrapOp, DecryptOp])
return alg.key_unwrap(kek, self.payload)
def __repr__(self) -> str:
phdr, uhdr = self._hdr_repr()
return f'<COSE_Recipient: [{phdr}, {uhdr}, {utils.truncate(self._payload)}, {str(self.recipients)}]>'