czar/angel
1
1
Fork 0
Code Issues 2 Pull requests Projects 1 Releases Wiki Activity

feat: make uploaded file naming better

Browse source
This commit is contained in:
Czar 2022-08-17 21:01:11 -03:00
commit 1351a944d0
Signed by: czar
GPG key ID: 2DBA9558E9277C37
1 changed files with 15 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

20
main.py
View file

@ -5,6 +5,7 @@ import random
import configparser import configparser
import re import re
import io import io
import os
import mimetypes import mimetypes
import asyncio import asyncio
from collections import defaultdict from collections import defaultdict
@ -12,7 +13,7 @@ from PythonSed import Sed
from slixmpp import ClientXMPP from slixmpp import ClientXMPP
from urllib.parse import urlparse, parse_qs, urlunparse from urllib.parse import urlparse, parse_qs, urlunparse
from pantomime import normalize_mimetype from pantomime import normalize_mimetype
import cgi
sed_parse = re.compile("(?<!\\\\)[/#]") sed_parse = re.compile("(?<!\\\\)[/#]")
sed_cmd = re.compile("^s[/#].*[/#].*[/#]") sed_cmd = re.compile("^s[/#].*[/#].*[/#]")
@ -123,7 +124,6 @@ class AngelBot(ClientXMPP):
r = requests.get(url, stream=True, headers=headers, timeout=5) r = requests.get(url, stream=True, headers=headers, timeout=5)
if not r.ok: if not r.ok:
return return
ftype = normalize_mimetype(r.headers.get("content-type")) ftype = normalize_mimetype(r.headers.get("content-type"))
if ftype in html_files: if ftype in html_files:
data = "" data = ""
@ -157,15 +157,20 @@ class AngelBot(ClientXMPP):
return return
outfile.write(chunk) 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: except Exception:
... ...
async def embed_file(self, url, sender, mtype, ftype, outfile): async def embed_file(self, url, sender, mtype, ftype, fname, outfile):
ext = mimetypes.guess_extension(ftype)
filename = f"file{ext}"
furl = await self.plugin["xep_0363"].upload_file( 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 = self.make_message(sender)
message["body"] = furl message["body"] = furl
@ -227,7 +232,6 @@ class AngelBot(ClientXMPP):
self.register_plugin("xep_0084") self.register_plugin("xep_0084")
self.register_plugin("xep_0153") self.register_plugin("xep_0153")
self.register_plugin("xep_0363") self.register_plugin("xep_0363")
self.register_plugin("xep_0363")
self.add_event_handler("session_start", self.session_start) self.add_event_handler("session_start", self.session_start)
self.add_event_handler("message", self.message) self.add_event_handler("message", self.message)