#!/usr/bin/env python3

# Libervia plugin for Jingle Encrypted Transports
# Copyright (C) 2009-2022 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from base64 import b64encode
from functools import partial
import io
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from twisted.words.protocols.jabber import error, jid, xmlstream
from twisted.words.xish import domish
from wokkel import disco, iwokkel
from zope.interface import implementer
from cryptography.exceptions import AlreadyFinalized
from cryptography.hazmat import backends
from cryptography.hazmat.primitives import ciphers
from cryptography.hazmat.primitives.ciphers import Cipher, CipherContext, modes
from cryptography.hazmat.primitives.padding import PKCS7, PaddingContext

from sat.core import exceptions
from sat.core.constants import Const as C
from sat.core.core_types import SatXMPPEntity
from sat.core.i18n import _
from sat.core.log import getLogger
from sat.tools import xml_tools

try:
    import oldmemo
    import oldmemo.etree
except ImportError as import_error:
    raise exceptions.MissingModule(
        "You are missing one or more package required by the OMEMO plugin. Please"
        " download/install the pip packages 'oldmemo'."
    ) from import_error


log = getLogger(__name__)

IMPORT_NAME = "XEP-0391"

PLUGIN_INFO = {
    C.PI_NAME: "Jingle Encrypted Transports",
    C.PI_IMPORT_NAME: IMPORT_NAME,
    C.PI_TYPE: C.PLUG_TYPE_XEP,
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: ["XEP-0391", "XEP-0396"],
    C.PI_DEPENDENCIES: ["XEP-0166", "XEP-0384"],
    C.PI_MAIN: "JET",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""End-to-end encryption of Jingle transports"""),
}

NS_JET = "urn:xmpp:jingle:jet:0"
NS_JET_OMEMO = "urn:xmpp:jingle:jet-omemo:0"


class JET:
    namespace = NS_JET

    def __init__(self, host):
        log.info(_("XEP-0391 (Pubsub Attachments) plugin initialization"))
        host.registerNamespace("jet", NS_JET)
        self.host = host
        self._o = host.plugins["XEP-0384"]
        self._j = host.plugins["XEP-0166"]
        host.trigger.add(
            "XEP-0166_initiate_elt_built",
            self._on_initiate_elt_build
        )
        host.trigger.add(
            "XEP-0166_on_session_initiate",
            self._on_session_initiate
        )
        host.trigger.add(
            "XEP-0234_jingle_handler",
            self._add_encryption_filter
        )
        host.trigger.add(
            "XEP-0234_file_receiving_request_conf",
            self._add_encryption_filter
        )

    def getHandler(self, client):
        return JET_Handler()

    async def _on_initiate_elt_build(
        self,
        client: SatXMPPEntity,
        session: Dict[str, Any],
        iq_elt: domish.Element,
        jingle_elt: domish.Element
    ) -> bool:
        if client.encryption.get_namespace(
               session["peer_jid"].userhostJID()
           ) != self._o.NS_OLDMEMO:
            return True
        for content_elt in jingle_elt.elements(self._j.namespace, "content"):
            content_data = session["contents"][content_elt["name"]]
            security_elt = content_elt.addElement((NS_JET, "security"))
            security_elt["name"] = content_elt["name"]
            # XXX: for now only OLDMEMO is supported, thus we do it directly here. If some
            #   other are supported in the future, a plugin registering mechanism will be
            #   implemented.
            cipher = "urn:xmpp:ciphers:aes-128-gcm-nopadding"
            enc_type = "eu.siacs.conversations.axolotl"
            security_elt["cipher"] = cipher
            security_elt["type"] = enc_type
            encryption_data = content_data["encryption"] = {
                "cipher": cipher,
                "type": enc_type
            }
            session_manager = await self._o.get_session_manager(client.profile)
            try:
                messages, encryption_errors = await session_manager.encrypt(
                    frozenset({session["peer_jid"].userhost()}),
                    # the value seems to be the commonly used value
                    { self._o.NS_OLDMEMO: b" " },
                    backend_priority_order=[ self._o.NS_OLDMEMO ],
                    identifier = client.jid.userhost()
                )
            except Exception as e:
                log.error("Can't generate IV and keys: {e}")
                raise e
            message, plain_key_material = next(iter(messages.items()))
            iv, key = message.content.initialization_vector, plain_key_material.key
            content_data["encryption"].update({
                "iv": iv,
                "key": key
            })
            encrypted_elt = xml_tools.et_elt_2_domish_elt(
                oldmemo.etree.serialize_message(message)
            )
            security_elt.addChild(encrypted_elt)
        return True

    async def _on_session_initiate(
        self,
        client: SatXMPPEntity,
        session: Dict[str, Any],
        iq_elt: domish.Element,
        jingle_elt: domish.Element
    ) -> bool:
        if client.encryption.get_namespace(
               session["peer_jid"].userhostJID()
           ) != self._o.NS_OLDMEMO:
            return True
        for content_elt in jingle_elt.elements(self._j.namespace, "content"):
            content_data = session["contents"][content_elt["name"]]
            security_elt = next(content_elt.elements(NS_JET, "security"), None)
            if security_elt is None:
                continue
            encrypted_elt = next(
                security_elt.elements(self._o.NS_OLDMEMO, "encrypted"), None
            )
            if encrypted_elt is None:
                log.warning(
                    "missing <encrypted> element, can't decrypt: {security_elt.toXml()}"
                )
                continue
            session_manager = await self._o.get_session_manager(client.profile)
            try:
                message = await oldmemo.etree.parse_message(
                    xml_tools.domish_elt_2_et_elt(encrypted_elt, False),
                    session["peer_jid"].userhost(),
                    client.jid.userhost(),
                    session_manager
                )
                __, __, plain_key_material = await session_manager.decrypt(message)
            except Exception as e:
                log.warning(f"Can't get IV and key: {e}\n{security_elt.toXml()}")
                continue
            try:
                content_data["encryption"] = {
                    "cipher": security_elt["cipher"],
                    "type": security_elt["type"],
                    "iv": message.content.initialization_vector,
                    "key": plain_key_material.key
                }
            except KeyError as e:
                log.warning(f"missing data, can't decrypt: {e}")
                continue

        return True

    def __encrypt(
        self,
        data: bytes,
        encryptor: CipherContext,
        data_cb: Callable
    ) -> bytes:
        data_cb(data)
        if data:
            return encryptor.update(data)
        else:
            try:
                return encryptor.finalize() + encryptor.tag
            except AlreadyFinalized:
                return b''

    def __decrypt(
        self,
        data: bytes,
        buffer: list[bytes],
        decryptor: CipherContext,
        data_cb: Callable
    ) -> bytes:
        buffer.append(data)
        data = b''.join(buffer)
        buffer.clear()
        if len(data) > 16:
            decrypted = decryptor.update(data[:-16])
            data_cb(decrypted)
        else:
            decrypted = b''
        buffer.append(data[-16:])
        return decrypted

    def __decrypt_finalize(
        self,
        file_obj: io.BytesIO,
        buffer: list[bytes],
        decryptor: CipherContext,
    ) -> None:
        tag = b''.join(buffer)
        file_obj.write(decryptor.finalize_with_tag(tag))

    async def _add_encryption_filter(
        self,
        client: SatXMPPEntity,
        session: Dict[str, Any],
        content_data: Dict[str, Any],
        elt: domish.Element
    ) -> bool:
        file_obj = content_data["stream_object"].file_obj
        try:
            encryption_data=content_data["encryption"]
        except KeyError:
            return True
        cipher = ciphers.Cipher(
            ciphers.algorithms.AES(encryption_data["key"]),
            modes.GCM(encryption_data["iv"]),
            backend=backends.default_backend(),
        )
        if file_obj.mode == "wb":
            # we are receiving a file
            buffer = []
            decryptor = cipher.decryptor()
            file_obj.pre_close_cb = partial(
                self.__decrypt_finalize,
                file_obj=file_obj,
                buffer=buffer,
                decryptor=decryptor
            )
            file_obj.data_cb = partial(
                self.__decrypt,
                buffer=buffer,
                decryptor=decryptor,
                data_cb=file_obj.data_cb
            )
        else:
            # we are sending a file
            file_obj.data_cb = partial(
                self.__encrypt,
                encryptor=cipher.encryptor(),
                data_cb=file_obj.data_cb
            )

        return True


@implementer(iwokkel.IDisco)
class JET_Handler(xmlstream.XMPPHandler):

    def getDiscoInfo(self, requestor, service, nodeIdentifier=""):
        return [
            disco.DiscoFeature(NS_JET),
            disco.DiscoFeature(NS_JET_OMEMO),
        ]

    def getDiscoItems(self, requestor, service, nodeIdentifier=""):
        return []