import os
from time import time
from threading import RLock as Lock
import cherrypy
import requests
from hashlib import sha256 as do_hash
from . import discord
from .downloadpool import DownloadPool
from . import fsmanager
download_locks={}
def safe_mkdir(dirname):
try:
return os.mkdir(dirname)
except FileExistsError:
pass
def download_and_cache(url, filename):
cherrypy.log("Downloading attachment %s"%filename)
try:
resp=requests.get(url)
except requests.exceptions.RequestException as e:
cherrypy.log("Error downloading: %s"%repr(e))
return None, None
if not resp.status_code==200:
cherrypy.log("Attachment %s failed to download"%filename)
return None, None
hash=do_hash(resp.content).hexdigest()
try:
tmpname=fsmanager.hash2fname(hash,temp=True)
dirname=os.path.dirname(tmpname)
safe_mkdir(dirname)
with fsmanager.DataFile(open(tmpname,'wb')) as fd:
fd.write(resp.content)
fname=fsmanager.hash2fname(hash)
dirname=os.path.dirname(fname)
safe_mkdir(dirname)
os.rename(tmpname, fname)
return hash, fname
except FileExistsError:
pass
except OSError as e:
cherrypy.log("Error writing "+filename+" to disk: "+repr(e))
return None, None
return None, None
def download_uncached(hash):
if not os.path.isfile(fsmanager.hash2meta_fname(hash)):
return None, None
with download_locks.setdefault(hash, Lock()), fsmanager.MetaFile(hash) as metafd:
target=fsmanager.hash2fname(hash)
if os.path.isfile(target):
return hash, target
data=metafd.read()
if not "sources" in data:
cherrypy.log("No sources available for "+hash)
return None, None
for source in data["sources"] + data["sources"]: # try everything twice
match source["type"]:
case "discord":
newhash, fname=download_uncached_discord(
channel=source["channel"],
msgid=source["message"],
attachmentid=source["attachment"]
)
case _:
# NOTE: maybe log here?
return None, None
if newhash.startswith(hash):
return hash, fname
cherrypy.log("No working sources available for "+hash)
return None, None
def download_uncached_discord(channel, msgid, attachmentid):
status,data=discord.channel_message_get(channel_id=channel, message_id=msgid)
for attachment in data["attachments"]:
if attachment["id"]==str(attachmentid):
return download_and_cache(attachment["url"],attachment["filename"])
return (None, None)
def clear_cache():
ttl=int(os.getenv("BRIDGE_CACHE_TIME",3600*24*7))
cherrypy.log("Clearing cache")
c=0
before=int(time())-ttl
# TODO: Rewrite
#with _values.dbpool.get_connection() as conn, conn.cursor() as cur:
# cur.execute("SELECT hash FROM cache WHERE fetched<?",(before,))
# for (hash,) in cur:
# print(hash)
# path=os.path.join(hash[:2],hash)
# if os.path.exists(path):
# c+=1
# os.remove(path)
# cur.execute("DELETE FROM cache WHERE fetched<?",(before,))
# conn.commit()
#cherrypy.log("%d files removed"%c)
def on_ready(plugin:discord.DiscordWsManager, client:discord.DiscordWsClient):
download_pool=DownloadPool(download_and_cache)
cherrypy.engine.subscribe("stop",download_pool.stop)
@client.event()
def on_message_create(payload, data):
msgid=data["id"]
channel=data["channel_id"]
for attachment in data["attachments"]:
id=attachment["id"]
fname=attachment["filename"]
def callback(hash, disk_fname):
cherrypy.log("Writing meta for hash "+hash)
with fsmanager.MetaFile(hash) as metafd:
data=metafd.read()
if not "sources" in data:
data["sources"]=[]
data["sources"].append({
"type":"discord",
"message":msgid,
"channel":channel,
"attachment":id
})
metafd.write(data)
#cherrypy.log("Done")
download_pool.exec(args=(attachment["url"], fname),callback=callback)