267 lines
10 KiB
Python
267 lines
10 KiB
Python
import secrets
|
|
from typing import Any, Union
|
|
|
|
from .models import (
|
|
CompactEncryption,
|
|
BaseJSONEncryption,
|
|
GeneralJSONEncryption,
|
|
FlattenedJSONEncryption,
|
|
Recipient,
|
|
JWEAlgModel,
|
|
JWEEncModel,
|
|
JWEKeyAgreement,
|
|
JWEDirectEncryption,
|
|
JWEKeyEncryption,
|
|
JWEKeyWrapping,
|
|
)
|
|
from .registry import JWERegistry
|
|
from ..errors import (
|
|
JoseError,
|
|
DecodeError,
|
|
InvalidEncryptedKeyError,
|
|
InvalidExchangeKeyError,
|
|
ConflictAlgorithmError,
|
|
)
|
|
from ..util import (
|
|
json_b64encode,
|
|
urlsafe_b64encode,
|
|
)
|
|
|
|
__all__ = [
|
|
"EncryptionData",
|
|
"perform_encrypt",
|
|
"perform_decrypt",
|
|
]
|
|
|
|
EncryptionData = Union[CompactEncryption, GeneralJSONEncryption, FlattenedJSONEncryption]
|
|
|
|
|
|
def perform_encrypt(obj: EncryptionData, registry: JWERegistry) -> None:
|
|
enc = registry.get_enc(obj.protected["enc"])
|
|
cek, delayed_tasks = pre_encrypt_recipients(enc, obj.recipients, registry)
|
|
|
|
# Step 9, Generate a random JWE Initialization Vector of the correct size
|
|
# for the content encryption algorithm (if required for the algorithm);
|
|
# otherwise, let the JWE Initialization Vector be the empty octet sequence.
|
|
iv = enc.generate_iv()
|
|
|
|
# Step 10, Compute the encoded Initialization Vector value
|
|
# BASE64URL(JWE Initialization Vector).
|
|
obj.base64_segments["iv"] = urlsafe_b64encode(iv)
|
|
|
|
# Step 11, If a "zip" parameter was included, compress the plaintext using
|
|
# the specified compression algorithm and let M be the octet sequence
|
|
# representing the compressed plaintext; otherwise, let M be the octet
|
|
# sequence representing the plaintext.
|
|
assert obj.plaintext is not None
|
|
|
|
plaintext: bytes
|
|
if "zip" in obj.protected:
|
|
zip_ = registry.get_zip(obj.protected["zip"])
|
|
plaintext = zip_.compress(obj.plaintext)
|
|
else:
|
|
plaintext = obj.plaintext
|
|
|
|
# Step 13, Compute the Encoded Protected Header value BASE64URL(UTF8(JWE Protected Header)).
|
|
aad = json_b64encode(obj.protected)
|
|
|
|
# Step 14, Let the Additional Authenticated Data encryption parameter be
|
|
# ASCII(Encoded Protected Header). However, if a JWE AAD value is
|
|
# present (which can only be the case when using the JWE JSON Serialization),
|
|
# instead let the Additional Authenticated Data encryption parameter be
|
|
# ASCII(Encoded Protected Header || '.' || BASE64URL(JWE AAD)).
|
|
if isinstance(obj, BaseJSONEncryption) and obj.aad:
|
|
aad = aad + b"." + urlsafe_b64encode(obj.aad)
|
|
obj.base64_segments["aad"] = aad
|
|
|
|
# encrypting plaintext
|
|
ciphertext, tag = enc.encrypt(plaintext, cek, iv, aad)
|
|
|
|
# delay encrypting every recipient
|
|
post_encrypt_recipients(enc, delayed_tasks, cek, tag)
|
|
|
|
obj.base64_segments["ciphertext"] = urlsafe_b64encode(ciphertext)
|
|
obj.base64_segments["tag"] = urlsafe_b64encode(tag)
|
|
|
|
|
|
def perform_decrypt(obj: EncryptionData, registry: JWERegistry) -> None:
|
|
try:
|
|
_perform_decrypt(obj, registry)
|
|
except InvalidExchangeKeyError as error:
|
|
raise DecodeError(error.description)
|
|
|
|
|
|
def _perform_decrypt(obj: EncryptionData, registry: JWERegistry) -> None:
|
|
enc = registry.get_enc(obj.protected["enc"])
|
|
|
|
iv = obj.bytes_segments["iv"]
|
|
enc.check_iv(iv)
|
|
|
|
tag = obj.bytes_segments["tag"]
|
|
ciphertext = obj.bytes_segments["ciphertext"]
|
|
|
|
cek_set = set()
|
|
for recipient in obj.recipients:
|
|
try:
|
|
headers = recipient.headers()
|
|
registry.check_header(headers, True)
|
|
# Step 6, Determine the Key Management Mode employed by the algorithm
|
|
# specified by the "alg" (algorithm) Header Parameter.
|
|
alg = registry.get_alg(headers["alg"])
|
|
cek = decrypt_recipient(alg, enc, recipient, tag)
|
|
cek_set.add(cek)
|
|
except (AssertionError, JoseError) as error:
|
|
if registry.verify_all_recipients:
|
|
raise error
|
|
else:
|
|
if not registry.verify_all_recipients:
|
|
break
|
|
|
|
if not cek_set:
|
|
raise DecodeError("Invalid recipients")
|
|
|
|
if len(cek_set) > 1: # pragma: no cover
|
|
raise DecodeError("Multiple 'cek' found")
|
|
|
|
cek = cek_set.pop()
|
|
if len(cek) * 8 != enc.cek_size: # pragma: no cover
|
|
cek = secrets.token_bytes(enc.cek_size // 8)
|
|
|
|
if isinstance(obj, BaseJSONEncryption):
|
|
aad = json_b64encode(obj.protected)
|
|
if obj.aad:
|
|
aad = aad + b"." + obj.base64_segments["aad"]
|
|
else:
|
|
aad = obj.base64_segments["aad"]
|
|
|
|
msg = enc.decrypt(ciphertext, tag, cek, iv, aad)
|
|
if "zip" in obj.protected:
|
|
zip_ = registry.get_zip(obj.protected["zip"])
|
|
obj.plaintext = zip_.decompress(msg)
|
|
else:
|
|
obj.plaintext = msg
|
|
|
|
|
|
def pre_encrypt_recipients(
|
|
enc: JWEEncModel, recipients: list[Recipient[Any]], registry: JWERegistry
|
|
) -> tuple[bytes, list[tuple[JWEKeyAgreement, Recipient[Any]]]]:
|
|
cek: bytes = b""
|
|
delayed_tasks: list[tuple[JWEKeyAgreement, Recipient[Any]]] = []
|
|
for recipient in recipients:
|
|
alg = __prepare_recipient_algorithm(recipient, registry)
|
|
|
|
if alg.direct_mode:
|
|
if len(recipients) > 1:
|
|
raise ConflictAlgorithmError(f"Algorithm {alg.name} SHOULD have 1 recipient only")
|
|
cek = __pre_encrypt_direct_mode(alg, enc, recipient)
|
|
else:
|
|
if not cek:
|
|
# 2. When Key Wrapping, Key Encryption, or Key Agreement with Key
|
|
# Wrapping are employed, generate a random CEK value. See RFC
|
|
# 4086 [RFC4086] for considerations on generating random values.
|
|
# The CEK MUST have a length equal to that required for the
|
|
# content encryption algorithm.
|
|
cek = enc.generate_cek()
|
|
|
|
if isinstance(alg, JWEKeyAgreement):
|
|
delayed_tasks.append((alg, recipient))
|
|
else:
|
|
# 4. When Key Wrapping, or Key Encryption are employed, encrypt the CEK
|
|
# to the recipient and let the result be the JWE Encrypted Key.
|
|
assert isinstance(alg, (JWEKeyWrapping, JWEKeyEncryption))
|
|
recipient.encrypted_key = alg.encrypt_cek(cek, recipient)
|
|
return cek, delayed_tasks
|
|
|
|
|
|
def __prepare_recipient_algorithm(recipient: Recipient[Any], registry: JWERegistry) -> JWEAlgModel:
|
|
headers = recipient.headers()
|
|
registry.check_header(headers)
|
|
# 1. Determine the Key Management Mode employed by the algorithm used
|
|
# to determine the Content Encryption Key value. (This is the
|
|
# algorithm recorded in the "alg" (algorithm) Header Parameter of
|
|
# the resulting JWE.)
|
|
alg = registry.get_alg(headers["alg"])
|
|
|
|
if isinstance(alg, JWEKeyAgreement):
|
|
alg.prepare_ephemeral_key(recipient)
|
|
return alg
|
|
|
|
|
|
def __pre_encrypt_direct_mode(alg: JWEAlgModel, enc: JWEEncModel, recipient: Recipient[Any]) -> bytes:
|
|
cek: bytes
|
|
if isinstance(alg, JWEKeyAgreement):
|
|
# 3. When Direct Key Agreement is employed,
|
|
# let the CEK be the agreed upon key.
|
|
cek = alg.encrypt_agreed_upon_key(enc, recipient)
|
|
if len(cek) * 8 != enc.cek_size: # pragma: no cover
|
|
cek = secrets.token_bytes(enc.cek_size // 8)
|
|
else:
|
|
# 6. When Direct Encryption is employed, let the CEK be the shared
|
|
# symmetric key.
|
|
assert isinstance(alg, JWEDirectEncryption)
|
|
cek = alg.compute_cek(enc.cek_size, recipient)
|
|
|
|
# 5. When Direct Key Agreement or Direct Encryption are employed, let
|
|
# the JWE Encrypted Key be the empty octet sequence.
|
|
recipient.encrypted_key = b""
|
|
return cek
|
|
|
|
|
|
def post_encrypt_recipients(
|
|
enc: JWEEncModel, tasks: list[tuple[JWEKeyAgreement, Recipient[Any]]], cek: bytes, tag: bytes
|
|
) -> None:
|
|
for alg, recipient in tasks:
|
|
if alg.tag_aware:
|
|
agreed_upon_key = alg.encrypt_agreed_upon_key_with_tag(enc, recipient, tag)
|
|
else:
|
|
agreed_upon_key = alg.encrypt_agreed_upon_key(enc, recipient)
|
|
# 4. When Key Agreement with Key Wrapping is employed, encrypt the CEK
|
|
# to the recipient and let the result be the JWE Encrypted Key.
|
|
recipient.encrypted_key = alg.wrap_cek_with_auk(cek, agreed_upon_key)
|
|
|
|
|
|
def decrypt_recipient(alg: JWEAlgModel, enc: JWEEncModel, recipient: Recipient[Any], tag: bytes) -> bytes:
|
|
cek: bytes
|
|
if alg.direct_mode:
|
|
# 10. When Direct Key Agreement or Direct Encryption are employed,
|
|
# verify that the JWE Encrypted Key value is an empty octet
|
|
# sequence.
|
|
if recipient.encrypted_key: # pragma: no cover
|
|
raise InvalidEncryptedKeyError()
|
|
|
|
if isinstance(alg, JWEKeyAgreement):
|
|
# 8. When Direct Key Agreement is employed, let the CEK be the agreed upon key.
|
|
cek = alg.decrypt_agreed_upon_key(enc, recipient)
|
|
else:
|
|
# 11. When Direct Encryption is employed, let the CEK be the shared
|
|
# symmetric key.
|
|
assert isinstance(alg, JWEDirectEncryption)
|
|
cek = alg.compute_cek(enc.cek_size, recipient)
|
|
elif isinstance(alg, JWEKeyAgreement):
|
|
agreed_upon_key: bytes
|
|
if alg.tag_aware:
|
|
agreed_upon_key = alg.decrypt_agreed_upon_key_with_tag(enc, recipient, tag)
|
|
else:
|
|
agreed_upon_key = alg.decrypt_agreed_upon_key(enc, recipient)
|
|
|
|
# 8. When Key Agreement with Key Wrapping is employed, the agreed upon key
|
|
# will be used to decrypt the JWE Encrypted Key.
|
|
if recipient.encrypted_key is None: # pragma: no cover
|
|
raise DecodeError("Invalid encrypted key")
|
|
cek = alg.unwrap_cek_with_auk(recipient.encrypted_key, agreed_upon_key)
|
|
else:
|
|
# 9. When Key Wrapping, Key Encryption, or Key Agreement with Key
|
|
# Wrapping are employed, decrypt the JWE Encrypted Key to produce
|
|
# the CEK. The CEK MUST have a length equal to that required for
|
|
# the content encryption algorithm. Note that when there are
|
|
# multiple recipients, each recipient will only be able to decrypt
|
|
# JWE Encrypted Key values that were encrypted to a key in that
|
|
# recipient's possession. It is therefore normal to only be able
|
|
# to decrypt one of the per-recipient JWE Encrypted Key values to
|
|
# obtain the CEK value.
|
|
assert isinstance(alg, (JWEKeyWrapping, JWEKeyEncryption))
|
|
if recipient.encrypted_key is None:
|
|
raise DecodeError("Invalid encrypted key")
|
|
cek = alg.decrypt_cek(recipient)
|
|
return cek
|