From c2b77c7a01c52a6221540e42fd1dad3f828c01bf Mon Sep 17 00:00:00 2001
From: Lekuru <contact@lekuru.xyz>
Date: Fri, 22 Sep 2023 19:18:52 +0200
Subject: [PATCH] Add remove operations for s3 and files

---
 storage.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/storage.py b/storage.py
index 96f7046..3775e9b 100644
--- a/storage.py
+++ b/storage.py
@@ -15,6 +15,7 @@
 import config
 import boto3
 import utils
+import os
 import io
 
 class Storage:
@@ -252,6 +253,14 @@ def cache_replay(self, id: int, content: bytes):
             expiry=timedelta(hours=5)
         )
 
+    def remove_replay(self, id: int):
+        self.logger.debug(f'Removing replay with id "{id}"...')
+
+        if not config.S3_ENABLED:
+            return self.remove_file(f'/replays/{id}')
+        else:
+            return self.remove_from_s3('replays', str(id))
+
     def get_presigned_url(self, bucket: str, key: str, expiration: int = 900) -> Optional[str]:
         if not config.S3_ENABLED:
             return
@@ -295,6 +304,18 @@ def save_to_s3(self, content: bytes, key: str, bucket: str) -> bool:
 
         return True
 
+    def remove_from_s3(self, bucket: str, key: str) -> bool:
+        try:
+            self.s3.delete_object(
+                Bucket=bucket,
+                Key=key
+            )
+        except Exception as e:
+            self.logger.error(f'Failed to remove "{key}" from {bucket}: "{e}"')
+            return False
+
+        return True
+
     def get_from_cache(self, name: str) -> Optional[bytes]:
         return self.cache.get(name)
 
@@ -305,6 +326,15 @@ def get_file_content(self, filepath: str) -> Optional[bytes]:
         except Exception as e:
             self.logger.error(f'Failed to read file "{filepath}": {e}')
 
+    def remove_file(self, filepath: str) -> bool:
+        try:
+            os.remove(f'{config.DATA_PATH}/{filepath}')
+        except Exception as e:
+            self.logger.error(f'Failed to file "{filepath}": "{e}"')
+            return False
+
+        return True
+
     def get_from_s3(self, key: str, bucket: str) -> Optional[bytes]:
         buffer = io.BytesIO()