feat: make uploaded file naming better

This commit is contained in:
Czar 2022-08-17 21:01:11 -03:00
commit 1351a944d0
Signed by: czar
GPG key ID: 2DBA9558E9277C37

26
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)
@ -261,13 +265,13 @@ class AngelBot(ClientXMPP):
} }
vcard = self.plugin["xep_0054"].make_vcard() vcard = self.plugin["xep_0054"].make_vcard()
vcard["URL"] = "https://gt.kalli.st/czar/angel" vcard["URL"] = "https://gt.kalli.st/czar/angel"
asyncio.gather(self.plugin["xep_0054"].publish_vcard(vcard)) asyncio.gather(self.plugin["xep_0054"].publish_vcard(vcard))
asyncio.gather(self.plugin["xep_0084"].publish_avatar(avatar)) asyncio.gather(self.plugin["xep_0084"].publish_avatar(avatar))
asyncio.gather( asyncio.gather(
self.plugin["xep_0153"].set_avatar( self.plugin["xep_0153"].set_avatar(
avatar=avatar, avatar=avatar,