aboutsummarybugs & patchesrefslogtreecommitdiffstats
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)