caching and hashing done

Signed-off-by: Vosjedev <vosje@vosjedev.net>

1 files changed, 74 insertions, 0 deletions

diff --git a/src/discord_image_bridge/utils.py b/src/discord_image_bridge/utils.py
new file mode 100644
index 0000000..923e17f
--- /dev/null
+++ b/src/discord_image_bridge/utils.py
@@ -0,0 +1,74 @@
+import os
+import cherrypy
+import requests
+from hashlib import md5 as do_hash
+from .dbpool import DBPoolManager
+
+from . import discord
+
+
+def download_and_cache(url, filename):
+	resp=requests.get(url)
+	if resp.status_code==200:
+		hash=do_hash(resp.content).hexdigest()
+	try:
+		os.mkdir(hash[:2])
+		fname=hash[:2]+hash
+		with open(fname,'wb') as fd:
+			fd.write(resp.content)
+		return hash, fname
+	except OSError as e:
+		cherrypy.log("Error writing "+filename+" to disk: "+repr(e))
+		return None, None
+
+def on_ready(plugin:discord.DiscordWsManager, client:discord.DiscordWsClient):
+	dbpool:DBPoolManager=plugin.dbpool
+
+	with dbpool.get_connection() as conn, conn.cursor() as cur:
+		cur.execute(
+			"CREATE TABLE IF NOT EXISTS attachments ("
+			"seq INT NOT NULL,"
+			"filename VARCHAR(255) NOT NULL,"
+			"channel BIGINT UNSIGNED,"
+			"message BIGINT UNSIGNED,"
+			"id BIGINT UNSIGNED,"
+			"hash CHAR(32)"
+			")"
+		)
+		conn.commit()
+
+	@client.event()
+	def on_message_create(payload, data):
+		msgid=data["id"]
+		channel=data["channel_id"]
+		with dbpool.get_connection() as conn, conn.cursor() as cur:
+			for attachment in data["attachments"]:
+				id=attachment["id"]
+				fname=attachment["filename"]
+
+				cur.execute("SELECT id FROM attachments WHERE id=? AND filename=?",(id,fname))
+				if cur.fetchone():
+					cherrypy.log("Attachment with id %s and fname `%s' already downloaded"%(id,fname))
+					continue
+
+				cherrypy.log("Downloading attachment %s/%s %s"%(channel,id,fname))
+
+				hash, disk_fname=download_and_cache(attachment["url"], fname)
+
+				cur.execute("SELECT seq FROM attachments WHERE filename=? ORDER BY seq DESC LIMIT 1",(fname,))
+				nr=cur.fetchone()
+				if nr:
+					nr=nr[0]+1
+				else:
+					nr=0
+				cherrypy.log("Caching attachment %s/%s %s under %d"%(channel,id,fname,nr))
+
+				cur.execute("INSERT INTO attachments (seq, filename, channel, message, id, hash) VALUES (?,?,?,?,?,?)",(
+					nr,fname[:255],
+					channel,msgid,id,
+					hash
+					))
+
+			conn.commit()
+			
+