diff sat/plugins/plugin_misc_attach.py @ 3223:163014f09bf4

plugin attach: handle large images resizing: if `C.MESS_KEY_ATTACHMENTS_RESIZE` flag is set for an attachment, it will be checked, and if it is a large image, a resized copy will be uploaded.
author Goffi <goffi@goffi.org>
date Sun, 22 Mar 2020 14:04:16 +0100
parents 2ba602aef90e
children be6d91572633
line wrap: on
line diff
--- a/sat/plugins/plugin_misc_attach.py	Sun Mar 22 14:01:47 2020 +0100
+++ b/sat/plugins/plugin_misc_attach.py	Sun Mar 22 14:04:16 2020 +0100
@@ -16,15 +16,18 @@
 # 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 functools import partial
 from pathlib import Path
 from collections import namedtuple
 from twisted.internet import defer
+import mimetypes
+import tempfile
+import shutil
 from sat.core.i18n import _
 from sat.core import exceptions
 from sat.core.constants import Const as C
 from sat.core.log import getLogger
 from sat.tools import utils
+from sat.tools import image
 
 
 log = getLogger(__name__)
@@ -86,6 +89,46 @@
         log.debug(f"new attachments handler: {handler}")
 
     async def attachFiles(self, client, data):
+        """Main method to attach file
+
+        It will do generic pre-treatment, and call the suitable attachments handler
+        """
+        # we check attachment for pre-treatment like large image resizing
+        # media_type will be added if missing (and if it can be guessed from path)
+        attachments = data["extra"][C.MESS_KEY_ATTACHMENTS]
+        tmp_dirs_to_clean = []
+        for attachment in attachments:
+            if attachment.get(C.MESS_KEY_ATTACHMENTS_RESIZE, False):
+                path = Path(attachment["path"])
+                try:
+                    media_type = attachment[C.MESS_KEY_ATTACHMENTS_MEDIA_TYPE]
+                except KeyError:
+                    media_type = mimetypes.guess_type(path, strict=False)[0]
+                    if media_type is None:
+                        log.warning(
+                            _("Can't resize attachment of unknown type: {attachment}")
+                            .format(attachment=attachment))
+                        continue
+                    attachment[C.MESS_KEY_ATTACHMENTS_MEDIA_TYPE] = media_type
+
+                main_type = media_type.split('/')[0]
+                if main_type == "image":
+                    report = image.check(self.host, path)
+                    if report['too_large']:
+                        tmp_dir = Path(tempfile.mkdtemp())
+                        tmp_dirs_to_clean.append(tmp_dir)
+                        new_path = tmp_dir / path.name
+                        await image.resize(
+                            path, report["recommended_size"], dest=new_path)
+                        attachment["path"] = new_path
+                        log.info(
+                            _("Attachment {path!r} has been resized at {new_path!r}")
+                            .format(path=str(path), new_path=str(new_path)))
+                else:
+                    log.warning(
+                        _("Can't resize attachment of type {main_type!r}: {attachment}")
+                        .format(main_type=main_type, attachment=attachment))
+
         if client.encryption.isEncryptionRequested(data):
             handlers = self._attachments_handlers['encrypted']
         else:
@@ -103,6 +146,10 @@
 
         await utils.asDeferred(handler.attach, client, data)
 
+        for dir_path in tmp_dirs_to_clean:
+            log.debug(f"Cleaning temporary directory at {dir_path}")
+            shutil.rmtree(dir_path)
+
         return data
 
     async def uploadFiles(self, client, data, upload_cb=None):
@@ -186,13 +233,13 @@
 
         return data
 
-    def _attachFiles(self, client, data):
+    def _attachFiles(self, data, client):
         return defer.ensureDeferred(self.attachFiles(client, data))
 
     def _sendMessageTrigger(
         self, client, mess_data, pre_xml_treatments, post_xml_treatments):
         if mess_data['extra'].get(C.MESS_KEY_ATTACHMENTS):
-            post_xml_treatments.addCallback(partial(self._attachFiles, client))
+            post_xml_treatments.addCallback(self._attachFiles, client=client)
         return True
 
     async def defaultCanHandle(self, client, data):