e2ee api: Ensure only one encryption method is enabled at a time

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>

1 files changed, 15 insertions, 9 deletions

diff --git a/poezio/plugin_e2ee.py b/poezio/plugin_e2ee.py
index 0a1080f2..e4497113 100644
--- a/poezio/plugin_e2ee.py
+++ b/poezio/plugin_e2ee.py
@@ -10,7 +10,7 @@
     Interface for E2EE (End-to-end Encryption) plugins.
 """
 
-from typing import Optional, Union
+from typing import Callable, Dict, Optional, Union
 
 from slixmpp import InvalidJID, JID, Message
 from poezio.tabs import ConversationTab, DynamicConversationTab, PrivateTab, MucTab
@@ -40,6 +40,10 @@ class E2EEPlugin(BasePlugin):
     # Required.
     eme_ns = None  # type: Optional[str]
 
+    # Static map, to be able to limit to one encryption mechanism per tab at a
+    # time
+    _enabled_tabs = {}  # type: Dict[JID, Callable]
+
     def init(self):
         if self.encryption_name is None and self.encryption_short_name is None:
             raise NotImplementedError
@@ -59,8 +63,6 @@ class E2EEPlugin(BasePlugin):
         self.api.add_event_handler('private_msg', self._decrypt)
         self.api.add_event_handler('private_say', self._encrypt)
 
-        self._enabled_tabs = set()
-
         for tab_t in (DynamicConversationTab, PrivateTab, MucTab):
             self.api.add_tab_command(
                 tab_t,
@@ -100,26 +102,29 @@ class E2EEPlugin(BasePlugin):
         except InvalidJID:
             return ""
 
-        if jid in self._enabled_tabs:
+        if self._encryption_enabled(jid):
             return " " + self.encryption_short_name
         return ""
 
     def _toggle_tab(self, _input: str) -> None:
         jid = self.api.current_tab().jid  # type: JID
 
-        if jid in self._enabled_tabs:
-            self._enabled_tabs.remove(jid)
+        if self._encryption_enabled(jid):
+            del self._enabled_tabs[jid]
             self.api.information(
                 '{} encryption disabled for {}'.format(self.encryption_name, jid),
                 'Info',
             )
         else:
-            self._enabled_tabs.add(jid)
+            self._enabled_tabs[jid] = self.encrypt
             self.api.information(
                 '{} encryption enabled for {}'.format(self.encryption_name, jid),
                 'Info',
             )
 
+    def _encryption_enabled(self, jid: JID) -> bool:
+        return jid in self._enabled_tabs and self._enabled_tabs[jid] == self.encrypt
+
     def _decrypt(self, message: Message, tab: ChatTabs) -> None:
         if message.xml.find('{%s}%s' % (EME_NS, EME_TAG)) is None:
             return None
@@ -136,7 +141,7 @@ class E2EEPlugin(BasePlugin):
 
     def _encrypt(self, message: Message, tab: ChatTabs):
         jid = tab.jid
-        if jid not in self._enabled_tabs:
+        if not self._encryption_enabled(jid):
             return None
 
         log.debug('Sending %s message: %r', self.encryption_name, message['body'])
@@ -144,7 +149,8 @@ class E2EEPlugin(BasePlugin):
         message['eme']['namespace'] = self.eme_ns
         message['eme']['name'] = self.encryption_name
 
-        self.encrypt(message, tab)
+        # Call the enabled encrypt method
+        self._enabled_tabs[jid](message, tab)
 
         log.debug('Decrypted %s message: %r', self.encryption_name, message['body'])
         return None