from __future__ import annotations
from abc import ABC, abstractmethod
from collections import OrderedDict
import copy
import itertools
import json
from typing import Optional, Tuple, Type, TypeVar, cast
from .aead import AEAD
from .diffie_hellman_ratchet import DiffieHellmanRatchet
from .kdf import KDF
from .migrations import parse_double_ratchet_model
from .models import DoubleRatchetModel, SkippedMessageKeyModel
from .types import EncryptedMessage, Header, JSONObject, SkippedMessageKeys
__all__ = [
"DoubleRatchet"
]
DoubleRatchetTypeT = TypeVar("DoubleRatchetTypeT", bound="DoubleRatchet")
[docs]
class DoubleRatchet(ABC):
"""
Combining the symmetric-key ratchet and the Diffie-Hellman ratchet gives the Double Ratchet.
https://signal.org/docs/specifications/doubleratchet/#double-ratchet
Note:
In this implementation, the Diffie-Hellman ratchet already manages the symmetric-key ratchet
internally, see :class:`~doubleratchet.diffie_hellman_ratchet.DiffieHellmanRatchet` for details. The
Double Ratchet class adds message en-/decryption and offers a more convenient public API that handles
lost and out-of-order messages.
"""
def __init__(self) -> None:
# Just the type definitions here
self.__max_num_skipped_message_keys: int
self.__skipped_message_keys: SkippedMessageKeys
self.__aead: Type[AEAD]
self.__diffie_hellman_ratchet: DiffieHellmanRatchet
[docs]
@classmethod
async def encrypt_initial_message(
cls: Type[DoubleRatchetTypeT],
diffie_hellman_ratchet_class: Type[DiffieHellmanRatchet],
root_chain_kdf: Type[KDF],
message_chain_kdf: Type[KDF],
message_chain_constant: bytes,
dos_protection_threshold: int,
max_num_skipped_message_keys: int,
aead: Type[AEAD],
shared_secret: bytes,
recipient_ratchet_pub: bytes,
message: bytes,
associated_data: bytes
) -> Tuple[DoubleRatchetTypeT, EncryptedMessage]:
"""
Args:
diffie_hellman_ratchet_class: A non-abstract subclass of
:class:`~doubleratchet.diffie_hellman_ratchet.DiffieHellmanRatchet`.
root_chain_kdf: The KDF to use for the root chain. The KDF must be capable of deriving 64 bytes.
message_chain_kdf: The KDF to use for the sending and receiving chains. The KDF must be capable of
deriving 64 bytes.
message_chain_constant: The constant to feed into the sending and receiving KDF chains on each
step.
dos_protection_threshold: The maximum number of skipped message keys to calculate. If more than
that number of message keys are skipped, the keys are not calculated to prevent being DoSed.
max_num_skipped_message_keys: The maximum number of skipped message keys to store in case the lost
or out-of-order message comes in later. Older keys are discarded to make space for newer keys.
aead: The AEAD implementation to use for message en- and decryption.
shared_secret: A shared secret consisting of 32 bytes that was agreed on by means external to this
protocol.
recipient_ratchet_pub: The ratchet public key of the recipient.
message: The initial message.
associated_data: Additional data to authenticate without including it in the ciphertext.
Returns:
A configured instance of :class:`DoubleRatchet` ready to send and receive messages together with
the initial message.
"""
if dos_protection_threshold > max_num_skipped_message_keys:
raise ValueError(
"The `dos_protection_threshold` can't be bigger than `max_num_skipped_message_keys`."
)
if len(shared_secret) != 32:
raise ValueError("The shared secret must consist of 32 bytes.")
self = cls()
self.__max_num_skipped_message_keys = max_num_skipped_message_keys
self.__skipped_message_keys = OrderedDict()
self.__aead = aead
self.__diffie_hellman_ratchet = await diffie_hellman_ratchet_class.create(
None,
recipient_ratchet_pub,
root_chain_kdf,
shared_secret,
message_chain_kdf,
message_chain_constant,
dos_protection_threshold
)
message_key, header = await self.__diffie_hellman_ratchet.next_encryption_key()
ciphertext = await self.__aead.encrypt(
message,
message_key,
self._build_associated_data(associated_data, header)
)
return (self, EncryptedMessage(header=header, ciphertext=ciphertext))
[docs]
@classmethod
async def decrypt_initial_message(
cls: Type[DoubleRatchetTypeT],
diffie_hellman_ratchet_class: Type[DiffieHellmanRatchet],
root_chain_kdf: Type[KDF],
message_chain_kdf: Type[KDF],
message_chain_constant: bytes,
dos_protection_threshold: int,
max_num_skipped_message_keys: int,
aead: Type[AEAD],
shared_secret: bytes,
own_ratchet_priv: bytes,
message: EncryptedMessage,
associated_data: bytes
) -> Tuple[DoubleRatchetTypeT, bytes]:
"""
Args:
diffie_hellman_ratchet_class: A non-abstract subclass of
:class:`~doubleratchet.diffie_hellman_ratchet.DiffieHellmanRatchet`.
root_chain_kdf: The KDF to use for the root chain. The KDF must be capable of deriving 64 bytes.
message_chain_kdf: The KDF to use for the sending and receiving chains. The KDF must be capable of
deriving 64 bytes.
message_chain_constant: The constant to feed into the sending and receiving KDF chains on each
step.
dos_protection_threshold: The maximum number of skipped message keys to calculate. If more than
that number of message keys are skipped, the keys are not calculated to prevent being DoSed.
max_num_skipped_message_keys: The maximum number of skipped message keys to store in case the lost
or out-of-order message comes in later. Older keys are discarded to make space for newer keys.
aead: The AEAD implementation to use for message en- and decryption.
shared_secret: A shared secret that was agreed on by means external to this protocol.
own_ratchet_priv: The ratchet private key to use initially.
message: The encrypted initial message.
associated_data: Additional data to authenticate without including it in the ciphertext.
Returns:
A configured instance of :class:`DoubleRatchet` ready to send and receive messages together with
the decrypted initial message.
Raises:
AuthenticationFailedException: if the message could not be authenticated using the associated
data.
DecryptionFailedException: if the decryption failed for a different reason (e.g. invalid padding).
DoSProtectionException: if a huge number of message keys were skipped that have to be calculated
first before decrypting the message.
"""
if dos_protection_threshold > max_num_skipped_message_keys:
raise ValueError(
"The `dos_protection_threshold` can't be bigger than `max_num_skipped_message_keys`."
)
if len(shared_secret) != 32:
raise ValueError("The shared secret must consist of 32 bytes.")
self = cls()
self.__max_num_skipped_message_keys = max_num_skipped_message_keys
self.__aead = aead
self.__diffie_hellman_ratchet = await diffie_hellman_ratchet_class.create(
own_ratchet_priv,
message.header.ratchet_pub,
root_chain_kdf,
shared_secret,
message_chain_kdf,
message_chain_constant,
dos_protection_threshold
)
message_key, skipped_message_keys = \
await self.__diffie_hellman_ratchet.next_decryption_key(message.header)
# Even the first message might have skipped message keys. The number of keys can't cross thresholds,
# thus no FIFO discarding required.
self.__skipped_message_keys = skipped_message_keys
return (self, await self.__aead.decrypt(
message.ciphertext,
message_key,
self._build_associated_data(associated_data, message.header)
))
@property
def sending_chain_length(self) -> int:
"""
Returns:
The length of the sending chain of the internal symmetric-key ratchet, as exposed by the internal
Diffie-Hellman ratchet.
"""
return self.__diffie_hellman_ratchet.sending_chain_length
@property
def receiving_chain_length(self) -> Optional[int]:
"""
Returns:
The length of the receiving chain of the internal symmetric-key ratchet, if it exists, as exposed
by the internal Diffie-Hellman ratchet.
"""
return self.__diffie_hellman_ratchet.receiving_chain_length
####################
# abstract methods #
####################
[docs]
@staticmethod
@abstractmethod
def _build_associated_data(associated_data: bytes, header: Header) -> bytes:
"""
Args:
associated_data: The associated data to prepend to the output. If the associated data is not
guaranteed to be a parseable byte sequence, a length value should be prepended to ensure that
the output is parseable as a unique pair (associated data, header).
header: The message header to encode in a unique, reversible manner.
Returns:
A byte sequence encoding the associated data and the header in a unique, reversible way.
"""
#################
# serialization #
#################
@property
def model(self) -> DoubleRatchetModel:
"""
Returns:
The internal state of this :class:`DoubleRatchet` as a pydantic model.
"""
return DoubleRatchetModel(
diffie_hellman_ratchet=self.__diffie_hellman_ratchet.model,
skipped_message_keys=[ SkippedMessageKeyModel(
ratchet_pub=ratchet_pub,
index=index,
message_key=message_key
) for (ratchet_pub, index), message_key in self.__skipped_message_keys.items() ]
)
@property
def json(self) -> JSONObject:
"""
Returns:
The internal state of this :class:`DoubleRatchet` as a JSON-serializable Python object.
"""
return cast(JSONObject, json.loads(self.model.model_dump_json()))
[docs]
@classmethod
def from_model(
cls: Type[DoubleRatchetTypeT],
model: DoubleRatchetModel,
diffie_hellman_ratchet_class: Type[DiffieHellmanRatchet],
root_chain_kdf: Type[KDF],
message_chain_kdf: Type[KDF],
message_chain_constant: bytes,
dos_protection_threshold: int,
max_num_skipped_message_keys: int,
aead: Type[AEAD]
) -> DoubleRatchetTypeT:
"""
Args:
model: The pydantic model holding the internal state of a :class:`DoubleRatchet`, as produced
by :attr:`model`.
diffie_hellman_ratchet_class: A non-abstract subclass of
:class:`~doubleratchet.diffie_hellman_ratchet.DiffieHellmanRatchet`.
root_chain_kdf: The KDF to use for the root chain. The KDF must be capable of deriving 64 bytes.
message_chain_kdf: The KDF to use for the sending and receiving chains. The KDF must be capable of
deriving 64 bytes.
message_chain_constant: The constant to feed into the sending and receiving KDF chains on each
step.
dos_protection_threshold: The maximum number of skipped message keys to calculate. If more than
that number of message keys are skipped, the keys are not calculated to prevent being DoSed.
max_num_skipped_message_keys: The maximum number of skipped message keys to store in case the lost
or out-of-order message comes in later. Older keys are discarded to make space for newer keys.
aead: The AEAD implementation to use for message en- and decryption.
Returns:
A configured instance of :class:`DoubleRatchet`, with internal state restored from the model.
Raises:
InconsistentSerializationException: if the serialized data is structurally correct, but
incomplete. This can only happen when migrating an instance from pre-stable data that was
serialized before sending or receiving a single message. In this case, the serialized instance
is basically uninitialized and can be discarded/replaced with a new instance using
:meth:`encrypt_initial_message` or :meth:`decrypt_initial_message` without losing information.
Warning:
Migrations are not provided via the :attr:`model`/:meth:`from_model` API. Use
:attr:`json`/:meth:`from_json` instead. Refer to :ref:`serialization_and_migration` in the
documentation for details.
"""
if dos_protection_threshold > max_num_skipped_message_keys:
raise ValueError(
"The `dos_protection_threshold` can't be bigger than `max_num_skipped_message_keys`."
)
self = cls()
self.__max_num_skipped_message_keys = max_num_skipped_message_keys
self.__skipped_message_keys = OrderedDict(
((smk.ratchet_pub, smk.index), smk.message_key)
for smk
in model.skipped_message_keys
)
self.__aead = aead
self.__diffie_hellman_ratchet = diffie_hellman_ratchet_class.from_model(
model.diffie_hellman_ratchet,
root_chain_kdf,
message_chain_kdf,
message_chain_constant,
dos_protection_threshold
)
return self
[docs]
@classmethod
def from_json(
cls: Type[DoubleRatchetTypeT],
serialized: JSONObject,
diffie_hellman_ratchet_class: Type[DiffieHellmanRatchet],
root_chain_kdf: Type[KDF],
message_chain_kdf: Type[KDF],
message_chain_constant: bytes,
dos_protection_threshold: int,
max_num_skipped_message_keys: int,
aead: Type[AEAD]
) -> DoubleRatchetTypeT:
"""
Args:
serialized: A JSON-serializable Python object holding the internal state of a
:class:`DoubleRatchet`, as produced by :attr:`json`.
diffie_hellman_ratchet_class: A non-abstract subclass of
:class:`~doubleratchet.diffie_hellman_ratchet.DiffieHellmanRatchet`.
root_chain_kdf: The KDF to use for the root chain. The KDF must be capable of deriving 64 bytes.
message_chain_kdf: The KDF to use for the sending and receiving chains. The KDF must be capable of
deriving 64 bytes.
message_chain_constant: The constant to feed into the sending and receiving KDF chains on each
step.
dos_protection_threshold: The maximum number of skipped message keys to calculate. If more than
that number of message keys are skipped, the keys are not calculated to prevent being DoSed.
max_num_skipped_message_keys: The maximum number of skipped message keys to store in case the lost
or out-of-order message comes in later. Older keys are discarded to make space for newer keys.
aead: The AEAD implementation to use for message en- and decryption.
Returns:
A configured instance of :class:`DoubleRatchet`, with internal state restored from the serialized
data.
Raises:
InconsistentSerializationException: if the serialized data is structurally correct, but
incomplete. This can only happen when migrating an instance from pre-stable data that was
serialized before sending or receiving a single message. In this case, the serialized instance
is basically uninitialized and can be discarded/replaced with a new instance using
:meth:`encrypt_initial_message` or :meth:`decrypt_initial_message` without losing information.
"""
return cls.from_model(
parse_double_ratchet_model(serialized),
diffie_hellman_ratchet_class,
root_chain_kdf,
message_chain_kdf,
message_chain_constant,
dos_protection_threshold,
max_num_skipped_message_keys,
aead
)
#########################
# message en/decryption #
#########################
[docs]
async def encrypt_message(self, message: bytes, associated_data: bytes) -> EncryptedMessage:
"""
Args:
message: The message to encrypt.
associated_data: Additional data to authenticate without including it in the ciphertext.
Returns:
The encrypted message including the header to send to the recipient.
"""
message_key, header = await self.__diffie_hellman_ratchet.next_encryption_key()
ciphertext = await self.__aead.encrypt(
message,
message_key,
self._build_associated_data(associated_data, header)
)
return EncryptedMessage(header=header, ciphertext=ciphertext)
[docs]
async def decrypt_message(self, message: EncryptedMessage, associated_data: bytes) -> bytes:
"""
Args:
message: The encrypted message.
associated_data: Additional data to authenticate without including it in the ciphertext.
Returns:
The message plaintext, after decrypting and authenticating the ciphertext.
Raises:
AuthenticationFailedException: if the message could not be authenticated using the associated
data.
DecryptionFailedException: if the decryption failed for a different reason (e.g. invalid padding).
DoSProtectionException: if a huge number of message keys were skipped that have to be calculated
first before decrypting the message.
DuplicateMessageException: if this message appears to be a duplicate.
"""
# Be careful to only keep changes to the internal state on decryption success. To do so, work with a
# clone of the Diffie-Hellman ratchet, discard the clone on failure or replace the original with the
# clone on success.
# https://signal.org/docs/specifications/doubleratchet/#decrypting-messages
diffie_hellman_ratchet = copy.deepcopy(self.__diffie_hellman_ratchet)
skipped_message_keys: Optional[SkippedMessageKeys] = None
skipped_message_key_key = (message.header.ratchet_pub, message.header.sending_chain_length)
# Get the message key, either from the skipped message keys or from the Diffie-Hellman ratchet clone
message_key: bytes
try:
message_key = self.__skipped_message_keys[skipped_message_key_key]
except KeyError:
message_key, skipped_message_keys = \
await diffie_hellman_ratchet.next_decryption_key(message.header)
# Decrypt the message (or at least attempt to do so). At this point, the internal state of this
# instance remains untouched.
plaintext = await self.__aead.decrypt(
message.ciphertext,
message_key,
self._build_associated_data(associated_data, message.header)
)
# Following decryption success, apply relevant changes to the internal state.
# In case a skipped message key was used, remove it.
self.__skipped_message_keys.pop(skipped_message_key_key, None)
# Store new skipped message keys and limit their number.
if skipped_message_keys is not None:
self.__skipped_message_keys.update(skipped_message_keys)
self.__skipped_message_keys = OrderedDict(itertools.islice(
self.__skipped_message_keys.items(),
max(len(self.__skipped_message_keys) - self.__max_num_skipped_message_keys, 0),
None
))
# Store the clone.
self.__diffie_hellman_ratchet = diffie_hellman_ratchet
return plaintext