# HG changeset patch
# User Goffi <goffi@goffi.org>
# Date 1582046238 -3600
# Node ID c90f27ce52b0025324ba59d6bf9adec032216f13
# Parent  343b8076e967e4dba631b768d433696093a7bcf6
plugin aesgcm: look for "aesgcm" links in body to use them as attachments

diff -r 343b8076e967 -r c90f27ce52b0 sat/plugins/plugin_sec_aesgcm.py
--- a/sat/plugins/plugin_sec_aesgcm.py	Tue Feb 18 18:17:14 2020 +0100
+++ b/sat/plugins/plugin_sec_aesgcm.py	Tue Feb 18 18:17:18 2020 +0100
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# SAT plugin for handling AES-GCM file encryption
+# SàT plugin for handling AES-GCM file encryption
 # Copyright (C) 2009-2020 Jérôme Poisson (goffi@goffi.org)
 
 # This program is free software: you can redistribute it and/or modify
@@ -16,8 +16,11 @@
 # 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/>.
 
+import re
 from textwrap import dedent
 from functools import partial
+from urllib.parse import urlparse
+import mimetypes
 import secrets
 from cryptography.hazmat.primitives import ciphers
 from cryptography.hazmat.primitives.ciphers import modes
@@ -47,6 +50,10 @@
     """)),
 }
 
+AESGCM_RE = re.compile(
+    r'aesgcm:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9'
+    r'()@:%_\+.~#?&\/\/=]*)')
+
 
 class AESGCM(object):
 
@@ -58,6 +65,7 @@
         )
         host.trigger.add("XEP-0363_upload_size", self._uploadSizeTrigger)
         host.trigger.add("XEP-0363_upload", self._uploadTrigger)
+        host.trigger.add("messageReceived", self._messageReceivedTrigger)
 
     async def download(self, client, uri_parsed, dest_path, options):
         fragment = bytes.fromhex(uri_parsed.fragment)
@@ -198,3 +206,53 @@
         # with data_cb we encrypt the file on the fly
         sat_file.data_cb = partial(self._encrypt, encryptor=encryptor)
         return True
+
+
+    def _popAESGCMLinks(self, match, links):
+        link = match.group()
+        if link not in links:
+            links.append(link)
+        return ""
+
+    def _checkAESGCMAttachments(self, client, data):
+        if not data.get('message'):
+            return data
+        links = []
+
+        for lang, message in list(data['message'].items()):
+            message = AESGCM_RE.sub(
+                partial(self._popAESGCMLinks, links=links),
+                message)
+            if links:
+                message = message.strip()
+                if not message:
+                    del data['message'][lang]
+                else:
+                    data['message'][lang] = message
+                mess_encrypted = client.encryption.isEncrypted(data)
+                attachments = data['extra'].setdefault(C.MESS_KEY_ATTACHMENTS, [])
+                for link in links:
+                    path = urlparse(link).path
+                    attachment = {
+                        "url": link,
+                    }
+                    media_type = mimetypes.guess_type(path, strict=False)[0]
+                    if media_type is not None:
+                        attachment[C.MESS_KEY_MEDIA_TYPE] = media_type
+
+                    if mess_encrypted:
+                        # we don't add the encrypted flag if the message itself is not
+                        # encrypted, because the decryption key is part of the link,
+                        # so sending it over unencrypted channel is like having no
+                        # encryption at all.
+                        attachment['encrypted'] = True
+                    attachments.append(attachment)
+
+        return data
+
+    def _messageReceivedTrigger(self, client, message_elt, post_treat):
+        # we use a post_treat callback instead of "message_parse" trigger because we need
+        # to check if the "encrypted" flag is set to decide if we add the same flag to the
+        # attachment
+        post_treat.addCallback(partial(self._checkAESGCMAttachments, client))
+        return True