#!/usr/bin/env python2
# -*- coding: utf-8 -*-

# SAT plugin for OTR encryption
# Copyright (C) 2009-2016 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/>.

# XXX: thanks to Darrik L Mazey for his documentation (https://blog.darmasoft.net/2013/06/30/using-pure-python-otr.html)
#      this implentation is based on it

from sat.core.i18n import _, D_
from sat.core.constants import Const as C
from sat.core.log import getLogger
from sat.core import exceptions
log = getLogger(__name__)
from sat.tools import xml_tools
from twisted.words.protocols.jabber import jid
from twisted.python import failure
from twisted.internet import defer
from sat.memory import persistent
import potr
import copy
import time
import uuid


PLUGIN_INFO = {
    "name": "OTR",
    "import_name": "OTR",
    "type": "SEC",
    "protocols": [],
    "dependencies": [],
    "main": "OTR",
    "handler": "no",
    "description": _(u"""Implementation of OTR""")
}

NS_OTR = "otr_plugin"
PRIVATE_KEY = "PRIVATE KEY"
OTR_MENU = D_(u'OTR')
AUTH_TXT = D_(u"To authenticate your correspondent, you need to give your below fingerprint *BY AN EXTERNAL CANAL* (i.e. not in this chat), and check that the one he gives you is the same as below. If there is a mismatch, there can be a spy between you!")
DROP_TXT = D_(u"You private key is used to encrypt messages for your correspondent, nobody except you must know it, if you are in doubt, you should drop it!\n\nAre you sure you want to drop your private key?")
# NO_LOG_AND = D_(u"/!\\Your history is not logged anymore, and")   # FIXME: not used at the moment
NO_ADV_FEATURES = D_(u"Most of advanced features are disabled !")

DEFAULT_POLICY_FLAGS = {
  'ALLOW_V1':False,
  'ALLOW_V2':True,
  'REQUIRE_ENCRYPTION':True,
}

OTR_STATE_TRUSTED = 'trusted'
OTR_STATE_UNTRUSTED = 'untrusted'
OTR_STATE_UNENCRYPTED = 'unencrypted'
OTR_STATE_ENCRYPTED = 'encrypted'


class Context(potr.context.Context):
    def __init__(self, host, account, other_jid):
        super(Context, self).__init__(account, other_jid)
        self.host = host

    def getPolicy(self, key):
        if key in DEFAULT_POLICY_FLAGS:
            return DEFAULT_POLICY_FLAGS[key]
        else:
            return False

    def inject(self, msg_str, appdata=None):
        assert isinstance(self.peer, jid.JID)
        msg = msg_str.decode('utf-8')
        client = self.user.client
        log.debug(u'inject(%s, appdata=%s, to=%s)' % (msg, appdata, self.peer))
        mess_data = {
                     'from': client.jid,
                     'to': self.peer,
                     'uid': unicode(uuid.uuid4()),
                     'message': {'': msg},
                     'subject': {},
                     'type': 'chat',
                     'extra': {},
                     'timestamp': time.time(),
                    }
        self.host.generateMessageXML(mess_data)
        client.xmlstream.send(mess_data['xml'])

    def setState(self, state):
        client = self.user.client
        old_state = self.state
        super(Context, self).setState(state)
        log.debug(u"setState: %s (old_state=%s)" % (state, old_state))

        if state == potr.context.STATE_PLAINTEXT:
            feedback = _(u"/!\\ conversation with %(other_jid)s is now UNENCRYPTED") % {'other_jid': self.peer.full()}
            self.host.bridge.otrState(OTR_STATE_UNENCRYPTED, self.peer.full(), client.profile)
        elif state == potr.context.STATE_ENCRYPTED:
            try:
                trusted = self.getCurrentTrust()
            except TypeError:
                trusted = False
            trusted_str = _(u"trusted") if trusted else _(u"untrusted")

            if old_state == potr.context.STATE_ENCRYPTED:
                feedback = D_(u"{trusted} OTR conversation with {other_jid} REFRESHED").format(
                    trusted = trusted_str,
                    other_jid = self.peer.full())
            else:
                feedback = D_(u"{trusted} encrypted OTR conversation started with {other_jid}\n{extra_info}").format(
                    trusted = trusted_str,
                    other_jid = self.peer.full(),
                    extra_info = NO_ADV_FEATURES)
            self.host.bridge.otrState(OTR_STATE_ENCRYPTED, self.peer.full(), client.profile)
        elif state == potr.context.STATE_FINISHED:
            feedback = D_(u"OTR conversation with {other_jid} is FINISHED").format(other_jid = self.peer.full())
            self.host.bridge.otrState(OTR_STATE_UNENCRYPTED, self.peer.full(), client.profile)
        else:
            log.error(D_(u"Unknown OTR state"))
            return

        self.host.bridge.messageNew(uid=unicode(uuid.uuid4()),
                                    timestamp=time.time(),
                                    from_jid=client.jid.full(),
                                    to_jid=self.peer.full(),
                                    message={u'': feedback},
                                    subject={},
                                    mess_type=C.MESS_TYPE_INFO,
                                    extra={},
                                    profile=client.profile)

    def disconnect(self):
        """Disconnect the session."""
        if self.state != potr.context.STATE_PLAINTEXT:
            super(Context, self).disconnect()

    def finish(self):
        """Finish the session - avoid to send any message but the user still has to end the session himself."""
        if self.state == potr.context.STATE_ENCRYPTED:
            self.processTLVs([potr.proto.DisconnectTLV()])


class Account(potr.context.Account):
    #TODO: manage trusted keys: if a fingerprint is not used anymore, we have no way to remove it from database yet (same thing for a correspondent jid)

    def __init__(self, host, client):
        log.debug(u"new account: %s" % client.jid)
        if not client.jid.resource:
            log.warning("Account created without resource")
        super(Account, self).__init__(unicode(client.jid), "xmpp", 1024)
        self.host = host
        self.client = client

    def loadPrivkey(self):
        log.debug(u"loadPrivkey")
        return self.privkey

    def savePrivkey(self):
        log.debug(u"savePrivkey")
        if self.privkey is None:
            raise exceptions.InternalError(_(u"Save is called but privkey is None !"))
        priv_key = self.privkey.serializePrivateKey().encode('hex')
        d = self.host.memory.encryptValue(priv_key, self.client.profile)
        def save_encrypted_key(encrypted_priv_key):
            self.client._otr_data[PRIVATE_KEY] = encrypted_priv_key
        d.addCallback(save_encrypted_key)

    def loadTrusts(self):
        trust_data = self.client._otr_data.get('trust', {})
        for jid_, jid_data in trust_data.iteritems():
            for fingerprint, trust_level in jid_data.iteritems():
                log.debug(u'setting trust for {jid}: [{fingerprint}] = "{trust_level}"'.format(jid=jid_, fingerprint=fingerprint, trust_level=trust_level))
                self.trusts.setdefault(jid.JID(jid_), {})[fingerprint] = trust_level

    def saveTrusts(self):
        log.debug(u"saving trusts for {profile}".format(profile=self.client.profile))
        log.debug(u"trusts = {}".format(self.client._otr_data['trust']))
        self.client._otr_data.force('trust')

    def setTrust(self, other_jid, fingerprint, trustLevel):
        try:
            trust_data = self.client._otr_data['trust']
        except KeyError:
            trust_data = {}
            self.client._otr_data['trust'] = trust_data
        jid_data = trust_data.setdefault(other_jid.full(), {})
        jid_data[fingerprint] = trustLevel
        super(Account, self).setTrust(other_jid, fingerprint, trustLevel)


class ContextManager(object):

    def __init__(self, host, client):
        self.host = host
        self.account = Account(host, client)
        self.contexts = {}

    def startContext(self, other_jid):
        assert isinstance(other_jid, jid.JID)
        context = self.contexts.setdefault(other_jid, Context(self.host, self.account, other_jid))
        return context

    def getContextForUser(self, other):
        log.debug(u"getContextForUser [%s]" % other)
        if not other.resource:
            log.warning(u"getContextForUser called with a bare jid: %s" % other.full())
        return self.startContext(other)


class OTR(object):

    def __init__(self, host):
        log.info(_(u"OTR plugin initialization"))
        self.host = host
        self.context_managers = {}
        self.skipped_profiles = set()  # FIXME: OTR should not be skipped per profile, this need to be refactored
        host.trigger.add("MessageReceived", self.MessageReceivedTrigger, priority=100000)
        host.trigger.add("messageSend", self.messageSendTrigger, priority=100000)
        host.bridge.addMethod("skipOTR", ".plugin", in_sign='s', out_sign='', method=self._skipOTR)  # FIXME: must be removed, must be done on per-message basis
        host.bridge.addSignal("otrState", ".plugin", signature='sss')  # args: state, destinee_jid, profile
        host.importMenu((OTR_MENU, D_(u"Start/Refresh")), self._otrStartRefresh, security_limit=0, help_string=D_(u"Start or refresh an OTR session"), type_=C.MENU_SINGLE)
        host.importMenu((OTR_MENU, D_(u"End session")), self._otrSessionEnd, security_limit=0, help_string=D_(u"Finish an OTR session"), type_=C.MENU_SINGLE)
        host.importMenu((OTR_MENU, D_(u"Authenticate")), self._otrAuthenticate, security_limit=0, help_string=D_(u"Authenticate user/see your fingerprint"), type_=C.MENU_SINGLE)
        host.importMenu((OTR_MENU, D_(u"Drop private key")), self._dropPrivKey, security_limit=0, type_=C.MENU_SINGLE)
        host.trigger.add("presenceReceived", self._presenceReceivedTrigger)

    def _skipOTR(self, profile):
        """Tell the backend to not handle OTR for this profile.

        @param profile (str): %(doc_profile)s
        """
        # FIXME: should not be done per profile but per message, using extra data
        #        for message received, profile wide hook may be need, but client
        #        should be used anyway instead of a class attribute
        self.skipped_profiles.add(profile)

    @defer.inlineCallbacks
    def profileConnected(self, profile):
        if profile in self.skipped_profiles:
            return
        client = self.host.getClient(profile)
        ctxMng = client._otr_context_manager = ContextManager(self.host, client)
        client._otr_data = persistent.PersistentBinaryDict(NS_OTR, profile)
        yield client._otr_data.load()
        encrypted_priv_key = client._otr_data.get(PRIVATE_KEY, None)
        if encrypted_priv_key is not None:
            priv_key = yield self.host.memory.decryptValue(encrypted_priv_key, profile)
            ctxMng.account.privkey = potr.crypt.PK.parsePrivateKey(priv_key.decode('hex'))[0]
        else:
            ctxMng.account.privkey = None
        ctxMng.account.loadTrusts()

    def profileDisconnected(self, profile):
        if profile in self.skipped_profiles:
            self.skipped_profiles.remove(profile)
            return
        client = self.host.getClient(profile)
        for context in client._otr_context_manager.contexts.values():
            context.disconnect()
        del client._otr_context_manager

    def _otrStartRefresh(self, menu_data, profile):
        """Start or refresh an OTR session

        @param menu_data: %(menu_data)s
        @param profile: %(doc_profile)s
        """
        client = self.host.getClient(profile)
        try:
            to_jid = jid.JID(menu_data['jid'])
        except KeyError:
            log.error(_(u"jid key is not present !"))
            return defer.fail(exceptions.DataError)
        self.startRefresh(client, to_jid)
        return {}

    def startRefresh(self, client, to_jid):
        """Start or refresh an OTR session

        @param to_jid(jid.JID): jid to start encrypted session with
        """
        if not to_jid.resource:
            to_jid.resource = self.host.memory.getMainResource(client, to_jid) # FIXME: temporary and unsecure, must be changed when frontends are refactored
        otrctx = client._otr_context_manager.getContextForUser(to_jid)
        query = otrctx.sendMessage(0, '?OTRv?')
        otrctx.inject(query)

    def _otrSessionEnd(self, menu_data, profile):
        """End an OTR session

        @param menu_data: %(menu_data)s
        @param profile: %(doc_profile)s
        """
        client = self.host.getClient(profile)
        try:
            to_jid = jid.JID(menu_data['jid'])
        except KeyError:
            log.error(_(u"jid key is not present !"))
            return defer.fail(exceptions.DataError)
        self.endSession(client, to_jid)
        return {}

    def endSession(self, client, to_jid):
        """End an OTR session"""
        if not to_jid.resource:
            to_jid.resource = self.host.memory.getMainResource(client, to_jid) # FIXME: temporary and unsecure, must be changed when frontends are refactored
        otrctx = client._otr_context_manager.getContextForUser(to_jid)
        otrctx.disconnect()
        return {}

    def _otrAuthenticate(self, menu_data, profile):
        """End an OTR session

        @param menu_data: %(menu_data)s
        @param profile: %(doc_profile)s
        """
        client = self.host.getClient(profile)
        try:
            to_jid = jid.JID(menu_data['jid'])
        except KeyError:
            log.error(_(u"jid key is not present !"))
            return defer.fail(exceptions.DataError)
        return self.authenticate(client, to_jid)

    def authenticate(self, client, to_jid):
        """Authenticate other user and see our own fingerprint"""
        if not to_jid.resource:
            to_jid.resource = self.host.memory.getMainResource(client, to_jid) # FIXME: temporary and unsecure, must be changed when frontends are refactored
        ctxMng = client._otr_context_manager
        otrctx = ctxMng.getContextForUser(to_jid)
        priv_key = ctxMng.account.privkey

        if priv_key is None:
            # we have no private key yet
            dialog = xml_tools.XMLUI(C.XMLUI_DIALOG,
                     dialog_opt = {C.XMLUI_DATA_TYPE: C.XMLUI_DIALOG_MESSAGE,
                                   C.XMLUI_DATA_MESS: _(u"You have no private key yet, start an OTR conversation to have one"),
                                   C.XMLUI_DATA_LVL: C.XMLUI_DATA_LVL_WARNING
                                  },
                     title = _(u"No private key"),
                     )
            return {'xmlui': dialog.toXml()}

        other_fingerprint = otrctx.getCurrentKey()

        if other_fingerprint is None:
            # we have a private key, but not the fingerprint of our correspondent
            dialog = xml_tools.XMLUI(C.XMLUI_DIALOG,
                     dialog_opt = {C.XMLUI_DATA_TYPE: C.XMLUI_DIALOG_MESSAGE,
                                   C.XMLUI_DATA_MESS: _(u"Your fingerprint is:\n{fingerprint}\n\nStart an OTR conversation to have your correspondent one.").format(fingerprint=priv_key),
                                   C.XMLUI_DATA_LVL: C.XMLUI_DATA_LVL_INFO
                                  },
                     title = _(u"Fingerprint"),
                     )
            return {'xmlui': dialog.toXml()}

        def setTrust(raw_data, profile):
            # This method is called when authentication form is submited
            data = xml_tools.XMLUIResult2DataFormResult(raw_data)
            if data['match'] == 'yes':
                otrctx.setCurrentTrust(OTR_STATE_TRUSTED)
                note_msg = _(u"Your correspondent {correspondent} is now TRUSTED")
                self.host.bridge.otrState(OTR_STATE_TRUSTED, to_jid.full(), client.profile)
            else:
                otrctx.setCurrentTrust('')
                note_msg = _(u"Your correspondent {correspondent} is now UNTRUSTED")
                self.host.bridge.otrState(OTR_STATE_UNTRUSTED, to_jid.full(), client.profile)
            note = xml_tools.XMLUI(C.XMLUI_DIALOG, dialog_opt = {
                                   C.XMLUI_DATA_TYPE: C.XMLUI_DIALOG_NOTE,
                                   C.XMLUI_DATA_MESS: note_msg.format(correspondent=otrctx.peer)}
                                   )
            return {'xmlui': note.toXml()}

        submit_id = self.host.registerCallback(setTrust, with_data=True, one_shot=True)
        trusted = bool(otrctx.getCurrentTrust())

        xmlui = xml_tools.XMLUI(C.XMLUI_FORM, title=_('Authentication (%s)') % to_jid.full(), submit_id=submit_id)
        xmlui.addText(_(AUTH_TXT))
        xmlui.addDivider()
        xmlui.addText(D_(u"Your own fingerprint is:\n{fingerprint}").format(fingerprint=priv_key))
        xmlui.addText(D_(u"Your correspondent fingerprint should be:\n{fingerprint}").format(fingerprint=other_fingerprint))
        xmlui.addDivider('blank')
        xmlui.changeContainer('pairs')
        xmlui.addLabel(D_(u'Is your correspondent fingerprint the same as here ?'))
        xmlui.addList("match", [('yes', _('yes')),('no', _('no'))], ['yes' if trusted else 'no'])
        return {'xmlui': xmlui.toXml()}

    def _dropPrivKey(self, menu_data, profile):
        """Drop our private Key

        @param menu_data: %(menu_data)s
        @param profile: %(doc_profile)s
        """
        client = self.host.getClient(profile)
        try:
            to_jid = jid.JID(menu_data['jid'])
            if not to_jid.resource:
                to_jid.resource = self.host.memory.getMainResource(client, to_jid) # FIXME: temporary and unsecure, must be changed when frontends are refactored
        except KeyError:
            log.error(_(u"jid key is not present !"))
            return defer.fail(exceptions.DataError)

        ctxMng = client._otr_context_manager
        if ctxMng.account.privkey is None:
            return {'xmlui': xml_tools.note(_(u"You don't have a private key yet !")).toXml()}

        def dropKey(data, profile):
            if C.bool(data['answer']):
                # we end all sessions
                for context in ctxMng.contexts.values():
                    context.disconnect()
                ctxMng.account.privkey = None
                ctxMng.account.getPrivkey()  # as account.privkey is None, getPrivkey will generate a new key, and save it
                return {'xmlui': xml_tools.note(D_(u"Your private key has been dropped")).toXml()}
            return {}

        submit_id = self.host.registerCallback(dropKey, with_data=True, one_shot=True)

        confirm = xml_tools.XMLUI(C.XMLUI_DIALOG, title=_(u'Confirm private key drop'), dialog_opt = {'type': C.XMLUI_DIALOG_CONFIRM, 'message': _(DROP_TXT)}, submit_id = submit_id)
        return {'xmlui': confirm.toXml()}

    def _receivedTreatment(self, data, client):
        from_jid = data['from']
        log.debug(u"_receivedTreatment [from_jid = %s]" % from_jid)
        otrctx = client._otr_context_manager.getContextForUser(from_jid)

        try:
            message = data['message'].itervalues().next() # FIXME: Q&D fix for message refactoring, message is now a dict
            res = otrctx.receiveMessage(message.encode('utf-8'))
        except potr.context.UnencryptedMessage:
            encrypted = False
            if otrctx.state == potr.context.STATE_ENCRYPTED:
                log.warning(u"Received unencrypted message in an encrypted context (from {jid})".format(
                    jid = from_jid.full()))

                feedback=D_(u"WARNING: received unencrypted data in a supposedly encrypted context"),
                self.host.bridge.messageNew(uid=unicode(uuid.uuid4()),
                                            timestamp=time.time(),
                                            from_jid=from_jid.full(),
                                            to_jid=client.jid.full(),
                                            message={u'': feedback},
                                            subject={},
                                            mess_type=C.MESS_TYPE_INFO,
                                            extra={},
                                            profile=client.profile)
        except StopIteration:
            return data
        else:
            encrypted = True

        if encrypted:
            if res[0] != None:
                # decrypted messages handling.
                # receiveMessage() will return a tuple, the first part of which will be the decrypted message
                data['message'] = {'':res[0].decode('utf-8')} # FIXME: Q&D fix for message refactoring, message is now a dict
                # TODO: add skip history as an option, but by default we don't skip it
                # raise failure.Failure(exceptions.SkipHistory()) # we send the decrypted message to frontends, but we don't want it in history
            else:
                log.warning(u"An encrypted message was expected, but got {}".format(data['message']))
                raise failure.Failure(exceptions.CancelError('Cancelled by OTR')) # no message at all (no history, no signal)
        return data

    def _receivedTreatmentForSkippedProfiles(self, data):
        """This profile must be skipped because the frontend manages OTR itself,

        but we still need to check if the message must be stored in history or not
        """
        # XXX: FIXME: this should not be done on a per-profile basis, but  per-message
        try:
            message = data['message'].itervalues().next().encode('utf-8') # FIXME: Q&D fix for message refactoring, message is now a dict
        except StopIteration:
            return data
        if message.startswith(potr.proto.OTRTAG):
            raise failure.Failure(exceptions.SkipHistory())
        return data

    def MessageReceivedTrigger(self, client, message_elt, post_treat):
        if message_elt.getAttribute('type') == C.MESS_TYPE_GROUPCHAT:
            # OTR is not possible in group chats
            return True
        if client.profile in self.skipped_profiles:
            post_treat.addCallback(self._receivedTreatmentForSkippedProfiles)
        else:
            post_treat.addCallback(self._receivedTreatment, client)
        return True

    def _messageSendOTR(self, mess_data, client):
        to_jid = copy.copy(mess_data['to'])
        if not to_jid.resource:
            to_jid.resource = self.host.memory.getMainResource(client, to_jid) # FIXME: full jid may not be known
        otrctx = client._otr_context_manager.getContextForUser(to_jid)
        if otrctx.state != potr.context.STATE_PLAINTEXT:
            if otrctx.state == potr.context.STATE_ENCRYPTED:
                log.debug(u"encrypting message")
                try:
                    msg = mess_data['message']['']
                except KeyError:
                    try:
                        msg = mess_data['message'].itervalues().next()
                    except StopIteration:
                        log.warning(u"No message found")
                        for key, value in mess_data['extra'].iteritems():
                            if key.startswith('rich') or key.startswith('xhtml'):
                                log.error(u'received rich content while OTR encryption is activated, cancelling')
                                raise failure.Failure(exceptions.CancelError())
                        return mess_data
                otrctx.sendMessage(0, msg.encode('utf-8'))
                self.host.messageAddToHistory(mess_data, client)
                self.host.messageSendToBridge(mess_data, client)
            else:
                feedback = D_(u"Your message was not sent because your correspondent closed the encrypted conversation on his/her side. Either close your own side, or refresh the session.")
                self.host.bridge.messageNew(uid=unicode(uuid.uuid4()),
                                            timestamp=time.time(),
                                            from_jid=to_jid.full(),
                                            to_jid=client.jid.full(),
                                            message={u'': feedback},
                                            subject={},
                                            mess_type=C.MESS_TYPE_INFO,
                                            extra={},
                                            profile=client.profile)
            # we stop treatment here for OTR, as message is already sent
            raise failure.Failure(exceptions.CancelError())
        else:
            return mess_data

    def messageSendTrigger(self, client, mess_data, pre_xml_treatments, post_xml_treatments):
        if mess_data['type'] == 'groupchat':
            return True
        if client.profile in self.skipped_profiles:  # FIXME: should not be done on a per-profile basis
            return True
        pre_xml_treatments.addCallback(self._messageSendOTR, client)
        return True

    def _presenceReceivedTrigger(self, entity, show, priority, statuses, profile):
        if show != C.PRESENCE_UNAVAILABLE:
            return True
        client = self.host.getClient(profile)
        if not entity.resource:
            try:
                entity.resource = self.host.memory.getMainResource(client, entity)  # FIXME: temporary and unsecure, must be changed when frontends are refactored
            except exceptions.UnknownEntityError:
                return True  #  entity was not connected
        if entity in client._otr_context_manager.contexts:
            otrctx = client._otr_context_manager.getContextForUser(entity)
            otrctx.disconnect()
        return True