# This import from future (theoretically) enables sphinx_autodoc_typehints to handle type aliases better
from __future__ import annotations

import base64
import secrets
from typing import Dict, Optional, Tuple, cast
from typing_extensions import Final

import doubleratchet
from doubleratchet.recommended import (
    aead_aes_hmac,
    diffie_hellman_ratchet_curve25519,
    HashFunction,
    kdf_hkdf,
    kdf_separate_hmacs
)
from doubleratchet.recommended.crypto_provider_impl import CryptoProviderImpl
import google.protobuf.message
import x3dh
import x3dh.identity_key_pair

from omemo.backend import Backend, DecryptionFailed, KeyExchangeFailed
from omemo.bundle import Bundle
from omemo.identity_key_pair import IdentityKeyPair, IdentityKeyPairSeed
from omemo.message import Content, EncryptedKeyMaterial, PlainKeyMaterial, KeyExchange
from omemo.session import Initiation, Session
from omemo.storage import Storage
from omemo.types import JSONType

# https://github.com/PyCQA/pylint/issues/4987
from .twomemo_pb2 import (  # pylint: disable=no-name-in-module
    OMEMOAuthenticatedMessage,
    OMEMOKeyExchange,
    OMEMOMessage
)


__all__ = [
    "Twomemo",
    "NAMESPACE",
    "AEADImpl",
    "BundleImpl",
    "ContentImpl",
    "DoubleRatchetImpl",
    "EncryptedKeyMaterialImpl",
    "KeyExchangeImpl",
    "MessageChainKDFImpl",
    "PlainKeyMaterialImpl",
    "RootChainKDFImpl",
    "SessionImpl",
    "StateImpl"
]


NAMESPACE: Final = "urn:xmpp:omemo:2"


class RootChainKDFImpl(kdf_hkdf.KDF):
    """
    The root chain KDF implementation used by this version of the specification.
    """

    @staticmethod
    def _get_hash_function() -> HashFunction:
        return HashFunction.SHA_256

    @staticmethod
    def _get_info() -> bytes:
        return "OMEMO Root Chain".encode("ASCII")


class MessageChainKDFImpl(kdf_separate_hmacs.KDF):
    """
    The message chain KDF implementation used by this version of the specification.
    """

    @staticmethod
    def _get_hash_function() -> HashFunction:
        return HashFunction.SHA_256


class AEADImpl(aead_aes_hmac.AEAD):
    """
    The AEAD used by this backend as part of the Double Ratchet. While this implementation derives from
    :class:`doubleratchet.recommended.aead_aes_hmac.AEAD`, it actually doesn't use any of its code. This is
    due to a minor difference in the way the associated data is built. The derivation only has symbolic value.

    Can only be used with :class:`DoubleRatchetImpl`, due to the reliance on a certain structure of the
    associated data.
    """

    AUTHENTICATION_TAG_TRUNCATED_LENGTH: Final = 16

    @staticmethod
    def _get_hash_function() -> HashFunction:
        return HashFunction.SHA_256

    @staticmethod
    def _get_info() -> bytes:
        return "OMEMO Message Key Material".encode("ASCII")

    @classmethod
    async def encrypt(cls, plaintext: bytes, key: bytes, associated_data: bytes) -> bytes:
        hash_function = cls._get_hash_function()

        encryption_key, authentication_key, iv = await cls.__derive(key, hash_function, cls._get_info())

        # Encrypt the plaintext using AES-256 (the 256 bit are implied by the key size) in CBC mode and the
        # previously created key and IV, after padding it with PKCS#7
        ciphertext = await CryptoProviderImpl.aes_cbc_encrypt(encryption_key, iv, plaintext)

        # Parse the associated data
        associated_data, header = cls.__parse_associated_data(associated_data)

        # Build an OMEMOMessage including the header and the ciphertext
        omemo_message = OMEMOMessage(
            n=header.sending_chain_length,
            pn=header.previous_sending_chain_length,
            dh_pub=header.ratchet_pub,
            ciphertext=ciphertext
        ).SerializeToString()

        # Calculate the authentication tag over the associated data and the OMEMOMessage, truncate the
        # authentication tag to AUTHENTICATION_TAG_TRUNCATED_LENGTH bytes
        auth = (await CryptoProviderImpl.hmac_calculate(
            authentication_key,
            hash_function,
            associated_data + omemo_message
        ))[:AEADImpl.AUTHENTICATION_TAG_TRUNCATED_LENGTH]

        # Serialize the authentication tag with the OMEMOMessage in an OMEMOAuthenticatedMessage.
        return OMEMOAuthenticatedMessage(mac=auth, message=omemo_message).SerializeToString()

    @classmethod
    async def decrypt(cls, ciphertext: bytes, key: bytes, associated_data: bytes) -> bytes:
        hash_function = cls._get_hash_function()

        decryption_key, authentication_key, iv = await cls.__derive(key, hash_function, cls._get_info())

        # Parse the associated data
        associated_data, header = cls.__parse_associated_data(associated_data)

        # Parse the ciphertext as an OMEMOAuthenticatedMessage
        try:
            omemo_authenticated_message = OMEMOAuthenticatedMessage.FromString(ciphertext)
        except google.protobuf.message.DecodeError as e:
            raise doubleratchet.DecryptionFailedException() from e

        # Calculate and verify the authentication tag
        new_auth = (await CryptoProviderImpl.hmac_calculate(
            authentication_key,
            hash_function,
            associated_data + omemo_authenticated_message.message
        ))[:AEADImpl.AUTHENTICATION_TAG_TRUNCATED_LENGTH]

        if new_auth != omemo_authenticated_message.mac:
            raise doubleratchet.aead.AuthenticationFailedException("Authentication tags do not match.")

        # Parse the OMEMOMessage contained in the OMEMOAuthenticatedMessage
        try:
            omemo_message = OMEMOMessage.FromString(omemo_authenticated_message.message)
        except google.protobuf.message.DecodeError as e:
            raise doubleratchet.DecryptionFailedException() from e

        # Make sure that the headers match as a little additional consistency check
        if header != doubleratchet.Header(omemo_message.dh_pub, omemo_message.pn, omemo_message.n):
            raise doubleratchet.aead.AuthenticationFailedException("Header mismatch.")

        # Decrypt the plaintext using AES-256 (the 256 bit are implied by the key size) in CBC mode and the
        # previously created key and IV, and unpad the resulting plaintext with PKCS#7
        return await CryptoProviderImpl.aes_cbc_decrypt(decryption_key, iv, omemo_message.ciphertext)

    @staticmethod
    async def __derive(key: bytes, hash_function: HashFunction, info: bytes) -> Tuple[bytes, bytes, bytes]:
        # Prepare the salt, a zero-filled byte sequence with the size of the hash digest
        salt = b"\x00" * hash_function.hash_size

        # Derive 80 bytes
        hkdf_out = await CryptoProviderImpl.hkdf_derive(
            hash_function=hash_function,
            length=80,
            salt=salt,
            info=info,
            key_material=key
        )

        # Split these 80 bytes into three parts
        return hkdf_out[:32], hkdf_out[32:64], hkdf_out[64:]

    @staticmethod
    def __parse_associated_data(associated_data: bytes) -> Tuple[bytes, doubleratchet.Header]:
        """
        Parse the associated data as built by :meth:`DoubleRatchetImpl._build_associated_data`.

        Args:
            associated_data: The associated data.

        Returns:
            The original associated data and the header used to build it.

        Raises:
            DecryptionFailedException: if the data is malformed.
        """

        associated_data_length = StateImpl.IDENTITY_KEY_ENCODING_LENGTH * 2

        try:
            omemo_message = OMEMOMessage.FromString(associated_data[associated_data_length:])
        except google.protobuf.message.DecodeError as e:
            raise doubleratchet.DecryptionFailedException() from e

        associated_data = associated_data[:associated_data_length]

        return associated_data, doubleratchet.Header(omemo_message.dh_pub, omemo_message.pn, omemo_message.n)


class DoubleRatchetImpl(doubleratchet.DoubleRatchet):
    """
    The Double Ratchet implementation used by this version of the specification.
    """

    MESSAGE_CHAIN_CONSTANT: Final = b"\x02\x01"

    @staticmethod
    def _build_associated_data(associated_data: bytes, header: doubleratchet.Header) -> bytes:
        return associated_data + OMEMOMessage(
            n=header.sending_chain_length,
            pn=header.previous_sending_chain_length,
            dh_pub=header.ratchet_pub
        ).SerializeToString()


class StateImpl(x3dh.BaseState):
    """
    The X3DH state implementation used by this version of the specification.
    """

    INFO: Final = "OMEMO X3DH".encode("ASCII")
    IDENTITY_KEY_ENCODING_LENGTH: Final = 32

    @staticmethod
    def _encode_public_key(key_format: x3dh.IdentityKeyFormat, pub: bytes) -> bytes:
        return pub


class BundleImpl(Bundle):
    """
    :class:`~omemo.bundle.Bundle` implementation as a simple storage type.
    """

    def __init__(
        self,
        bare_jid: str,
        device_id: int,
        bundle: x3dh.Bundle,
        signed_pre_key_id: int,
        pre_key_ids: Dict[bytes, int]
    ) -> None:
        """
        Args:
            bare_jid: The bare JID this bundle belongs to.
            device_id: The device id of the specific device this bundle belongs to.
            bundle: The bundle to store in this instance.
            signed_pre_key_id: The id of the signed pre key referenced in the bundle.
            pre_key_ids: A dictionary that maps each pre key referenced in the bundle to its id.
        """

        self.__bare_jid = bare_jid
        self.__device_id = device_id
        self.__bundle = bundle
        self.__signed_pre_key_id = signed_pre_key_id
        self.__pre_key_ids = dict(pre_key_ids)

    @property
    def namespace(self) -> str:
        return NAMESPACE

    @property
    def bare_jid(self) -> str:
        return self.__bare_jid

    @property
    def device_id(self) -> int:
        return self.__device_id

    @property
    def identity_key(self) -> bytes:
        return self.__bundle.identity_key

    def __eq__(self, other: object) -> bool:
        if isinstance(other, BundleImpl):
            return (
                other.bare_jid == self.bare_jid
                and other.device_id == self.device_id
                and other.bundle == self.bundle
                and other.signed_pre_key_id == self.signed_pre_key_id
                and other.pre_key_ids == self.pre_key_ids
            )

        return False

    def __hash__(self) -> int:
        return hash((
            self.bare_jid,
            self.device_id,
            self.bundle,
            self.signed_pre_key_id,
            frozenset(self.pre_key_ids.items())
        ))

    @property
    def bundle(self) -> x3dh.Bundle:
        """
        Returns:
            The bundle held by this instance.
        """

        return self.__bundle

    @property
    def signed_pre_key_id(self) -> int:
        """
        Returns:
            The id of the signed pre key referenced in the bundle.
        """

        return self.__signed_pre_key_id

    @property
    def pre_key_ids(self) -> Dict[bytes, int]:
        """
        Returns:
            A dictionary that maps each pre key referenced in the bundle to its id.
        """

        return dict(self.__pre_key_ids)


class ContentImpl(Content):
    """
    :class:`~omemo.message.Content` implementation as a simple storage type.
    """

    def __init__(self, ciphertext: bytes) -> None:
        """
        Args:
            ciphertext: The ciphertext to store in this instance.

        Note:
            For empty OMEMO messages as per the specification, the ciphertext is set to an empty byte string.
        """

        self.__ciphertext = ciphertext

    @property
    def empty(self) -> bool:
        return self.__ciphertext == b""

    @staticmethod
    def make_empty() -> ContentImpl:
        """
        Returns:
            An "empty" instance, i.e. one that corresponds to an empty OMEMO message as per the specification.
            The ciphertext stored in empty instances is a byte string of zero length.
        """

        return ContentImpl(b"")

    @property
    def ciphertext(self) -> bytes:
        """
        Returns:
            The ciphertext held by this instance.
        """

        return self.__ciphertext


class EncryptedKeyMaterialImpl(EncryptedKeyMaterial):
    """
    :class:`~omemo.message.EncryptedKeyMaterial` implementation as a simple storage type.
    """

    def __init__(
        self,
        bare_jid: str,
        device_id: int,
        encrypted_message: doubleratchet.EncryptedMessage
    ) -> None:
        """
        Args:
            bare_jid: The bare JID of the other party.
            device_id: The device id of the specific device of the other party.
            encrypted_message: The encrypted Double Ratchet message to store in this instance.
        """

        self.__bare_jid = bare_jid
        self.__device_id = device_id
        self.__encrypted_message = encrypted_message

    @property
    def bare_jid(self) -> str:
        return self.__bare_jid

    @property
    def device_id(self) -> int:
        return self.__device_id

    @property
    def encrypted_message(self) -> doubleratchet.EncryptedMessage:
        """
        Returns:
            The encrypted Double Ratchet message held by this instance.
        """

        return self.__encrypted_message

    def serialize(self) -> bytes:
        """
        Returns:
            A serialized OMEMOAuthenticatedMessage message structure representing the content of this
            instance.
        """

        # The ciphertext field contains the result of :meth:`AEADImpl.encrypt`, which is a serialized
        # OMEMOAuthenticatedMessage with all fields already correctly set, thus it can be used here as is.
        return self.__encrypted_message.ciphertext

    @staticmethod
    def parse(authenticated_message: bytes, bare_jid: str, device_id: int) -> EncryptedKeyMaterialImpl:
        """
        Args:
            authenticated_message: A serialized OMEMOAuthenticatedMessage message structure.
            bare_jid: The bare JID of the other party.
            device_id: The device id of the specific device of the other party.

        Returns:
            An instance of this class, parsed from the OMEMOAuthenticatedMessage.

        Raises:
            ValueError: if the data is malformed.
        """

        # Parse the OMEMOAuthenticatedMessage and OMEMOMessage structures to extract the header.
        try:
            message = OMEMOMessage.FromString(OMEMOAuthenticatedMessage.FromString(
                authenticated_message
            ).message)
        except google.protobuf.message.DecodeError as e:
            raise ValueError() from e

        return EncryptedKeyMaterialImpl(
            bare_jid,
            device_id,
            doubleratchet.EncryptedMessage(
                doubleratchet.Header(message.dh_pub, message.pn, message.n),
                authenticated_message
            )
        )


class PlainKeyMaterialImpl(PlainKeyMaterial):
    """
    :class:`~omemo.message.PlainKeyMaterial` implementation as a simple storage type.
    """

    KEY_LENGTH: Final = 32

    def __init__(self, key: bytes, auth_tag: bytes) -> None:
        """
        Args:
            key: The key to store in this instance.
            auth_tag: The authentication tag to store in this instance.

        Note:
            For empty OMEMO messages as per the specification, the key is set to :attr:`KEY_LENGTH`
            zero-bytes, and the auth tag is set to an empty byte string.
        """

        self.__key = key
        self.__auth_tag = auth_tag

    @property
    def key(self) -> bytes:
        """
        Returns:
            The key held by this instance.
        """

        return self.__key

    @property
    def auth_tag(self) -> bytes:
        """
        Returns:
            The authentication tag held by this instance.
        """

        return self.__auth_tag

    @staticmethod
    def make_empty() -> PlainKeyMaterialImpl:
        """
        Returns:
            An "empty" instance, i.e. one that corresponds to an empty OMEMO message as per the specification.
            The key stored in empty instances is a byte string of :attr:`KEY_LENGTH` zero-bytes, and the auth
            tag is an empty byte string.
        """

        return PlainKeyMaterialImpl(b"\x00" * PlainKeyMaterialImpl.KEY_LENGTH, b"")


class KeyExchangeImpl(KeyExchange):
    """
    :class:`~omemo.message.KeyExchange` implementation as a simple storage type.

    There are two kinds of instances:

    - Completely filled instances
    - Partially filled instances received via network

    Empty fields are filled with filler values such that the data types and lengths still match expectations.
    """

    def __init__(self, header: x3dh.Header, signed_pre_key_id: int, pre_key_id: int) -> None:
        """
        Args:
            header: The header to store in this instance.
            signed_pre_key_id: The id of the signed pre key referenced in the header.
            pre_key_id: The id of the pre key referenced in the header.
        """

        self.__header = header
        self.__signed_pre_key_id = signed_pre_key_id
        self.__pre_key_id = pre_key_id

    @property
    def identity_key(self) -> bytes:
        return self.__header.identity_key

    def builds_same_session(self, other: KeyExchange) -> bool:
        # The signed pre key id and pre key id are enough for uniqueness; ignoring the actual signed pre key
        # and pre key bytes here makes it possible to compare network instances with completely filled
        # instances.
        return isinstance(other, KeyExchangeImpl) and (
            other.header.identity_key == self.header.identity_key
            and other.header.ephemeral_key == self.header.ephemeral_key
            and other.signed_pre_key_id == self.signed_pre_key_id
            and other.pre_key_id == self.pre_key_id
        )

    @property
    def header(self) -> x3dh.Header:
        """
        Returns:
            The header held by this instance.
        """

        return self.__header

    @property
    def signed_pre_key_id(self) -> int:
        """
        Returns:
            The id of the signed pre key referenced in the header.
        """

        return self.__signed_pre_key_id

    @property
    def pre_key_id(self) -> int:
        """
        Returns:
            The id of the pre key referenced in the header.
        """

        return self.__pre_key_id

    def is_network_instance(self) -> bool:
        """
        Returns:
            Returns whether this is a network instance. A network instance has all fields filled except for
            the signed pre key and pre key byte data. The missing byte data can be restored by looking it up
            from storage using the respective ids.
        """

        return self.__header.signed_pre_key == b"" and self.__header.pre_key == b""

    def serialize(self, authenticated_message: bytes) -> bytes:
        """
        Args:
            authenticated_message: The serialized OMEMOAuthenticatedMessage message structure to include with
                the key exchange information.

        Returns:
            A serialized OMEMOKeyExchange message structure representing the content of this instance.

        Raises:
            ValueError: if the serialized OMEMOAuthenticatedMessage is malformed.
        """

        try:
            authenticated_message_parsed = OMEMOAuthenticatedMessage.FromString(authenticated_message)
        except google.protobuf.message.DecodeError as e:
            raise ValueError() from e

        return OMEMOKeyExchange(
            pk_id=self.__pre_key_id,
            spk_id=self.__signed_pre_key_id,
            ik=self.__header.identity_key,
            ek=self.__header.ephemeral_key,
            message=authenticated_message_parsed
        ).SerializeToString()

    @staticmethod
    def parse(key_exchange: bytes) -> Tuple[KeyExchangeImpl, bytes]:
        """
        Args:
            key_exchange: A serialized OMEMOKeyExchange message structure.

        Returns:
            An instance of this class, parsed from the OMEMOKeyExchange, and the serialized
            OMEMOAuthenticatedMessage extracted from the OMEMOKeyExchange.

        Raises:
            ValueError: if the data is malformed.

        Warning:
            The OMEMOKeyExchange message structure only contains the ids of the signed pre key and the pre key
            used for the key exchange, not the full public keys. Since the job of this method is just parsing,
            the X3DH header is initialized without the public keys here, and the code using instances of this
            class has to handle the public key lookup from the ids. Use :attr:`header_filled` to check whether
            the header is filled with the public keys.
        """

        try:
            parsed = OMEMOKeyExchange.FromString(key_exchange)
        except google.protobuf.message.DecodeError as e:
            raise ValueError() from e

        return KeyExchangeImpl(
            x3dh.Header(parsed.ik, parsed.ek, b"", b""),
            parsed.spk_id,
            parsed.pk_id
        ), parsed.message.SerializeToString()


class SessionImpl(Session):
    """
    :class:`~omemo.session.Session` implementation as a simple storage type.
    """

    def __init__(
        self,
        bare_jid: str,
        device_id: int,
        initiation: Initiation,
        key_exchange: KeyExchangeImpl,
        associated_data: bytes,
        double_ratchet: DoubleRatchetImpl,
        confirmed: bool = False
    ):
        """
        Args:
            bare_jid: The bare JID of the other party.
            device_id: The device id of the specific device of the other party.
            initiation: Whether this session was built through active or passive session initiation.
            key_exchange: The key exchange information to store in this instance.
            associated_data: The associated data to store in this instance.
            double_ratchet: The Double Ratchet to store in this instance.
            confirmed: Whether the session was confirmed, i.e. whether a message was decrypted after actively
                initiating the session. Leave this at the default value for passively initiated sessions.
        """

        self.__bare_jid = bare_jid
        self.__device_id = device_id
        self.__initiation = initiation
        self.__key_exchange = key_exchange
        self.__associated_data = associated_data
        self.__double_ratchet = double_ratchet
        self.__confirmed = confirmed

    @property
    def namespace(self) -> str:
        return NAMESPACE

    @property
    def bare_jid(self) -> str:
        return self.__bare_jid

    @property
    def device_id(self) -> int:
        return self.__device_id

    @property
    def initiation(self) -> Initiation:
        return self.__initiation

    @property
    def confirmed(self) -> bool:
        return self.__confirmed

    @property
    def key_exchange(self) -> KeyExchangeImpl:
        return self.__key_exchange

    @property
    def receiving_chain_length(self) -> Optional[int]:
        return self.__double_ratchet.receiving_chain_length

    @property
    def sending_chain_length(self) -> int:
        return self.__double_ratchet.sending_chain_length

    @property
    def associated_data(self) -> bytes:
        """
        Returns:
            The associated data held by this instance.
        """

        return self.__associated_data

    @property
    def double_ratchet(self) -> DoubleRatchetImpl:
        """
        Returns:
            The Double Ratchet held by this instance.
        """

        return self.__double_ratchet

    def confirm(self) -> None:
        """
        Mark this session as confirmed.
        """

        self.__confirmed = True


class Twomemo(Backend):
    """
    :class:`~omemo.backend.Backend` implementation providing OMEMO in the `urn:xmpp:omemo:2` namespace.
    """

    def __init__(
        self,
        storage: Storage,
        max_num_per_session_skipped_keys: int = 1000,
        max_num_per_message_skipped_keys: Optional[int] = None
    ) -> None:
        """
        Args:
            storage: The storage to store backend-specific data in. Note that all data keys are prefixed with
                the backend namespace to avoid name clashes between backends.
            max_num_per_session_skipped_keys: The maximum number of skipped message keys to keep around per
                session. Once the maximum is reached, old message keys are deleted to make space for newer
                ones. Accessible via :attr:`max_num_per_session_skipped_keys`.
            max_num_per_message_skipped_keys: The maximum number of skipped message keys to accept in a single
                message. When set to ``None`` (the default), this parameter defaults to the per-session
                maximum (i.e. the value of the ``max_num_per_session_skipped_keys`` parameter). This parameter
                may only be 0 if the per-session maximum is 0, otherwise it must be a number between 1 and the
                per-session maximum. Accessible via :attr:`max_num_per_message_skipped_keys`.
        """

        super().__init__(max_num_per_session_skipped_keys, max_num_per_message_skipped_keys)

        self.__storage = storage

    async def __get_state(self) -> StateImpl:
        """
        Returns:
            The loaded or newly created X3DH state.
        """

        def check_type(value: JSONType) -> x3dh.types.JSONObject:
            if isinstance(value, dict):
                return cast(x3dh.types.JSONObject, value)

            raise TypeError(
                f"Stored StateImpl under key /{self.namespace}/x3dh corrupt: not a JSON object: {value}"
            )

        state, _ = (await self.__storage.load(
            f"/{self.namespace}/x3dh"
        )).fmap(check_type).fmap(lambda serialized: StateImpl.from_json(
            serialized,
            x3dh.IdentityKeyFormat.ED_25519,
            x3dh.HashFunction.SHA_256,
            StateImpl.INFO
        )).maybe((None, False))

        if state is None:
            identity_key_pair = await IdentityKeyPair.get(self.__storage)

            state = StateImpl.create(
                x3dh.IdentityKeyFormat.ED_25519,
                x3dh.HashFunction.SHA_256,
                StateImpl.INFO,
                (
                    x3dh.identity_key_pair.IdentityKeyPairSeed(identity_key_pair.seed)
                    if isinstance(identity_key_pair, IdentityKeyPairSeed)
                    else x3dh.identity_key_pair.IdentityKeyPairPriv(identity_key_pair.as_priv().priv)
                )
            )

            await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

        return state

    @property
    def namespace(self) -> str:
        return NAMESPACE

    async def load_session(self, bare_jid: str, device_id: int) -> Optional[SessionImpl]:
        def check_type(value: JSONType) -> doubleratchet.types.JSONObject:
            if isinstance(value, dict):
                return cast(doubleratchet.types.JSONObject, value)

            raise TypeError(
                f"Stored DoubleRatchetImpl under key"
                f" /{self.namespace}/{bare_jid}/{device_id}/double_ratchet corrupt: not a JSON object:"
                f" {value}"
            )

        try:
            double_ratchet = (await self.__storage.load(
                f"/{self.namespace}/{bare_jid}/{device_id}/double_ratchet"
            )).fmap(check_type).fmap(lambda serialized: DoubleRatchetImpl.from_json(
                serialized,
                diffie_hellman_ratchet_curve25519.DiffieHellmanRatchet,
                RootChainKDFImpl,
                MessageChainKDFImpl,
                DoubleRatchetImpl.MESSAGE_CHAIN_CONSTANT,
                self.max_num_per_message_skipped_keys,
                self.max_num_per_session_skipped_keys,
                AEADImpl
            )).maybe(None)
        except doubleratchet.InconsistentSerializationException:
            return None

        if double_ratchet is None:
            return None

        initiation = Initiation((await self.__storage.load_primitive(
            f"/{self.namespace}/{bare_jid}/{device_id}/initiation",
            str
        )).from_just())

        identity_key = (await self.__storage.load_bytes(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/identity_key"
        )).from_just()

        ephemeral_key = (await self.__storage.load_bytes(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/ephemeral_key"
        )).from_just()

        signed_pre_key = (await self.__storage.load_bytes(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/signed_pre_key"
        )).from_just()

        signed_pre_key_id = (await self.__storage.load_primitive(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/signed_pre_key_id",
            int
        )).from_just()

        pre_key = (await self.__storage.load_bytes(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/pre_key"
        )).from_just()

        pre_key_id = (await self.__storage.load_primitive(
            f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/pre_key_id",
            int
        )).from_just()

        associated_data = (await self.__storage.load_bytes(
            f"/{self.namespace}/{bare_jid}/{device_id}/associated_data"
        )).from_just()

        confirmed = (await self.__storage.load_primitive(
            f"/{self.namespace}/{bare_jid}/{device_id}/confirmed",
            bool
        )).from_just()

        return SessionImpl(bare_jid, device_id, initiation, KeyExchangeImpl(
            x3dh.Header(identity_key, ephemeral_key, signed_pre_key, pre_key),
            signed_pre_key_id,
            pre_key_id
        ), associated_data, double_ratchet, confirmed)

    async def store_session(self, session: Session) -> None:
        assert isinstance(session, SessionImpl)

        assert session.key_exchange.header.pre_key is not None

        await self.__storage.store(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/initiation",
            session.initiation.name
        )

        await self.__storage.store_bytes(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/identity_key",
            session.key_exchange.header.identity_key
        )

        await self.__storage.store_bytes(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/ephemeral_key",
            session.key_exchange.header.ephemeral_key
        )

        await self.__storage.store_bytes(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/signed_pre_key",
            session.key_exchange.header.signed_pre_key
        )

        await self.__storage.store(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/signed_pre_key_id",
            session.key_exchange.signed_pre_key_id
        )

        await self.__storage.store_bytes(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/pre_key",
            session.key_exchange.header.pre_key
        )

        await self.__storage.store(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/key_exchange/pre_key_id",
            session.key_exchange.pre_key_id
        )

        await self.__storage.store_bytes(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/associated_data",
            session.associated_data
        )

        await self.__storage.store(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/double_ratchet",
            session.double_ratchet.json
        )

        await self.__storage.store(
            f"/{self.namespace}/{session.bare_jid}/{session.device_id}/confirmed",
            session.confirmed
        )

        # Keep track of bare JIDs with stored sessions
        bare_jids = set((await self.__storage.load_list(f"/{self.namespace}/bare_jids", str)).maybe([]))
        bare_jids.add(session.bare_jid)
        await self.__storage.store(f"/{self.namespace}/bare_jids", list(bare_jids))

        # Keep track of device ids with stored sessions
        device_ids = set((await self.__storage.load_list(
            f"/{self.namespace}/{session.bare_jid}/device_ids",
            int
        )).maybe([]))
        device_ids.add(session.device_id)
        await self.__storage.store(f"/{self.namespace}/{session.bare_jid}/device_ids", list(device_ids))

    async def build_session_active(
        self,
        bare_jid: str,
        device_id: int,
        bundle: Bundle,
        plain_key_material: PlainKeyMaterial
    ) -> Tuple[SessionImpl, EncryptedKeyMaterialImpl]:
        assert isinstance(bundle, BundleImpl)
        assert isinstance(plain_key_material, PlainKeyMaterialImpl)

        try:
            state = await self.__get_state()
            shared_secret, associated_data, header = await state.get_shared_secret_active(bundle.bundle)
        except x3dh.KeyAgreementException as e:
            raise KeyExchangeFailed() from e

        assert header.pre_key is not None

        double_ratchet, encrypted_message = await DoubleRatchetImpl.encrypt_initial_message(
            diffie_hellman_ratchet_curve25519.DiffieHellmanRatchet,
            RootChainKDFImpl,
            MessageChainKDFImpl,
            DoubleRatchetImpl.MESSAGE_CHAIN_CONSTANT,
            self.max_num_per_message_skipped_keys,
            self.max_num_per_session_skipped_keys,
            AEADImpl,
            shared_secret,
            bundle.bundle.signed_pre_key,
            plain_key_material.key + plain_key_material.auth_tag,
            associated_data
        )

        session = SessionImpl(
            bare_jid,
            device_id,
            Initiation.ACTIVE,
            KeyExchangeImpl(
                header,
                bundle.signed_pre_key_id,
                bundle.pre_key_ids[header.pre_key]
            ),
            associated_data,
            double_ratchet
        )

        encrypted_key_material = EncryptedKeyMaterialImpl(bare_jid, device_id, encrypted_message)

        return session, encrypted_key_material

    async def build_session_passive(
        self,
        bare_jid: str,
        device_id: int,
        key_exchange: KeyExchange,
        encrypted_key_material: EncryptedKeyMaterial
    ) -> Tuple[SessionImpl, PlainKeyMaterialImpl]:
        assert isinstance(key_exchange, KeyExchangeImpl)
        assert isinstance(encrypted_key_material, EncryptedKeyMaterialImpl)

        state = await self.__get_state()

        if key_exchange.is_network_instance():
            # Perform lookup of the signed pre key and pre key public keys in case the header is not filled
            signed_pre_keys_by_id = { v: k for k, v in (await self.__get_signed_pre_key_ids()).items() }
            if key_exchange.signed_pre_key_id not in signed_pre_keys_by_id:
                raise KeyExchangeFailed(f"No signed pre key with id {key_exchange.signed_pre_key_id} known.")

            pre_keys_by_id = { v: k for k, v in (await self.__get_pre_key_ids()).items() }
            if key_exchange.pre_key_id not in pre_keys_by_id:
                raise KeyExchangeFailed(f"No pre key with id {key_exchange.pre_key_id} known.")

            # Update the key exchange information with the filled header
            key_exchange = KeyExchangeImpl(
                x3dh.Header(
                    key_exchange.header.identity_key,
                    key_exchange.header.ephemeral_key,
                    signed_pre_keys_by_id[key_exchange.signed_pre_key_id],
                    pre_keys_by_id[key_exchange.pre_key_id]
                ),
                key_exchange.signed_pre_key_id,
                key_exchange.pre_key_id
            )

        try:
            shared_secret, associated_data, signed_pre_key = await state.get_shared_secret_passive(
                key_exchange.header
            )
        except x3dh.KeyAgreementException as e:
            raise KeyExchangeFailed() from e

        try:
            double_ratchet, decrypted_message = await DoubleRatchetImpl.decrypt_initial_message(
                diffie_hellman_ratchet_curve25519.DiffieHellmanRatchet,
                RootChainKDFImpl,
                MessageChainKDFImpl,
                DoubleRatchetImpl.MESSAGE_CHAIN_CONSTANT,
                self.max_num_per_message_skipped_keys,
                self.max_num_per_session_skipped_keys,
                AEADImpl,
                shared_secret,
                signed_pre_key.priv,
                encrypted_key_material.encrypted_message,
                associated_data
            )
        except Exception as e:
            raise DecryptionFailed(
                "Decryption of the initial message as part of passive session building failed."
            ) from e

        session = SessionImpl(
            bare_jid,
            device_id,
            Initiation.PASSIVE,
            key_exchange,
            associated_data,
            double_ratchet
        )

        plain_key_material = PlainKeyMaterialImpl(
            decrypted_message[:PlainKeyMaterialImpl.KEY_LENGTH],
            decrypted_message[PlainKeyMaterialImpl.KEY_LENGTH:]
        )

        return session, plain_key_material

    async def encrypt_plaintext(self, plaintext: bytes) -> Tuple[ContentImpl, PlainKeyMaterialImpl]:
        # Generate KEY_LENGTH bytes of cryptographically secure random data for the key
        key = secrets.token_bytes(PlainKeyMaterialImpl.KEY_LENGTH)

        # Derive 80 bytes from the key using HKDF-SHA-256
        key_material = await CryptoProviderImpl.hkdf_derive(
            hash_function=HashFunction.SHA_256,
            length=80,
            salt=b"\x00" * 32,
            info="OMEMO Payload".encode("ASCII"),
            key_material=key
        )

        # Split those 80 bytes into an encryption key, authentication key and an initialization vector
        encryption_key = key_material[:32]
        authentication_key = key_material[32:64]
        initialization_vector = key_material[64:]

        # Encrypt the plaintext using AES-256 (the 256 bit are implied by the key size) in CBC mode and the
        # previously created key and IV, after padding it with PKCS#7
        ciphertext = await CryptoProviderImpl.aes_cbc_encrypt(
            encryption_key,
            initialization_vector,
            plaintext
        )

        # Calculate the authentication tag and truncate it to AUTHENTICATION_TAG_TRUNCATED_LENGTH bytes
        auth_tag = (await CryptoProviderImpl.hmac_calculate(
            authentication_key,
            HashFunction.SHA_256,
            ciphertext
        ))[:AEADImpl.AUTHENTICATION_TAG_TRUNCATED_LENGTH]

        return ContentImpl(ciphertext), PlainKeyMaterialImpl(key, auth_tag)

    async def encrypt_empty(self) -> Tuple[ContentImpl, PlainKeyMaterialImpl]:
        return ContentImpl.make_empty(), PlainKeyMaterialImpl.make_empty()

    async def encrypt_key_material(
        self,
        session: Session,
        plain_key_material: PlainKeyMaterial
    ) -> EncryptedKeyMaterialImpl:
        assert isinstance(session, SessionImpl)
        assert isinstance(plain_key_material, PlainKeyMaterialImpl)

        return EncryptedKeyMaterialImpl(
            session.bare_jid,
            session.device_id,
            await session.double_ratchet.encrypt_message(
                plain_key_material.key + plain_key_material.auth_tag,
                session.associated_data
            )
        )

    async def decrypt_plaintext(self, content: Content, plain_key_material: PlainKeyMaterial) -> bytes:
        assert isinstance(content, ContentImpl)
        assert isinstance(plain_key_material, PlainKeyMaterialImpl)

        assert not content.empty

        # Derive 80 bytes from the key using HKDF-SHA-256
        key_material = await CryptoProviderImpl.hkdf_derive(
            hash_function=HashFunction.SHA_256,
            length=80,
            salt=b"\x00" * 32,
            info="OMEMO Payload".encode("ASCII"),
            key_material=plain_key_material.key
        )

        # Split those 80 bytes into an encryption key, authentication key and an initialization vector
        decryption_key = key_material[:32]
        authentication_key = key_material[32:64]
        initialization_vector = key_material[64:]

        # Calculate and verify the authentication tag after truncating it to
        # AUTHENTICATION_TAG_TRUNCATED_LENGTH bytes
        auth_tag = (await CryptoProviderImpl.hmac_calculate(
            authentication_key,
            HashFunction.SHA_256,
            content.ciphertext
        ))[:AEADImpl.AUTHENTICATION_TAG_TRUNCATED_LENGTH]

        if auth_tag != plain_key_material.auth_tag:
            raise DecryptionFailed("Authentication tag verification failed.")

        # Decrypt the plaintext using AES-256 (the 256 bit are implied by the key size) in CBC mode and the
        # previously created key and IV, and unpad the resulting plaintext with PKCS#7
        return await CryptoProviderImpl.aes_cbc_decrypt(
            decryption_key,
            initialization_vector,
            content.ciphertext
        )

    async def decrypt_key_material(
        self,
        session: Session,
        encrypted_key_material: EncryptedKeyMaterial
    ) -> PlainKeyMaterialImpl:
        assert isinstance(session, SessionImpl)
        assert isinstance(encrypted_key_material, EncryptedKeyMaterialImpl)

        try:
            decrypted_message = await session.double_ratchet.decrypt_message(
                encrypted_key_material.encrypted_message,
                session.associated_data
            )
        except Exception as e:
            raise DecryptionFailed("Key material decryption failed.") from e

        session.confirm()

        return PlainKeyMaterialImpl(
            decrypted_message[:PlainKeyMaterialImpl.KEY_LENGTH],
            decrypted_message[PlainKeyMaterialImpl.KEY_LENGTH:]
        )

    async def signed_pre_key_age(self) -> int:
        return (await self.__get_state()).signed_pre_key_age()

    async def rotate_signed_pre_key(self) -> None:
        state = await self.__get_state()

        state.rotate_signed_pre_key()

        await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

    async def hide_pre_key(self, session: Session) -> bool:
        assert isinstance(session, SessionImpl)

        # This method is only called with KeyExchangeImpl instances that have the pre key byte data set. We do
        # not have to worry about the field containing a filler value and the assertion is merely there to
        # satisfy the type system.
        assert session.key_exchange.header.pre_key is not None

        state = await self.__get_state()

        hidden = state.hide_pre_key(session.key_exchange.header.pre_key)

        await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

        return hidden

    async def delete_pre_key(self, session: Session) -> bool:
        assert isinstance(session, SessionImpl)

        # This method is only called with KeyExchangeImpl instances that have the pre key byte data set. We do
        # not have to worry about the field containing a filler value and the assertion is merely there to
        # satisfy the type system.
        assert session.key_exchange.header.pre_key is not None

        state = await self.__get_state()

        deleted = state.delete_pre_key(session.key_exchange.header.pre_key)

        await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

        return deleted

    async def delete_hidden_pre_keys(self) -> None:
        state = await self.__get_state()

        state.delete_hidden_pre_keys()

        await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

    async def get_num_visible_pre_keys(self) -> int:
        return (await self.__get_state()).get_num_visible_pre_keys()

    async def generate_pre_keys(self, num_pre_keys: int) -> None:
        state = await self.__get_state()

        state.generate_pre_keys(num_pre_keys)

        await self.__storage.store(f"/{self.namespace}/x3dh", state.json)

    async def get_bundle(self, bare_jid: str, device_id: int) -> BundleImpl:
        bundle = (await self.__get_state()).bundle

        return BundleImpl(
            bare_jid,
            device_id,
            bundle,
            (await self.__get_signed_pre_key_ids())[bundle.signed_pre_key],
            {
                pre_key: pre_key_id
                for pre_key, pre_key_id
                in (await self.__get_pre_key_ids()).items()
                if pre_key in bundle.pre_keys
            }
        )

    async def purge(self) -> None:
        for bare_jid in (await self.__storage.load_list(f"/{self.namespace}/bare_jids", str)).maybe([]):
            await self.purge_bare_jid(bare_jid)

        await self.__storage.delete(f"/{self.namespace}/bare_jids")
        await self.__storage.delete(f"/{self.namespace}/x3dh")
        await self.__storage.delete(f"/{self.namespace}/signed_pre_key_ids")
        await self.__storage.delete(f"/{self.namespace}/pre_key_ids")
        await self.__storage.delete(f"/{self.namespace}/pre_key_id_counter")

    async def purge_bare_jid(self, bare_jid: str) -> None:
        storage = self.__storage

        for device_id in (await storage.load_list(f"/{self.namespace}/{bare_jid}/device_ids", int)).maybe([]):
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/initiation")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/identity_key")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/ephemeral_key")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/signed_pre_key")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/signed_pre_key_id")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/pre_key")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/key_exchange/pre_key_id")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/associated_data")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/double_ratchet")
            await storage.delete(f"/{self.namespace}/{bare_jid}/{device_id}/confirmed")

        await storage.delete(f"/{self.namespace}/{bare_jid}/device_ids")

        bare_jids = set((await storage.load_list(f"/{self.namespace}/bare_jids", str)).maybe([]))
        bare_jids.remove(bare_jid)
        await storage.store(f"/{self.namespace}/bare_jids", list(bare_jids))

    async def __get_signed_pre_key_ids(self) -> Dict[bytes, int]:
        """
        Assigns an id to each signed pre key currently available in the X3DH state, both the current signed
        pre key and the old signed pre key that is kept around for one more rotation period. Once assigned to
        a signed pre key, its id will never change.

        Returns:
            The mapping from signed pre key to id.
        """

        state = await self.__get_state()

        signed_pre_key = state.bundle.signed_pre_key
        old_signed_pre_key = state.old_signed_pre_key

        # Load the existing signed pre key ids from the storage
        signed_pre_key_ids = {
            base64.b64decode(signed_pre_key_b64): signed_pre_key_id
            for signed_pre_key_b64, signed_pre_key_id
            in (await self.__storage.load_dict(
                f"/{self.namespace}/signed_pre_key_ids",
                int
            )).maybe({}).items()
        }

        # Take note of the highest id that was assigned, default to 0 if no ids were assigned yet
        signed_pre_key_id_counter = max(
            signed_pre_key_id
            for _, signed_pre_key_id
            in signed_pre_key_ids.items()
        ) if len(signed_pre_key_ids) > 0 else 0

        # Prepare the dictionary to hold updated signed pre key ids
        new_signed_pre_key_ids: Dict[bytes, int] = {}

        # Assign the next highest id to the signed pre key, if there is no id assigned to it yet.
        signed_pre_key_id_counter += 1
        new_signed_pre_key_ids[signed_pre_key] = signed_pre_key_ids.get(
            signed_pre_key,
            signed_pre_key_id_counter
        )

        # Assign the next highest id to the old signed pre key, if there is no id assigned to it yet. This
        # should never happen, since the old signed pre key should have been assigned an id when it was the
        # (non-old) signed pre key, however there might be edge cases of the signed pre key rotating twice
        # before the assigned ids are updated.
        if old_signed_pre_key is not None:
            signed_pre_key_id_counter += 1
            new_signed_pre_key_ids[old_signed_pre_key] = signed_pre_key_ids.get(
                old_signed_pre_key,
                signed_pre_key_id_counter
            )

        # If the ids have changed, store them
        if new_signed_pre_key_ids != signed_pre_key_ids:
            await self.__storage.store(f"/{self.namespace}/signed_pre_key_ids", {
                base64.b64encode(signed_pre_key).decode("ASCII"): signed_pre_key_id
                for signed_pre_key, signed_pre_key_id
                in new_signed_pre_key_ids.items()
            })

        return new_signed_pre_key_ids

    async def __get_pre_key_ids(self) -> Dict[bytes, int]:
        """
        Assigns an id to each pre key currently available in the X3DH state, both hidden and visible pre keys.
        Once assigned to a pre key, its id will never change.

        Returns:
            The mapping from pre key to id.
        """

        state = await self.__get_state()

        pre_keys = state.bundle.pre_keys | state.hidden_pre_keys

        # Load the existing pre key ids from the storage
        pre_key_ids = {
            base64.b64decode(pre_key_b64): pre_key_id
            for pre_key_b64, pre_key_id
            in (await self.__storage.load_dict(f"/{self.namespace}/pre_key_ids", int)).maybe({}).items()
        }

        # Load the pre key id counter from the storage
        pre_key_id_counter = (await self.__storage.load_primitive(
            f"/{self.namespace}/pre_key_id_counter",
            int
        )).maybe(0)

        # Prepare the dictionary to hold updated pre key ids
        new_pre_key_ids: Dict[bytes, int] = {}

        # Assign the next highest id to each pre key if there is no existing id assigned to it
        for pre_key in pre_keys:
            pre_key_id_counter += 1
            new_pre_key_ids[pre_key] = pre_key_ids.get(pre_key, pre_key_id_counter)

        # If the ids have changed, store them
        if new_pre_key_ids != pre_key_ids:
            await self.__storage.store(f"/{self.namespace}/pre_key_ids", {
                base64.b64encode(pre_key).decode("ASCII"): pre_key_id
                for pre_key, pre_key_id
                in new_pre_key_ids.items()
            })

            await self.__storage.store(f"/{self.namespace}/pre_key_id_counter", pre_key_id_counter)

        return new_pre_key_ids