346 lines
11 KiB
Python
346 lines
11 KiB
Python
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]
|