From: Kylie Czar <czar@kalli.st>
Date: Thu, 18 Aug 2022 00:01:11 +0000 (-0300)
Subject: feat: make uploaded file naming better
X-Git-Url: https://git.patchysicp.xyz/?a=commitdiff_plain;h=1351a944d091c91ce8e844da7510f9ec385b186f;p=angel.git

feat: make uploaded file naming better
---

diff --git a/main.py b/main.py
index c9549a2..334cb2b 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@ import random
 import configparser
 import re
 import io
+import os
 import mimetypes
 import asyncio
 from collections import defaultdict
@@ -12,7 +13,7 @@ from PythonSed import Sed
 from slixmpp import ClientXMPP
 from urllib.parse import urlparse, parse_qs, urlunparse
 from pantomime import normalize_mimetype
-
+import cgi
 
 sed_parse = re.compile("(?<!\\\\)[/#]")
 sed_cmd = re.compile("^s[/#].*[/#].*[/#]")
@@ -123,7 +124,6 @@ class AngelBot(ClientXMPP):
         r = requests.get(url, stream=True, headers=headers, timeout=5)
         if not r.ok:
             return
-
         ftype = normalize_mimetype(r.headers.get("content-type"))
         if ftype in html_files:
             data = ""
@@ -157,15 +157,20 @@ class AngelBot(ClientXMPP):
                         return
                     outfile.write(chunk)
 
-                await self.embed_file(url, sender, mtype, ftype, outfile)
+                content_disposition = r.headers.get("content-disposition")
+
+                _, params = cgi.parse_header(content_disposition)
+
+                filename = params.get("filename")
+                ext = os.path.splitext(filename)[1] if filename else None or ".bin"
+                fname = filename or uri.path.strip("/").split("/")[-1] or f"file{ext}"
+                await self.embed_file(url, sender, mtype, ftype, fname, outfile)
             except Exception:
                 ...
 
-    async def embed_file(self, url, sender, mtype, ftype, outfile):
-        ext = mimetypes.guess_extension(ftype)
-        filename = f"file{ext}"
+    async def embed_file(self, url, sender, mtype, ftype, fname, outfile):
         furl = await self.plugin["xep_0363"].upload_file(
-            filename, content_type=ftype, input_file=outfile
+            fname, content_type=ftype, input_file=outfile
         )
         message = self.make_message(sender)
         message["body"] = furl
@@ -227,7 +232,6 @@ class AngelBot(ClientXMPP):
         self.register_plugin("xep_0084")
         self.register_plugin("xep_0153")
         self.register_plugin("xep_0363")
-        self.register_plugin("xep_0363")
 
         self.add_event_handler("session_start", self.session_start)
         self.add_event_handler("message", self.message)
@@ -261,13 +265,13 @@ class AngelBot(ClientXMPP):
         }
 
         vcard = self.plugin["xep_0054"].make_vcard()
-        
+
         vcard["URL"] = "https://gt.kalli.st/czar/angel"
-        
+
         asyncio.gather(self.plugin["xep_0054"].publish_vcard(vcard))
 
         asyncio.gather(self.plugin["xep_0084"].publish_avatar(avatar))
-        
+
         asyncio.gather(
             self.plugin["xep_0153"].set_avatar(
                 avatar=avatar,