From bb6e222255c7dbbb3b6f2c4ff14c54a49ee6b899 Mon Sep 17 00:00:00 2001
From: dogeystamp <dogeystamp@disroot.org>
Date: Mon, 8 May 2023 18:49:36 -0400
Subject: [PATCH] /files/: clean up files

---
 sachet/server/files/views.py |  1 +
 sachet/server/models.py      | 18 +++++++++++++++++-
 tests/test_files.py          |  5 +++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/sachet/server/files/views.py b/sachet/server/files/views.py
index e4486f6..679c9f6 100644
--- a/sachet/server/files/views.py
+++ b/sachet/server/files/views.py
@@ -117,6 +117,7 @@ class FileContentAPI(MethodView):
 
         if upload.completed:
             share.initialized = True
+            db.session.delete(upload)
             db.session.commit()
             return jsonify(dict(status="success", message="Upload completed.")), 201
         else:
diff --git a/sachet/server/models.py b/sachet/server/models.py
index 8ad2d24..adef63f 100644
--- a/sachet/server/models.py
+++ b/sachet/server/models.py
@@ -6,6 +6,7 @@ from bitmask import Bitmask
 from marshmallow import fields, ValidationError
 from flask import request, jsonify, url_for, current_app
 from sqlalchemy_utils import UUIDType
+from sqlalchemy import event
 import uuid
 
 
@@ -282,6 +283,13 @@ class Share(db.Model):
     def get_handle(self):
         return storage.get_file(str(self.share_id))
 
+    @classmethod
+    def __declare_last__(cls):
+        @event.listens_for(cls, "before_delete")
+        def share_before_delete(mapper, connection, share):
+            file = share.get_handle()
+            file.delete()
+
 
 class Upload(db.Model):
     """Upload instance for a given file.
@@ -325,7 +333,8 @@ class Upload(db.Model):
     completed = db.Column(db.Boolean, nullable=False, default=False)
 
     chunks = db.relationship(
-        "Chunk", backref=db.backref("upload"), order_by="Chunk.chunk_id"
+        "Chunk", backref=db.backref("upload"), order_by="Chunk.chunk_id",
+        cascade="all, delete"
     )
 
     def __init__(self, upload_id, total_chunks, share_id):
@@ -408,3 +417,10 @@ class Chunk(db.Model):
         file = storage.get_file(self.filename)
         with file.open(mode="wb") as f:
             f.write(data)
+
+    @classmethod
+    def __declare_last__(cls):
+        @event.listens_for(cls, "before_delete")
+        def chunk_before_delete(mapper, connection, chunk):
+            file = storage.get_file(chunk.filename)
+            file.delete()
diff --git a/tests/test_files.py b/tests/test_files.py
index d720aee..9c8ec6e 100644
--- a/tests/test_files.py
+++ b/tests/test_files.py
@@ -1,6 +1,8 @@
 import pytest
+from os.path import basename
 from io import BytesIO
 from werkzeug.datastructures import FileStorage
+from sachet.server import storage
 import uuid
 
 """Test file share endpoints."""
@@ -55,6 +57,9 @@ class TestSuite:
         )
         assert resp.status_code == 404
 
+        for f in storage.list_files():
+            assert basename(url) not in f.name
+
     def test_modification(self, client, users, auth, rand, upload):
         # create share
         resp = client.post(