from __future__ import annotations
from abc import ABC, abstractmethod
import json
import time
import secrets
from typing import FrozenSet, Optional, Set, Tuple, Type, TypeVar, cast
import xeddsa
from .crypto_provider import HashFunction
from .crypto_provider_cryptography import CryptoProviderImpl
from .identity_key_pair import IdentityKeyPair, IdentityKeyPairSeed
from .migrations import parse_base_state_model
from .models import BaseStateModel
from .pre_key_pair import PreKeyPair
from .signed_pre_key_pair import SignedPreKeyPair
from .types import Bundle, IdentityKeyFormat, Header, JSONObject
__all__ = [
"KeyAgreementException",
"BaseState"
]
[docs]
class KeyAgreementException(Exception):
"""
Exception raised by :meth:`BaseState.get_shared_secret_active` and
:meth:`BaseState.get_shared_secret_passive` in case of an error related to the key agreement operation.
"""
BaseStateTypeT = TypeVar("BaseStateTypeT", bound="BaseState")
[docs]
class BaseState(ABC):
"""
This class is the core of this X3DH implementation. It offers methods to manually manage the X3DH state
and perform key agreements with other parties.
Warning:
This class requires manual state management, including e.g. signed pre key rotation, pre key
hiding/deletion and refills. The subclass :class:`~x3dh.state.State` automates those
management/maintenance tasks and should be preferred if external/manual management is not explicitly
wanted.
"""
def __init__(self) -> None:
# Just the type definitions here
self.__identity_key_format: IdentityKeyFormat
self.__hash_function: HashFunction
self.__info: bytes
self.__identity_key: IdentityKeyPair
self.__signed_pre_key: SignedPreKeyPair
self.__old_signed_pre_key: Optional[SignedPreKeyPair]
self.__pre_keys: Set[PreKeyPair]
self.__hidden_pre_keys: Set[PreKeyPair]
[docs]
@classmethod
def create(
cls: Type[BaseStateTypeT],
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes,
identity_key_pair: Optional[IdentityKeyPair] = None
) -> BaseStateTypeT:
"""
Args:
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
identity_key_pair: If set, use the given identity key pair instead of generating a new one.
Returns:
A configured instance of :class:`~x3dh.base_state.BaseState`. Note that an identity key pair and a
signed pre key are generated, but no pre keys. Use :meth:`generate_pre_keys` to generate some.
"""
self = cls()
self.__identity_key_format = identity_key_format
self.__hash_function = hash_function
self.__info = info
self.__identity_key = identity_key_pair or IdentityKeyPairSeed(secrets.token_bytes(32))
self.__signed_pre_key = self.__generate_spk()
self.__old_signed_pre_key = None
self.__pre_keys = set()
self.__hidden_pre_keys = set()
return self
####################
# abstract methods #
####################
[docs]
@staticmethod
@abstractmethod
def _encode_public_key(key_format: IdentityKeyFormat, pub: bytes) -> bytes:
"""
Args:
key_format: The format in which this public key is serialized.
pub: The public key.
Returns:
An encoding of the public key, possibly including information about the curve and type of key,
though this is application defined. Note that two different public keys must never result in the
same byte sequence, uniqueness of the public keys must be preserved.
"""
raise NotImplementedError("Create a subclass of BaseState and implement `_encode_public_key`.")
#################
# serialization #
#################
@property
def model(self) -> BaseStateModel:
"""
Returns:
The internal state of this :class:`BaseState` as a pydantic model. Note that pre keys hidden using
:meth:`hide_pre_key` are not considered part of the state.
"""
return BaseStateModel(
identity_key=self.__identity_key.model,
signed_pre_key=self.__signed_pre_key.model,
old_signed_pre_key=None if self.__old_signed_pre_key is None else self.__old_signed_pre_key.model,
pre_keys=frozenset(pre_key.priv for pre_key in self.__pre_keys)
)
@property
def json(self) -> JSONObject:
"""
Returns:
The internal state of this :class:`BaseState` as a JSON-serializable Python object. Note that pre
keys hidden using :meth:`hide_pre_key` are not considered part of the state.
"""
return cast(JSONObject, json.loads(self.model.model_dump_json()))
[docs]
@classmethod
def from_model(
cls: Type[BaseStateTypeT],
model: BaseStateModel,
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes
) -> BaseStateTypeT:
"""
Args:
model: The pydantic model holding the internal state of a :class:`BaseState`, as produced by
:attr:`model`.
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
Returns:
A configured instance of :class:`BaseState`, with internal state restored from the model.
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.
"""
self = cls()
self.__identity_key_format = identity_key_format
self.__hash_function = hash_function
self.__info = info
self.__identity_key = IdentityKeyPair.from_model(model.identity_key)
self.__signed_pre_key = SignedPreKeyPair.from_model(model.signed_pre_key)
self.__old_signed_pre_key = (
None
if model.old_signed_pre_key is None
else SignedPreKeyPair.from_model(model.old_signed_pre_key)
)
self.__pre_keys = { PreKeyPair(pre_key) for pre_key in model.pre_keys }
self.__hidden_pre_keys = set()
return self
[docs]
@classmethod
def from_json(
cls: Type[BaseStateTypeT],
serialized: JSONObject,
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes
) -> Tuple[BaseStateTypeT, bool]:
"""
Args:
serialized: A JSON-serializable Python object holding the internal state of a :class:`BaseState`,
as produced by :attr:`json`.
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
Returns:
A configured instance of :class:`BaseState`, with internal state restored from the serialized
data, and a flag that indicates whether the bundle needs to be published. The latter was part of
the pre-stable serialization format.
"""
model, bundle_needs_publish = parse_base_state_model(serialized)
self = cls.from_model(
model,
identity_key_format,
hash_function,
info
)
return self, bundle_needs_publish
#################################
# key generation and management #
#################################
def __generate_spk(self) -> SignedPreKeyPair:
"""
Returns:
A newly generated signed pre key.
"""
# Get the own identity key in the format required for signing, forcing the sign bit if necessary to
# comply with XEdDSA
identity_key = self.__identity_key.as_priv().priv
if self.__identity_key_format is IdentityKeyFormat.CURVE_25519:
identity_key = xeddsa.priv_force_sign(identity_key, False)
# Generate the private key of the new signed pre key
priv = secrets.token_bytes(32)
# Sign the encoded public key of the new signed pre key
sig = xeddsa.ed25519_priv_sign(
identity_key,
self._encode_public_key(IdentityKeyFormat.CURVE_25519, xeddsa.priv_to_curve25519_pub(priv))
)
# Add the current timestamp
return SignedPreKeyPair(priv=priv, sig=sig, timestamp=int(time.time()))
@property
def old_signed_pre_key(self) -> Optional[bytes]:
"""
Returns:
The old signed pre key, if there is one.
"""
return None if self.__old_signed_pre_key is None else self.__old_signed_pre_key.pub
[docs]
def signed_pre_key_age(self) -> int:
"""
Returns:
The age of the signed pre key, i.e. the time elapsed since it was last rotated, in seconds.
"""
return int(time.time()) - self.__signed_pre_key.timestamp
[docs]
def rotate_signed_pre_key(self) -> None:
"""
Rotate the signed pre key. Keep the old signed pre key around for one additional rotation period, i.e.
until this method is called again.
"""
self.__old_signed_pre_key = self.__signed_pre_key
self.__signed_pre_key = self.__generate_spk()
@property
def hidden_pre_keys(self) -> FrozenSet[bytes]:
"""
Returns:
The currently hidden pre keys.
"""
return frozenset(pre_key.pub for pre_key in self.__hidden_pre_keys)
[docs]
def hide_pre_key(self, pre_key_pub: bytes) -> bool:
"""
Hide a pre key from the bundle returned by :attr:`bundle` and pre key count returned by
:meth:`get_num_visible_pre_keys`, but keep the pre key for cryptographic operations. Hidden pre keys
are not included in the serialized state as returned by :attr:`model` and :attr:`json`.
Args:
pre_key_pub: The pre key to hide.
Returns:
Whether the pre key was visible before and is hidden now.
"""
hidden_pre_keys = frozenset(filter(lambda pre_key: pre_key.pub == pre_key_pub, self.__pre_keys))
self.__pre_keys -= hidden_pre_keys
self.__hidden_pre_keys |= hidden_pre_keys
return len(hidden_pre_keys) > 0
[docs]
def delete_pre_key(self, pre_key_pub: bytes) -> bool:
"""
Delete a pre key.
Args:
pre_key_pub: The pre key to delete. Can be visible or hidden.
Returns:
Whether the pre key existed before and is deleted now.
"""
deleted_pre_keys = frozenset(filter(
lambda pre_key: pre_key.pub == pre_key_pub,
self.__pre_keys | self.__hidden_pre_keys
))
self.__pre_keys -= deleted_pre_keys
self.__hidden_pre_keys -= deleted_pre_keys
return len(deleted_pre_keys) > 0
[docs]
def delete_hidden_pre_keys(self) -> None:
"""
Delete all pre keys that were previously hidden using :meth:`hide_pre_key`.
"""
self.__hidden_pre_keys = set()
[docs]
def get_num_visible_pre_keys(self) -> int:
"""
Returns:
The number of visible pre keys available. The number returned here matches the number of pre keys
included in the bundle returned by :attr:`bundle`.
"""
return len(self.__pre_keys)
[docs]
def generate_pre_keys(self, num_pre_keys: int) -> None:
"""
Generate and store pre keys.
Args:
num_pre_keys: The number of pre keys to generate.
"""
for _ in range(num_pre_keys):
self.__pre_keys.add(PreKeyPair(priv=secrets.token_bytes(32)))
@property
def bundle(self) -> Bundle:
"""
Returns:
The bundle, i.e. the public information of this state.
"""
identity_key = self.__identity_key.as_priv().priv
return Bundle(
identity_key=(
xeddsa.priv_to_curve25519_pub(identity_key)
if self.__identity_key_format is IdentityKeyFormat.CURVE_25519
else xeddsa.priv_to_ed25519_pub(identity_key)
),
signed_pre_key=self.__signed_pre_key.pub,
signed_pre_key_sig=self.__signed_pre_key.sig,
pre_keys=frozenset(pre_key.pub for pre_key in self.__pre_keys)
)
#################
# key agreement #
#################
[docs]
async def get_shared_secret_active(
self,
bundle: Bundle,
associated_data_appendix: bytes = b"",
require_pre_key: bool = True
) -> Tuple[bytes, bytes, Header]:
"""
Perform an X3DH key agreement, actively.
Args:
bundle: The bundle of the passive party.
associated_data_appendix: Additional information to append to the associated data, like usernames,
certificates or other identifying information.
require_pre_key: Use this flag to abort the key agreement if the bundle does not contain a pre
key.
Returns:
The shared secret and associated data shared between both parties, and the header required by the
other party to complete the passive part of the key agreement.
Raises:
KeyAgreementException: If an error occurs during the key agreement. The exception message will
contain (human-readable) details.
"""
# Check whether a pre key is required but not included
if len(bundle.pre_keys) == 0 and require_pre_key:
raise KeyAgreementException("This bundle does not contain a pre key.")
# Get the identity key of the other party in the format required for signature verification
other_identity_key = bundle.identity_key
if self.__identity_key_format is IdentityKeyFormat.CURVE_25519:
other_identity_key = xeddsa.curve25519_pub_to_ed25519_pub(other_identity_key, False)
# Verify the signature on the signed pre key of the other party
if not xeddsa.ed25519_verify(
bundle.signed_pre_key_sig,
other_identity_key,
self._encode_public_key(IdentityKeyFormat.CURVE_25519, bundle.signed_pre_key)
):
raise KeyAgreementException("The signature of the signed pre key could not be verified.")
# All pre-checks successful.
# Choose a pre key if available
pre_key = None if len(bundle.pre_keys) == 0 else secrets.choice(list(bundle.pre_keys))
# Generate the ephemeral key required for the key agreement
ephemeral_key = secrets.token_bytes(32)
# Get the own identity key in the format required for X25519
own_identity_key = self.__identity_key.as_priv().priv
# Get the identity key of the other party in the format required for X25519
other_identity_key = bundle.identity_key
if self.__identity_key_format is IdentityKeyFormat.ED_25519:
other_identity_key = xeddsa.ed25519_pub_to_curve25519_pub(other_identity_key)
# Calculate the three to four Diffie-Hellman shared secrets that become the input of HKDF in the next
# step
dh1 = xeddsa.x25519(own_identity_key, bundle.signed_pre_key)
dh2 = xeddsa.x25519(ephemeral_key, other_identity_key)
dh3 = xeddsa.x25519(ephemeral_key, bundle.signed_pre_key)
dh4 = b"" if pre_key is None else xeddsa.x25519(ephemeral_key, pre_key)
# Prepare salt and padding
salt = b"\x00" * self.__hash_function.hash_size
padding = b"\xFF" * 32
# Use HKDF to derive the final shared secret
shared_secret = await CryptoProviderImpl.hkdf_derive(
self.__hash_function,
32,
salt,
self.__info,
padding + dh1 + dh2 + dh3 + dh4
)
# Build the associated data for further use by other protocols
associated_data = (
self._encode_public_key(self.__identity_key_format, self.bundle.identity_key)
+ self._encode_public_key(self.__identity_key_format, bundle.identity_key)
+ associated_data_appendix
)
# Build the header required by the other party to complete the passive part of the key agreement
header = Header(
identity_key=self.bundle.identity_key,
ephemeral_key=xeddsa.priv_to_curve25519_pub(ephemeral_key),
pre_key=pre_key,
signed_pre_key=bundle.signed_pre_key
)
return shared_secret, associated_data, header
[docs]
async def get_shared_secret_passive(
self,
header: Header,
associated_data_appendix: bytes = b"",
require_pre_key: bool = True
) -> Tuple[bytes, bytes, SignedPreKeyPair]:
"""
Perform an X3DH key agreement, passively.
Args:
header: The header received from the active party.
associated_data_appendix: Additional information to append to the associated data, like usernames,
certificates or other identifying information.
require_pre_key: Use this flag to abort the key agreement if the active party did not use a pre
key.
Returns:
The shared secret and the associated data shared between both parties, and the signed pre key pair
that was used during the key exchange, for use by follow-up protocols.
Raises:
KeyAgreementException: If an error occurs during the key agreement. The exception message will
contain (human-readable) details.
"""
# Check whether the signed pre key used by this initiation is still available
signed_pre_key: Optional[SignedPreKeyPair] = None
if header.signed_pre_key == self.__signed_pre_key.pub:
# The current signed pre key was used
signed_pre_key = self.__signed_pre_key
if self.__old_signed_pre_key is not None and header.signed_pre_key == self.__old_signed_pre_key.pub:
# The old signed pre key was used
signed_pre_key = self.__old_signed_pre_key
if signed_pre_key is None:
raise KeyAgreementException(
"This key agreement attempt uses a signed pre key that is not available any more."
)
# Check whether a pre key is required but not used
if header.pre_key is None and require_pre_key:
raise KeyAgreementException("This key agreement attempt does not use a pre key.")
# If a pre key was used, check whether it is still available
pre_key: Optional[bytes] = None
if header.pre_key is not None:
pre_key = next((
pre_key.priv
for pre_key
in self.__pre_keys | self.__hidden_pre_keys
if pre_key.pub == header.pre_key
), None)
if pre_key is None:
raise KeyAgreementException(
"This key agreement attempt uses a pre key that is not available any more."
)
# Get the own identity key in the format required for X25519
own_identity_key = self.__identity_key.as_priv().priv
# Get the identity key of the other party in the format required for X25519
other_identity_key = header.identity_key
if self.__identity_key_format is IdentityKeyFormat.ED_25519:
other_identity_key = xeddsa.ed25519_pub_to_curve25519_pub(other_identity_key)
# Calculate the three to four Diffie-Hellman shared secrets that become the input of HKDF in the next
# step
dh1 = xeddsa.x25519(signed_pre_key.priv, other_identity_key)
dh2 = xeddsa.x25519(own_identity_key, header.ephemeral_key)
dh3 = xeddsa.x25519(signed_pre_key.priv, header.ephemeral_key)
dh4 = b"" if pre_key is None else xeddsa.x25519(pre_key, header.ephemeral_key)
# Prepare salt and padding
salt = b"\x00" * self.__hash_function.hash_size
padding = b"\xFF" * 32
# Use HKDF to derive the final shared secret
shared_secret = await CryptoProviderImpl.hkdf_derive(
self.__hash_function,
32,
salt,
self.__info,
padding + dh1 + dh2 + dh3 + dh4
)
# Build the associated data for further use by other protocols
associated_data = (
self._encode_public_key(self.__identity_key_format, header.identity_key)
+ self._encode_public_key(self.__identity_key_format, self.bundle.identity_key)
+ associated_data_appendix
)
return shared_secret, associated_data, signed_pre_key