From 06e4e480c14a614cbf00646271fbc9abc6908fbb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Sat, 19 Mar 2022 17:06:12 +0100 Subject: xep_0454: keep original filename extension if available MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- slixmpp/plugins/xep_0454/__init__.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/slixmpp/plugins/xep_0454/__init__.py b/slixmpp/plugins/xep_0454/__init__.py index cecde480..d9292a13 100644 --- a/slixmpp/plugins/xep_0454/__init__.py +++ b/slixmpp/plugins/xep_0454/__init__.py @@ -26,6 +26,11 @@ class InvalidURL(Exception): """Raised for URLs that either aren't HTTPS or already contain a fragment.""" +EXTENSIONS_MAP = { + 'jpeg': 'jpg', + 'text': 'txt', +} + class XEP_0454(BasePlugin): """ XEP-0454: OMEMO Media Sharing @@ -118,6 +123,14 @@ class XEP_0454(BasePlugin): raise InvalidURL return 'aesgcm://' + url.removeprefix('https://') + '#' + fragment + @staticmethod + def map_extensions(ext: str) -> str: + """ + Apply conversions to extensions to reduce the number of + variations, (e.g., JPEG -> jpg). + """ + return EXTENSIONS_MAP.get(ext, ext).lower() + async def upload_file( self, filename: Path, @@ -142,8 +155,10 @@ class XEP_0454(BasePlugin): payload, fragment = self.encrypt(input_file, filename) # Prepare kwargs for upload_file call - filename = urandom(12).hex() # Random filename to hide user-provided path - kwargs['filename'] = filename + new_filename = urandom(12).hex() # Random filename to hide user-provided path + if filename.suffix: + new_filename += self.map_extensions(filename.suffix) + kwargs['filename'] = new_filename input_enc = BytesIO(payload) kwargs['input_file'] = input_enc -- cgit v1.2.3