From 15fa2f3231da79f61cd98600340eac72e18a3f11 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sat, 23 Mar 2019 21:19:30 -0300
Subject: [PATCH] channel.messages: add foreign key deletions before message
 delete

 - images: export try_unlink, and allow it to handle pathlib.Path
---
 litecord/blueprints/channel/messages.py | 35 +++++++++++++++++++++++++
 litecord/images.py                      | 15 +++++++----
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/litecord/blueprints/channel/messages.py b/litecord/blueprints/channel/messages.py
index 2e15825..faca969 100644
--- a/litecord/blueprints/channel/messages.py
+++ b/litecord/blueprints/channel/messages.py
@@ -18,6 +18,7 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 
 import json
+from pathlib import Path
 
 from PIL import Image
 from quart import Blueprint, request, current_app as app, jsonify
@@ -35,6 +36,7 @@ from litecord.utils import pg_set_json
 from litecord.embed.sanitizer import fill_embed
 from litecord.embed.messages import process_url_embed
 from litecord.blueprints.channel.dm_checks import dm_pre_check
+from litecord.images import try_unlink
 
 
 log = Logger(__name__)
@@ -515,6 +517,37 @@ async def edit_message(channel_id, message_id):
     return jsonify(message)
 
 
+async def _del_msg_fkeys(message_id: int):
+    attachs = await app.db.fetch("""
+    SELECT id FROM attachments
+    WHERE message_id = $1
+    """, message_id)
+
+    attachs = [r['id'] for r in attachs]
+
+    attachments = Path('./attachments')
+    for attach_id in attachs:
+        # anything starting with the given attachment shall be
+        # deleted, because there may be resizes of the original
+        # attachment laying around.
+        for filepath in attachments.glob(f'{attach_id}*'):
+            try_unlink(filepath)
+
+    # after trying to delete all available attachments, delete
+    # them from the database.
+
+    # take the chance and delete all the data from the other tables too!
+
+    tables = ['attachments', 'message_webhook_info',
+              'message_reactions', 'channel_pins']
+
+    for table in tables:
+        await app.db.execute(f"""
+        DELETE FROM {table}
+        WHERE message_id = $1
+        """, message_id)
+
+
 @bp.route('/<int:channel_id>/messages/<int:message_id>', methods=['DELETE'])
 async def delete_message(channel_id, message_id):
     user_id = await token_check()
@@ -535,6 +568,8 @@ async def delete_message(channel_id, message_id):
     if not can_delete:
         raise Forbidden('You can not delete this message')
 
+    await _del_msg_fkeys(message_id)
+
     await app.db.execute("""
     DELETE FROM messages
     WHERE messages.id = $1
diff --git a/litecord/images.py b/litecord/images.py
index 8db1e2c..ec37469 100644
--- a/litecord/images.py
+++ b/litecord/images.py
@@ -22,7 +22,7 @@ import mimetypes
 import asyncio
 import base64
 import tempfile
-from typing import Optional
+from typing import Optional, Union
 
 from dataclasses import dataclass
 from hashlib import sha256
@@ -205,9 +205,14 @@ def _invalid(kwargs: dict) -> Optional[Icon]:
     return Icon(None, None, '')
 
 
-def _try_unlink(path: str):
+def try_unlink(path: Union[Path, str]):
+    """Try unlinking a file. Does not do anything if the file
+    does not exist."""
     try:
-        os.remove(path)
+        if isinstance(path, Path):
+            path.unlink()
+        else:
+            os.remove(path)
     except FileNotFoundError:
         pass
 
@@ -258,8 +263,8 @@ async def resize_gif(raw_data: bytes, target: tuple) -> tuple:
     output_handler.close()
 
     # delete the files we created with mkstemp
-    _try_unlink(input_path)
-    _try_unlink(output_path)
+    try_unlink(input_path)
+    try_unlink(output_path)
 
     # reseek, save to raw_data, reseek again.
     # TODO: remove raw_data altogether as its inefficient