From b7f79744d078810cf8d6bb820aba5a22d6d65cff Mon Sep 17 00:00:00 2001
From: green <41323182+greeeen-dev@users.noreply.github.com>
Date: Sat, 7 Dec 2024 00:39:54 +0100
Subject: [PATCH] fix data backup

---
 unifier.py          | 26 ++++++++++++++++++--------
 utils/compressor.py |  5 ++++-
 utils/secrets.py    |  4 ++--
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/unifier.py b/unifier.py
index 1966d21a..ec00ad67 100644
--- a/unifier.py
+++ b/unifier.py
@@ -17,6 +17,7 @@
 """
 
 import nextcord
+import ujson
 from nextcord import Interaction, ApplicationError
 from nextcord.ext import commands
 import aiohttp
@@ -245,17 +246,22 @@ def __init__(self, *args, **kwargs):
         self.__encrypted = kwargs.get('encrypt', False)
 
         # Ensure necessary keys exist
-        self.update({'rooms': {}, 'emojis': [], 'nicknames': {}, 'blocked': {}, 'banned': {},
-                     'moderators': [], 'avatars': {}, 'experiments': {}, 'experiments_info': {}, 'colors': {},
-                     'external_bridge': [], 'modlogs': {}, 'trusted': [], 'report_threads': {}, 'fullbanned': [],
-                     'exp': {}, 'appealban': [], 'languages': {}, 'settings': {}, 'invites': {}, 'underattack': [],
-                     'rooms_count': {}, 'connections_count': {}, 'allocations_override': {}, 'filters': {}}
-                    )
+        self.update(self.base)
         self.threads = []
 
         # Load data
         self.load_data()
 
+    @property
+    def base(self):
+        return {
+            'rooms': {}, 'emojis': [], 'nicknames': {}, 'blocked': {}, 'banned': {},
+            'moderators': [], 'avatars': {}, 'experiments': {}, 'experiments_info': {}, 'colors': {},
+            'external_bridge': [], 'modlogs': {}, 'trusted': [], 'report_threads': {}, 'fullbanned': [],
+            'exp': {}, 'appealban': [], 'languages': {}, 'settings': {}, 'invites': {}, 'underattack': [],
+            'rooms_count': {}, 'connections_count': {}, 'allocations_override': {}, 'filters': {}
+        }
+
     @property
     def save_lock(self):
         return self.__save_lock
@@ -281,20 +287,24 @@ def load_data(self):
             self.update(data)
         except FileNotFoundError:
             pass  # If the file is not found, initialize an empty dictionary
+        except json.JSONDecodeError:
+            pass  # If the file is corrupted, initialize an empty dictionary
 
     def save(self):
         tosave = {}
         for key in self.keys():
+            if not key in self.base.keys():
+                continue
             tosave.update({key: self[key]})
         if self.__save_lock:
             return
         if self.__encrypted:
-            data = jsontools.dumps_bytes(tosave)
+            data = jsontools.dumps_bytes(dict(tosave))
             self.__secure_storage.save(data, self.file_path)
         else:
             with open(self.file_path, 'w') as file:
                 # noinspection PyTypeChecker
-                json.dump(tosave, file, indent=4)
+                json.dump(dict(tosave), file, indent=4)
 
         return
 
diff --git a/utils/compressor.py b/utils/compressor.py
index 33280516..3deed0c8 100644
--- a/utils/compressor.py
+++ b/utils/compressor.py
@@ -36,10 +36,13 @@ def compress(data: bytes, filename: Optional[str], chunk_size: int, level: int =
     with compressor.stream_writer(target) as f:
         for i in range(0, len(data), chunk_size):
             f.write(data[i:i + chunk_size])
+        f.flush()
+        target.seek(0)
+        result = target.read()
 
     # return bytes if filename is None
     if not filename:
-        return target.getvalue()
+        return result
 
 def decompress(file: Union[bytes, str], chunk_size: int) -> bytes:
     """Decompresses a file compressed using Zstandard."""
diff --git a/utils/secrets.py b/utils/secrets.py
index f4a3cb72..17ab2d72 100644
--- a/utils/secrets.py
+++ b/utils/secrets.py
@@ -95,7 +95,7 @@ def __init__(self, password):
     def encrypt(self, data):
         data, tag, nonce, salt = self.__encryptor.encrypt(data, self.__password)
         return {
-            'data': data,
+            'data': base64.b64encode(data).decode('ascii'),
             'tag': tag,
             'nonce': nonce,
             'salt': salt
@@ -469,4 +469,4 @@ def load(self, filename):
         with open(filename, 'r') as file:
             data = json.load(file)
 
-        return self.__rawencryptor.decrypt(data['data'], data['tag'], data['nonce'], data['salt'])
+        return self.__rawencryptor.decrypt(base64.b64decode(data['data']), data['tag'], data['nonce'], data['salt'])