first commit
This commit is contained in:
linkitapps
@@ -0,0 +1,75 @@
|
||||
from .models import CompactEncryption, Recipient
|
||||
from .registry import JWERegistry
|
||||
from .._keys import Key
|
||||
from ..errors import (
|
||||
MissingAlgorithmError,
|
||||
MissingEncryptionError,
|
||||
DecodeError,
|
||||
)
|
||||
from ..util import (
|
||||
json_b64decode,
|
||||
urlsafe_b64encode,
|
||||
urlsafe_b64decode,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"represent_compact",
|
||||
"extract_compact",
|
||||
]
|
||||
|
||||
|
||||
def represent_compact(obj: CompactEncryption) -> bytes:
|
||||
assert obj.recipient is not None
|
||||
encrypted_key = obj.recipient.encrypted_key
|
||||
assert encrypted_key is not None
|
||||
return b".".join(
|
||||
[
|
||||
obj.base64_segments["aad"],
|
||||
urlsafe_b64encode(encrypted_key),
|
||||
obj.base64_segments["iv"],
|
||||
obj.base64_segments["ciphertext"],
|
||||
obj.base64_segments["tag"],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def extract_compact(value: bytes, registry: JWERegistry) -> CompactEncryption:
|
||||
parts = value.split(b".")
|
||||
if len(parts) != 5:
|
||||
raise ValueError("Invalid JSON Web Encryption")
|
||||
|
||||
header_segment, ek_segment, iv_segment, ciphertext_segment, tag_segment = parts
|
||||
registry.validate_protected_header_size(header_segment)
|
||||
registry.validate_encrypted_key_size(ek_segment)
|
||||
registry.validate_initialization_vector_size(iv_segment)
|
||||
registry.validate_ciphertext_size(ciphertext_segment)
|
||||
registry.validate_auth_tag_size(tag_segment)
|
||||
try:
|
||||
protected = json_b64decode(header_segment)
|
||||
if "alg" not in protected:
|
||||
raise MissingAlgorithmError()
|
||||
if "enc" not in protected:
|
||||
raise MissingEncryptionError()
|
||||
except (TypeError, ValueError):
|
||||
raise DecodeError("Invalid header")
|
||||
|
||||
obj = CompactEncryption(protected)
|
||||
obj.base64_segments.update(
|
||||
{
|
||||
"aad": header_segment,
|
||||
"iv": iv_segment,
|
||||
"ciphertext": ciphertext_segment,
|
||||
"tag": tag_segment,
|
||||
}
|
||||
)
|
||||
obj.bytes_segments.update(
|
||||
{
|
||||
"iv": urlsafe_b64decode(iv_segment),
|
||||
"ciphertext": urlsafe_b64decode(ciphertext_segment),
|
||||
"tag": urlsafe_b64decode(tag_segment),
|
||||
}
|
||||
)
|
||||
recipient: Recipient[Key] = Recipient(obj)
|
||||
recipient.encrypted_key = urlsafe_b64decode(ek_segment)
|
||||
obj.recipient = recipient
|
||||
return obj
|
||||
@@ -0,0 +1,142 @@
|
||||
from typing import Any
|
||||
from .models import (
|
||||
BaseJSONEncryption,
|
||||
GeneralJSONEncryption,
|
||||
FlattenedJSONEncryption,
|
||||
Recipient,
|
||||
)
|
||||
from .registry import JWERegistry
|
||||
from .types import (
|
||||
JSONRecipientDict,
|
||||
GeneralJSONSerialization,
|
||||
FlattenedJSONSerialization,
|
||||
)
|
||||
from ..util import (
|
||||
to_bytes,
|
||||
to_str,
|
||||
json_b64encode,
|
||||
json_b64decode,
|
||||
urlsafe_b64encode,
|
||||
urlsafe_b64decode,
|
||||
)
|
||||
from .._keys import Key
|
||||
|
||||
|
||||
__all__ = [
|
||||
"represent_general_json",
|
||||
"represent_flattened_json",
|
||||
"extract_general_json",
|
||||
"extract_flattened_json",
|
||||
]
|
||||
|
||||
|
||||
def represent_general_json(obj: GeneralJSONEncryption) -> GeneralJSONSerialization:
|
||||
data: GeneralJSONSerialization = __represent_json_serialization(obj)
|
||||
recipients = []
|
||||
for recipient in obj.recipients:
|
||||
item: JSONRecipientDict = {}
|
||||
assert recipient.header is not None
|
||||
assert recipient.encrypted_key is not None
|
||||
item["header"] = recipient.header
|
||||
item["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
|
||||
recipients.append(item)
|
||||
data["recipients"] = recipients
|
||||
return data
|
||||
|
||||
|
||||
def represent_flattened_json(obj: FlattenedJSONEncryption) -> FlattenedJSONSerialization:
|
||||
data: FlattenedJSONSerialization = __represent_json_serialization(obj)
|
||||
recipient = obj.recipients[0]
|
||||
assert recipient is not None
|
||||
assert recipient.header is not None
|
||||
assert recipient.encrypted_key is not None
|
||||
data["header"] = recipient.header
|
||||
data["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
|
||||
return data
|
||||
|
||||
|
||||
def __represent_json_serialization(obj: BaseJSONEncryption) -> Any:
|
||||
data: dict[str, Any] = {
|
||||
"protected": to_str(json_b64encode(obj.protected)),
|
||||
"iv": to_str(obj.base64_segments["iv"]),
|
||||
"ciphertext": to_str(obj.base64_segments["ciphertext"]),
|
||||
"tag": to_str(obj.base64_segments["tag"]),
|
||||
}
|
||||
if obj.aad:
|
||||
data["aad"] = to_str(urlsafe_b64encode(obj.aad))
|
||||
|
||||
if obj.unprotected:
|
||||
data["unprotected"] = obj.unprotected
|
||||
return data
|
||||
|
||||
|
||||
def extract_general_json(data: GeneralJSONSerialization, registry: JWERegistry) -> GeneralJSONEncryption:
|
||||
protected_segment = to_bytes(data["protected"])
|
||||
registry.validate_protected_header_size(protected_segment)
|
||||
protected = json_b64decode(protected_segment)
|
||||
|
||||
unprotected = data.get("unprotected")
|
||||
recipients = data["recipients"]
|
||||
registry.validate_recipient_count(recipients)
|
||||
base64_segments, bytes_segments, aad = __extract_segments(data, registry)
|
||||
|
||||
obj = GeneralJSONEncryption(protected, None, unprotected, aad)
|
||||
obj.base64_segments = base64_segments
|
||||
obj.bytes_segments = bytes_segments
|
||||
for item in recipients:
|
||||
recipient = __extract_recipient(obj, item, registry)
|
||||
obj.recipients.append(recipient)
|
||||
return obj
|
||||
|
||||
|
||||
def extract_flattened_json(data: FlattenedJSONSerialization, registry: JWERegistry) -> FlattenedJSONEncryption:
|
||||
protected_segment = to_bytes(data["protected"])
|
||||
registry.validate_protected_header_size(protected_segment)
|
||||
protected = json_b64decode(protected_segment)
|
||||
unprotected = data.get("unprotected")
|
||||
base64_segments, bytes_segments, aad = __extract_segments(data, registry)
|
||||
obj = FlattenedJSONEncryption(protected, None, unprotected, aad)
|
||||
obj.base64_segments = base64_segments
|
||||
obj.bytes_segments = bytes_segments
|
||||
recipient = __extract_recipient(obj, data, registry)
|
||||
obj.recipients.append(recipient)
|
||||
return obj
|
||||
|
||||
|
||||
def __extract_segments(
|
||||
data: GeneralJSONSerialization | FlattenedJSONSerialization,
|
||||
registry: JWERegistry,
|
||||
) -> tuple[dict[str, bytes], dict[str, bytes], bytes | None]:
|
||||
base64_segments: dict[str, bytes] = {
|
||||
"iv": to_bytes(data["iv"]),
|
||||
"ciphertext": to_bytes(data["ciphertext"]),
|
||||
"tag": to_bytes(data["tag"]),
|
||||
}
|
||||
registry.validate_initialization_vector_size(base64_segments["iv"])
|
||||
registry.validate_ciphertext_size(base64_segments["ciphertext"])
|
||||
registry.validate_auth_tag_size(base64_segments["tag"])
|
||||
bytes_segments: dict[str, bytes] = {
|
||||
"iv": urlsafe_b64decode(base64_segments["iv"]),
|
||||
"ciphertext": urlsafe_b64decode(base64_segments["ciphertext"]),
|
||||
"tag": urlsafe_b64decode(base64_segments["tag"]),
|
||||
}
|
||||
if "aad" in data:
|
||||
base64_segments["aad"] = to_bytes(data["aad"])
|
||||
aad = urlsafe_b64decode(base64_segments["aad"])
|
||||
bytes_segments["aad"] = aad
|
||||
else:
|
||||
aad = None
|
||||
return base64_segments, bytes_segments, aad
|
||||
|
||||
|
||||
def __extract_recipient(
|
||||
obj: FlattenedJSONEncryption | GeneralJSONEncryption,
|
||||
data: FlattenedJSONSerialization | JSONRecipientDict,
|
||||
registry: JWERegistry,
|
||||
) -> Recipient[Key]:
|
||||
recipient: Recipient[Key] = Recipient(obj, data.get("header"))
|
||||
if "encrypted_key" in data:
|
||||
ek_segment = to_bytes(data["encrypted_key"])
|
||||
registry.validate_encrypted_key_size(ek_segment)
|
||||
recipient.encrypted_key = urlsafe_b64decode(ek_segment)
|
||||
return recipient
|
||||
@@ -0,0 +1,266 @@
|
||||
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
|
||||
@@ -0,0 +1,345 @@
|
||||
import typing as t
|
||||
import secrets
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from ..registry import Header, HeaderRegistryDict
|
||||
from ..errors import InvalidKeyTypeError, InvalidKeyLengthError
|
||||
from .._keys import Key, ECKey, OctKey
|
||||
|
||||
__all__ = [
|
||||
"Recipient",
|
||||
"CompactEncryption",
|
||||
"BaseJSONEncryption",
|
||||
"GeneralJSONEncryption",
|
||||
"FlattenedJSONEncryption",
|
||||
"JWEEncModel",
|
||||
"JWEZipModel",
|
||||
"KeyManagement",
|
||||
"JWEDirectEncryption",
|
||||
"JWEKeyEncryption",
|
||||
"JWEKeyWrapping",
|
||||
"JWEKeyAgreement",
|
||||
"JWEAlgModel",
|
||||
]
|
||||
|
||||
KeyType = t.TypeVar("KeyType")
|
||||
|
||||
|
||||
class Recipient(t.Generic[KeyType]):
|
||||
def __init__(
|
||||
self,
|
||||
parent: t.Union["CompactEncryption", "GeneralJSONEncryption", "FlattenedJSONEncryption"],
|
||||
header: Header | None = None,
|
||||
recipient_key: KeyType | None = None,
|
||||
):
|
||||
self.__parent = parent
|
||||
self.header = header
|
||||
self.recipient_key = recipient_key
|
||||
self.sender_key: KeyType | None = None
|
||||
self.encrypted_key: bytes | None = None
|
||||
self.ephemeral_key: KeyType | None = None
|
||||
|
||||
def headers(self) -> Header:
|
||||
rv: Header = {}
|
||||
if isinstance(self.__parent, BaseJSONEncryption) and self.__parent.unprotected:
|
||||
rv.update(self.__parent.unprotected)
|
||||
if self.header:
|
||||
rv.update(self.header)
|
||||
|
||||
rv.update(self.__parent.protected)
|
||||
return rv
|
||||
|
||||
def add_header(self, k: str, v: t.Any) -> None:
|
||||
if isinstance(self.__parent, CompactEncryption):
|
||||
self.__parent.protected.update({k: v})
|
||||
elif self.header:
|
||||
self.header.update({k: v})
|
||||
else:
|
||||
self.header = {k: v}
|
||||
|
||||
def set_kid(self, kid: str) -> None:
|
||||
self.add_header("kid", kid)
|
||||
|
||||
|
||||
class CompactEncryption:
|
||||
"""An object to represent the JWE Compact Serialization. It is usually returned by
|
||||
``decrypt_compact`` method.
|
||||
"""
|
||||
|
||||
def __init__(self, protected: Header, plaintext: bytes | None = None):
|
||||
#: protected header in dict
|
||||
self.protected = protected
|
||||
#: the plaintext in bytes
|
||||
self.plaintext = plaintext
|
||||
self.recipient: Recipient[t.Any] | None = None
|
||||
self.bytes_segments: dict[str, bytes] = {} # store the decoded segments
|
||||
self.base64_segments: dict[str, bytes] = {} # store the encoded segments
|
||||
|
||||
def headers(self) -> Header:
|
||||
"""Returns the protected header values in dict."""
|
||||
return self.protected
|
||||
|
||||
def attach_recipient(self, key: Key, header: Header | None = None) -> None:
|
||||
"""Add a recipient to the JWE Compact Serialization. Please add a key that
|
||||
comply with the given "alg" value.
|
||||
|
||||
:param key: an instance of a key, e.g. (OctKey, RSAKey, ECKey, and etc)
|
||||
:param header: extra header in dict
|
||||
"""
|
||||
recipient = Recipient(self, None, key)
|
||||
if header:
|
||||
self.protected.update(header)
|
||||
self.recipient = recipient
|
||||
|
||||
@property
|
||||
def recipients(self) -> list[Recipient[t.Any]]:
|
||||
if self.recipient is not None:
|
||||
return [self.recipient]
|
||||
return []
|
||||
|
||||
|
||||
class BaseJSONEncryption(metaclass=ABCMeta):
|
||||
#: represents if the object is in flatten syntax
|
||||
flattened: t.ClassVar[bool]
|
||||
#: protected header in dict
|
||||
protected: Header
|
||||
#: the plaintext in bytes
|
||||
plaintext: t.Optional[bytes]
|
||||
#: unprotected header in dict
|
||||
unprotected: t.Optional[Header]
|
||||
#: an optional additional authenticated data
|
||||
aad: t.Optional[bytes]
|
||||
#: a list of recipients
|
||||
recipients: list[Recipient[t.Any]]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
protected: Header,
|
||||
plaintext: bytes | None = None,
|
||||
unprotected: Header | None = None,
|
||||
aad: bytes | None = None,
|
||||
):
|
||||
self.protected = protected
|
||||
self.plaintext = plaintext
|
||||
self.unprotected = unprotected
|
||||
self.aad = aad
|
||||
self.recipients = []
|
||||
self.bytes_segments: dict[str, bytes] = {} # store the decoded segments
|
||||
self.base64_segments: dict[str, bytes] = {} # store the encoded segments
|
||||
|
||||
@abstractmethod
|
||||
def add_recipient(self, header: Header | None = None, key: Key | None = None) -> None:
|
||||
"""Add a recipient to the JWE JSON Serialization. Please add a key that
|
||||
comply with the "alg" to this recipient.
|
||||
|
||||
:param header: recipient's own (unprotected) header
|
||||
:param key: an instance of a key, e.g. (OctKey, RSAKey, ECKey, and etc)
|
||||
"""
|
||||
|
||||
|
||||
class GeneralJSONEncryption(BaseJSONEncryption):
|
||||
"""An object to represent the JWE General JSON Serialization. It is used by
|
||||
``encrypt_json``, and it is usually returned by ``decrypt_json`` method.
|
||||
|
||||
To construct an object of ``GeneralJSONEncryption``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
protected = {"enc": "A128CBC-HS256"}
|
||||
plaintext = b"hello world"
|
||||
obj = GeneralJSONEncryption(protected, plaintext)
|
||||
# then add each recipient
|
||||
obj.add_recipient({"alg": "A128KW"})
|
||||
"""
|
||||
|
||||
flattened = False
|
||||
|
||||
def add_recipient(self, header: Header | None = None, key: Key | None = None) -> None:
|
||||
recipient = Recipient(self, header, key)
|
||||
if key and key.kid:
|
||||
recipient.set_kid(key.kid)
|
||||
self.recipients.append(recipient)
|
||||
|
||||
|
||||
class FlattenedJSONEncryption(BaseJSONEncryption):
|
||||
"""An object to represent the JWE Flattened JSON Serialization. It is used by
|
||||
``encrypt_json``, and it is usually returned by ``decrypt_json`` method.
|
||||
|
||||
To construct an object of ``FlattenedJSONEncryption``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
protected = {"enc": "A128CBC-HS256"}
|
||||
plaintext = b"hello world"
|
||||
obj = FlattenedJSONEncryption(protected, plaintext)
|
||||
# then add each recipient
|
||||
obj.add_recipient({"alg": "A128KW"})
|
||||
"""
|
||||
|
||||
flattened = True
|
||||
|
||||
def add_recipient(self, header: Header | None = None, key: Key | None = None) -> None:
|
||||
recipient = Recipient(self, header, key)
|
||||
if key and key.kid:
|
||||
recipient.set_kid(key.kid)
|
||||
self.recipients = [recipient]
|
||||
|
||||
|
||||
class JWEEncModel(metaclass=ABCMeta):
|
||||
name: str
|
||||
description: str
|
||||
recommended: bool = False
|
||||
algorithm_type: t.Literal["JWE"] = "JWE"
|
||||
algorithm_location: t.Literal["enc"] = "enc"
|
||||
|
||||
iv_size: int
|
||||
cek_size: int
|
||||
|
||||
def generate_cek(self) -> bytes:
|
||||
return secrets.token_bytes(self.cek_size // 8)
|
||||
|
||||
def generate_iv(self) -> bytes:
|
||||
return secrets.token_bytes(self.iv_size // 8)
|
||||
|
||||
def check_iv(self, iv: bytes) -> bytes:
|
||||
if len(iv) * 8 != self.iv_size: # pragma: no cover
|
||||
raise ValueError("Invalid 'iv' size")
|
||||
return iv
|
||||
|
||||
@abstractmethod
|
||||
def encrypt(self, plaintext: bytes, cek: bytes, iv: bytes, aad: bytes) -> tuple[bytes, bytes]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def decrypt(self, ciphertext: bytes, tag: bytes, cek: bytes, iv: bytes, aad: bytes) -> bytes:
|
||||
pass
|
||||
|
||||
|
||||
class JWEZipModel(metaclass=ABCMeta):
|
||||
name: str
|
||||
description: str
|
||||
recommended: bool = True
|
||||
algorithm_type: t.Literal["JWE"] = "JWE"
|
||||
algorithm_location: t.Literal["zip"] = "zip"
|
||||
|
||||
@abstractmethod
|
||||
def compress(self, s: bytes) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def decompress(self, s: bytes) -> bytes:
|
||||
pass
|
||||
|
||||
|
||||
class KeyManagement:
|
||||
name: str
|
||||
description: str
|
||||
recommended: bool = False
|
||||
key_size: int | None = None
|
||||
key_types: list[str]
|
||||
security_warning: str | None = None
|
||||
|
||||
algorithm_type: t.Literal["JWE"] = "JWE"
|
||||
algorithm_location: t.Literal["alg"] = "alg"
|
||||
more_header_registry: HeaderRegistryDict = {}
|
||||
|
||||
@property
|
||||
def direct_mode(self) -> bool:
|
||||
return self.key_size is None
|
||||
|
||||
def check_key_type(self, key: Key) -> None:
|
||||
if key.key_type not in self.key_types:
|
||||
raise InvalidKeyTypeError()
|
||||
|
||||
def prepare_recipient_header(self, recipient: Recipient[t.Any]) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class JWEDirectEncryption(KeyManagement, metaclass=ABCMeta):
|
||||
key_types = ["oct"]
|
||||
|
||||
@abstractmethod
|
||||
def compute_cek(self, size: int, recipient: Recipient[OctKey]) -> bytes:
|
||||
pass
|
||||
|
||||
|
||||
class JWEKeyEncryption(KeyManagement, metaclass=ABCMeta):
|
||||
@property
|
||||
def direct_mode(self) -> bool:
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def encrypt_cek(self, cek: bytes, recipient: Recipient[t.Any]) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def decrypt_cek(self, recipient: Recipient[t.Any]) -> bytes:
|
||||
pass
|
||||
|
||||
|
||||
class JWEKeyWrapping(KeyManagement, metaclass=ABCMeta):
|
||||
key_size: int
|
||||
key_types = ["oct"]
|
||||
|
||||
@property
|
||||
def direct_mode(self) -> bool:
|
||||
return False
|
||||
|
||||
def check_op_key(self, op_key: bytes) -> None:
|
||||
if len(op_key) * 8 != self.key_size:
|
||||
raise InvalidKeyLengthError(f"A key of size {self.key_size} bits MUST be used")
|
||||
|
||||
@abstractmethod
|
||||
def wrap_cek(self, cek: bytes, key: bytes) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def unwrap_cek(self, ek: bytes, key: bytes) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def encrypt_cek(self, cek: bytes, recipient: Recipient[OctKey]) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def decrypt_cek(self, recipient: Recipient[OctKey]) -> bytes:
|
||||
pass
|
||||
|
||||
|
||||
class JWEKeyAgreement(KeyManagement, metaclass=ABCMeta):
|
||||
key_types = ["EC", "OKP"]
|
||||
tag_aware: bool = False
|
||||
key_wrapping: t.Optional[JWEKeyWrapping]
|
||||
|
||||
def prepare_ephemeral_key(self, recipient: Recipient[ECKey]) -> None:
|
||||
recipient_key = recipient.recipient_key
|
||||
assert recipient_key is not None
|
||||
self.check_key_type(recipient_key)
|
||||
if recipient.ephemeral_key is None:
|
||||
ephemeral_key = recipient_key.generate_key(recipient_key.curve_name, private=True)
|
||||
recipient.ephemeral_key = ephemeral_key
|
||||
recipient.add_header("epk", recipient.ephemeral_key.as_dict(private=False))
|
||||
|
||||
@abstractmethod
|
||||
def encrypt_agreed_upon_key(self, enc: JWEEncModel, recipient: Recipient[ECKey]) -> bytes:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def decrypt_agreed_upon_key(self, enc: JWEEncModel, recipient: Recipient[ECKey]) -> bytes:
|
||||
pass
|
||||
|
||||
def wrap_cek_with_auk(self, cek: bytes, key: bytes) -> bytes:
|
||||
assert self.key_wrapping is not None
|
||||
return self.key_wrapping.wrap_cek(cek, key)
|
||||
|
||||
def unwrap_cek_with_auk(self, ek: bytes, key: bytes) -> bytes:
|
||||
assert self.key_wrapping is not None
|
||||
return self.key_wrapping.unwrap_cek(ek, key)
|
||||
|
||||
def encrypt_agreed_upon_key_with_tag(self, enc: JWEEncModel, recipient: Recipient[ECKey], tag: bytes) -> bytes:
|
||||
raise NotImplementedError()
|
||||
|
||||
def decrypt_agreed_upon_key_with_tag(self, enc: JWEEncModel, recipient: Recipient[ECKey], tag: bytes) -> bytes:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
JWEAlgModel = t.Union[JWEKeyEncryption, JWEKeyWrapping, JWEKeyAgreement, JWEDirectEncryption]
|
||||
@@ -0,0 +1,177 @@
|
||||
import warnings
|
||||
import typing as t
|
||||
from collections.abc import Collection
|
||||
from .models import JWEAlgModel, JWEEncModel, JWEZipModel
|
||||
from ..errors import (
|
||||
UnsupportedAlgorithmError,
|
||||
SecurityWarning,
|
||||
ExceededSizeError,
|
||||
)
|
||||
from ..registry import (
|
||||
Header,
|
||||
HeaderRegistryDict,
|
||||
JWE_HEADER_REGISTRY,
|
||||
check_supported_header,
|
||||
check_registry_header,
|
||||
check_crit_header,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"JWEAlgorithm",
|
||||
"JWERegistry",
|
||||
"default_registry",
|
||||
]
|
||||
|
||||
JWEAlgorithm = t.Union[JWEAlgModel, JWEEncModel, JWEZipModel]
|
||||
|
||||
AlgorithmsDict = t.TypedDict(
|
||||
"AlgorithmsDict",
|
||||
{
|
||||
"alg": dict[str, JWEAlgModel],
|
||||
"enc": dict[str, JWEEncModel],
|
||||
"zip": dict[str, JWEZipModel],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class JWERegistry:
|
||||
"""A registry for JSON Web Encryption to keep all the supported algorithms.
|
||||
An instance of ``JWERegistry`` is usually used together with methods in
|
||||
``joserfc.jwe``.
|
||||
|
||||
:param header_registry: extra header parameters registry
|
||||
:param algorithms: allowed algorithms to be used
|
||||
:param verify_all_recipients: validating all recipients in a JSON serialization
|
||||
:param strict_check_header: only allow header key in the registry to be used
|
||||
:param max_recipients: max number of recipients in JWE General JSON Serialization
|
||||
"""
|
||||
|
||||
algorithms: AlgorithmsDict = {
|
||||
"alg": {},
|
||||
"enc": {},
|
||||
"zip": {},
|
||||
}
|
||||
recommended: t.ClassVar[list[str]] = []
|
||||
|
||||
#: max protected header content's size in bytes
|
||||
max_protected_header_length: int = 1024
|
||||
#: max encrypted key's size in bytes
|
||||
max_encrypted_key_length: int = 1024
|
||||
#: max initialization vector's size in bytes
|
||||
max_initialization_vector_length: int = 64
|
||||
#: max ciphertext's size in bytes
|
||||
max_ciphertext_length: int = 65536 # 64KB
|
||||
#: max auth tag's size in bytes
|
||||
max_auth_tag_length: int = 64
|
||||
#: max number of recipients in JWE General JSON Serialization
|
||||
max_recipients: int = 16
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
header_registry: HeaderRegistryDict | None = None,
|
||||
algorithms: Collection[str] | None = None,
|
||||
verify_all_recipients: bool = True,
|
||||
strict_check_header: bool = True,
|
||||
max_recipients: int | None = None,
|
||||
):
|
||||
self.header_registry: HeaderRegistryDict = {}
|
||||
self.header_registry.update(JWE_HEADER_REGISTRY)
|
||||
if header_registry is not None:
|
||||
self.header_registry.update(header_registry)
|
||||
self.allowed = algorithms
|
||||
self.verify_all_recipients = verify_all_recipients
|
||||
self.strict_check_header = strict_check_header
|
||||
if max_recipients is not None:
|
||||
self.max_recipients = max_recipients
|
||||
|
||||
@classmethod
|
||||
def register(cls, model: JWEAlgorithm) -> None:
|
||||
cls.algorithms[model.algorithm_location][model.name] = model # type: ignore
|
||||
if model.recommended:
|
||||
cls.recommended.append(model.name)
|
||||
|
||||
def check_header(self, header: Header, check_more: bool = False) -> None:
|
||||
"""Check and validate the fields in header part of a JWS object."""
|
||||
check_crit_header(self.header_registry, header)
|
||||
check_registry_header(self.header_registry, header)
|
||||
|
||||
alg = self.get_alg(header["alg"])
|
||||
if alg.more_header_registry:
|
||||
check_registry_header(alg.more_header_registry, header, check_more)
|
||||
|
||||
if self.strict_check_header:
|
||||
allowed_registry = self.header_registry.copy()
|
||||
allowed_registry.update(alg.more_header_registry)
|
||||
check_supported_header(allowed_registry, header)
|
||||
elif self.strict_check_header:
|
||||
check_supported_header(self.header_registry, header)
|
||||
|
||||
def validate_protected_header_size(self, header: bytes) -> None:
|
||||
if header and len(header) > self.max_protected_header_length:
|
||||
raise ExceededSizeError(f"Header size exceeds {self.max_protected_header_length} bytes.")
|
||||
|
||||
def validate_encrypted_key_size(self, ek: bytes) -> None:
|
||||
if ek and len(ek) > self.max_encrypted_key_length:
|
||||
raise ExceededSizeError(f"Encrypted key size exceeds {self.max_encrypted_key_length} bytes.")
|
||||
|
||||
def validate_initialization_vector_size(self, iv: bytes) -> None:
|
||||
if iv and len(iv) > self.max_initialization_vector_length:
|
||||
raise ExceededSizeError(
|
||||
f"Initialization vector size exceeds {self.max_initialization_vector_length} bytes."
|
||||
)
|
||||
|
||||
def validate_ciphertext_size(self, ciphertext: bytes) -> None:
|
||||
if ciphertext and len(ciphertext) > self.max_ciphertext_length:
|
||||
raise ExceededSizeError(f"Ciphertext size exceeds {self.max_ciphertext_length} bytes.")
|
||||
|
||||
def validate_auth_tag_size(self, tag: bytes) -> None:
|
||||
if tag and len(tag) > self.max_auth_tag_length:
|
||||
raise ExceededSizeError(f"Auth tag size exceeds {self.max_auth_tag_length} bytes.")
|
||||
|
||||
def validate_recipient_count(self, recipients: Collection[t.Any]) -> None:
|
||||
if len(recipients) > self.max_recipients:
|
||||
raise ExceededSizeError(f"Recipient count exceeds {self.max_recipients}.")
|
||||
|
||||
def get_alg(self, name: str) -> JWEAlgModel:
|
||||
"""Get the allowed ("alg") algorithm instance of the given name.
|
||||
|
||||
:param name: value of the ``alg``, e.g. ``ECDH-ES``, ``A128KW``
|
||||
"""
|
||||
registry = self.algorithms["alg"]
|
||||
self._check_algorithm(name, registry)
|
||||
alg: JWEAlgModel = registry[name]
|
||||
if alg.security_warning:
|
||||
warnings.warn(alg.security_warning, SecurityWarning)
|
||||
return alg
|
||||
|
||||
def get_enc(self, name: str) -> JWEEncModel:
|
||||
"""Get the allowed ("enc") algorithm instance of the given name.
|
||||
|
||||
:param name: value of the ``enc``, e.g. ``A128CBC-HS256``, ``A128GCM``
|
||||
"""
|
||||
registry = self.algorithms["enc"]
|
||||
self._check_algorithm(name, registry)
|
||||
return registry[name]
|
||||
|
||||
def get_zip(self, name: str) -> JWEZipModel:
|
||||
"""Get the allowed ("zip") algorithm instance of the given name.
|
||||
|
||||
:param name: value of the ``zip``, e.g. ``DEF``
|
||||
"""
|
||||
registry = self.algorithms["zip"]
|
||||
self._check_algorithm(name, registry)
|
||||
return registry[name]
|
||||
|
||||
def _check_algorithm(self, name: str, registry: dict[str, t.Any]) -> None:
|
||||
if name not in registry:
|
||||
raise UnsupportedAlgorithmError(f"Algorithm of '{name}' is not supported")
|
||||
|
||||
if self.allowed:
|
||||
if name not in self.allowed:
|
||||
raise UnsupportedAlgorithmError(f"Algorithm of '{name}' is not allowed")
|
||||
else:
|
||||
if name not in self.recommended:
|
||||
raise UnsupportedAlgorithmError(f"Algorithm of '{name}' is not recommended")
|
||||
|
||||
|
||||
default_registry = JWERegistry()
|
||||
@@ -0,0 +1,33 @@
|
||||
from typing import Any, TypedDict
|
||||
|
||||
__all__ = [
|
||||
"JSONRecipientDict",
|
||||
"FlattenedJSONSerialization",
|
||||
"GeneralJSONSerialization",
|
||||
]
|
||||
|
||||
|
||||
class JSONRecipientDict(TypedDict, total=False):
|
||||
header: dict[str, Any]
|
||||
encrypted_key: str
|
||||
|
||||
|
||||
class GeneralJSONSerialization(TypedDict, total=False):
|
||||
protected: str
|
||||
unprotected: dict[str, Any]
|
||||
iv: str
|
||||
aad: str
|
||||
ciphertext: str
|
||||
tag: str
|
||||
recipients: list[JSONRecipientDict]
|
||||
|
||||
|
||||
class FlattenedJSONSerialization(TypedDict, total=False):
|
||||
protected: str
|
||||
unprotected: dict[str, Any]
|
||||
iv: str
|
||||
aad: str
|
||||
ciphertext: str
|
||||
tag: str
|
||||
header: dict[str, Any]
|
||||
encrypted_key: str
|
||||
Reference in New Issue
Block a user