Merge pull request #12 from DarthUdp/main

- Storage format switched to binary
- checksum added 
- File format versioning & feature flags for encryption

elara/elara.py

@@ -42,6 +42,8 @@ def __init__(self, path, commitdb, key_path=None):
         self.path = os.path.expanduser(path)
         self.commitdb = commitdb
         self.lru = LRU()
+        # this is in place to prevent opening incompatible databases between versions of the storage format
+        self.db_format_version = 0x0001
 
         # Since key file is generated first, invalid token error for pre existing open dbs
 
@@ -65,15 +67,15 @@ def __init__(self, path, commitdb, key_path=None):
 
     def _load(self):
         if self.key:
-            self.db = Util.readAndDecrypt(self)
+            self.db = Util.read_and_decrypt(self)
         else:
-            self.db = Util.readJSON(self)
+            self.db = Util.read_plain_db(self)
 
     def _dump(self):
         if self.key:
-            Util.encryptAndStore(self)  # Enclose in try-catch
+            Util.encrypt_and_store(self)  # Enclose in try-catch
         else:
-            Util.storeJSON(self)
+            Util.store_plain_db(self)
 
     def _autocommit(self):
         if self.commitdb:

elara/elarautil.py

@@ -4,62 +4,108 @@
 
 This source code is licensed under the BSD-style license found in the LICENSE file in the root directory of this source tree.
 """
+from typing import Dict
 
-from cryptography.fernet import Fernet
-import json
+import msgpack
 import os
-from .exceptions import FileAccessError, FileKeyError
+from zlib import crc32
+
+from cryptography.fernet import Fernet
+
+from .exceptions import FileAccessError, FileKeyError, LoadChecksumError, LoadIncompatibleDB
 
 
 class Util:
     @staticmethod
-    def readJSON(obj):
-        try:
-            curr_db = json.load(open(obj.path, "rb"))
-        except Exception:
-            # print("Read JSON error. File might be encrypted. Run in secure mode.")
-            raise FileAccessError(
-                "Read JSON error. File might be encrypted. Run in secure mode with key path."
-            )
-        return curr_db
+    def check_mag(mag):
+        return mag == b"ELDB"
 
     @staticmethod
-    def storeJSON(obj):
-        try:
-            json.dump(obj.db, open(obj.path, "wt"), indent=4)
-        except Exception:
-            raise FileAccessError(
-                "Store JSON error. File might be encrypted. Run in secure mode with key path."
-            )
+    def check_encrypted(version):
+        # if msb of version number is set the db is encrypted
+        return (version & (1 << 15)) != 0
+
+    @staticmethod
+    def read_plain_db(obj) -> Dict:
+        with open(obj.path, "rb") as fctx:
+            if not Util.check_mag(fctx.read(4)):
+                raise FileAccessError("File magic number not known")
+            version = int.from_bytes(fctx.read(2), "little", signed=False)
+            # check for encryption before trying anything
+            if Util.check_encrypted(version):
+                raise FileAccessError("This file is encrypted, run in secure mode")
+            checksum = int.from_bytes(fctx.read(4), "little", signed=False)
+            data = fctx.read()
+            calculated_checksum = crc32(data)
+            if calculated_checksum != checksum:
+                raise LoadChecksumError(
+                    f"calculated checksum: {calculated_checksum} is different from stored checksum {checksum}")
+            elif version != obj.db_format_version:
+                raise LoadIncompatibleDB(f"db format version {version} is incompatible with {obj.db_format_version}")
+            try:
+                curr_db = msgpack.unpackb(data)
+            except FileNotFoundError:
+                raise FileAccessError(
+                    "File not found"
+                )
+            return curr_db
+
+    @staticmethod
+    def store_plain_db(obj):
+        with open(obj.path, "wb") as fctx:
+            try:
+                data = msgpack.packb(obj.db)
+                buffer = b"ELDB"
+                buffer += obj.db_format_version.to_bytes(2, "little")
+                buffer += (crc32(data)).to_bytes(4, "little")
+                buffer += data
+                fctx.write(buffer)
+            except FileExistsError:
+                raise FileAccessError(
+                    "File already exists"
+                )
 
     @staticmethod
-    def readAndDecrypt(obj):
+    def read_and_decrypt(obj):
         if obj.key:
             fernet = Fernet(obj.key)
-            encrypted_data = None
             try:
-                with open(obj.path, "rb") as file:
-                    encrypted_data = file.read()
-            except Exception:
+                with open(obj.path, "rb") as fctx:
+                    if not Util.check_mag(fctx.read(4)):
+                        raise FileAccessError("File magic number not known")
+                    version = int.from_bytes(fctx.read(2), "little")
+                    if not Util.check_encrypted(version):
+                        raise FileAccessError("File is marked not encrypted, you might have a corrupt db")
+                    checksum = int.from_bytes(fctx.read(4), "little")
+                    encrypted_data = fctx.read()
+                    calculated_checksum = crc32(encrypted_data)
+                    if calculated_checksum != checksum:
+                        raise LoadChecksumError(
+                            f"calculated checksum: {calculated_checksum} is different from stored checksum {checksum}")
+            except FileNotFoundError:
                 raise FileAccessError("File open & read error")
             decrypted_data = fernet.decrypt(encrypted_data)
-            return json.loads(decrypted_data.decode("utf-8"))
+            return msgpack.unpackb(decrypted_data)
         else:
             return None
 
     @staticmethod
-    def encryptAndStore(obj):
+    def encrypt_and_store(obj):
         if obj.key:
             fernet = Fernet(obj.key)
-            db_snapshot = json.dumps(obj.db)
-            db_byte = db_snapshot.encode("utf-8")
-            encrypted_data = fernet.encrypt(db_byte)
+            db_snapshot = msgpack.packb(obj.db)
+            buffer = b"ELDB"
+            # set version msb
+            buffer += (obj.db_format_version | 1 << 15).to_bytes(2, "little")
+            encrypted_data = fernet.encrypt(db_snapshot)
+            buffer += crc32(encrypted_data).to_bytes(4, "little")
+            buffer += encrypted_data
             try:
                 with open(obj.path, "wb") as file:
-                    file.write(encrypted_data)
+                    file.write(buffer)
                     return True
-            except Exception:
-                raise FileAccessError("File open & write error")
+            except FileExistsError:
+                raise FileAccessError("File exists")
         else:
             return False
 

elara/exceptions.py

@@ -5,6 +5,7 @@
 This source code is licensed under the BSD-style license found in the LICENSE file in the root directory of this source tree.
 """
 
+
 # Add all custom exception classes here
 
 
@@ -24,3 +25,11 @@ def __init__(self, message):
 
     def __str__(self):
         return f"Error -> {self.message}"
+
+
+class LoadChecksumError(Exception):
+    pass
+
+
+class LoadIncompatibleDB(Exception):
+    pass

elara/shared.py

@@ -15,9 +15,9 @@
 
 def retdb(self):
     if self.key:
-        return Util.readAndDecrypt(self)
+        return Util.read_and_decrypt(self)
     else:
-        return Util.readJSON(self)
+        return Util.read_plain_db(self)
 
 
 def retmem(self):
@@ -83,7 +83,7 @@ def securedb(self, key_path=None):
             Util.keygen(new_key_path)
 
         self.key = Util.readkey(new_key_path)
-        Util.encryptAndStore(self)
+        Util.encrypt_and_store(self)
         return True
 
 
@@ -110,7 +110,7 @@ def updatekey(self, key_path=None):
         f.truncate(0)
         f.close
         self.key = Util.readkey(new_key_path)
-        Util.encryptAndStore(self)
+        Util.encrypt_and_store(self)
 
     else:
         raise FileKeyError("Update key Failed")

requirements.txt

@@ -1 +1,2 @@
 cryptography==3.4.7
+msgpack>=1.0.0

setup.py

@@ -1,39 +1,36 @@
 from distutils.core import setup
 
-with open('README.rst') as f:
+with open("README.rst") as f:
     long_description = f.read()
-    
+
 setup(
-  name = 'elara',         
-  packages = ['elara'],   
-  version = '0.3.0',      
-  license='three-clause BSD',        
-  description = 'Elara DB is an easy to use, lightweight NoSQL database written for python that can also be used as a fast in-memory cache for JSON-serializable data. Includes various methods to manipulate data structures in-memory, secure database files and export data.',   
-  long_description = long_description,
-  author = 'Saurabh Pujari',                   
-  author_email = 'saurabhpuj99@gmail.com',      
-  url = 'https://github.com/saurabh0719/elara',  
-  keywords = [
-    'database',
-    'key-value',
-    'storage',
-    'file storage',
-    'json storage',
-    'json database',
-    'key-value database' ,
-    'nosql',
-    'nosql database' 
-    'cache',
-    'file cache'
-    ],   
-  install_requires=[
-    'cryptography'
+    name="elara",
+    packages=["elara"],
+    version="0.3.0",
+    license="three-clause BSD",
+    description="Elara DB is an easy to use, lightweight NoSQL database written for python that can also be used as a fast in-memory cache for JSON-serializable data. Includes various methods to manipulate data structures in-memory, secure database files and export data.",
+    long_description=long_description,
+    author="Saurabh Pujari",
+    author_email="saurabhpuj99@gmail.com",
+    url="https://github.com/saurabh0719/elara",
+    keywords=[
+        "database",
+        "key-value",
+        "storage",
+        "file storage",
+        "json storage",
+        "json database",
+        "key-value database",
+        "nosql",
+        "nosql database" "cache",
+        "file cache",
+    ],
+    install_requires=["cryptography", "msgpack"],
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Topic :: Database",
+        "License :: OSI Approved :: BSD License",
+        "Programming Language :: Python",
     ],
-  classifiers=[
-    'Development Status :: 4 - Beta',      
-    'Intended Audience :: Developers',      
-    'Topic :: Database',
-    'License :: OSI Approved :: BSD License',   
-    'Programming Language :: Python'
-  ],
-)
+)

test/test_1.py

@@ -12,6 +12,24 @@ def test_exe(self):
         res = elara.exe("test.db", False)
         assert res is not None
 
+    def test_store_restore_data(self):
+        db = elara.exe("test.db")
+        db.set("test_key", "test_data:\"ði ıntəˈnæʃənəl fəˈnɛtık əsoʊsiˈeıʃn Y [ˈʏpsilɔn], Yen [jɛn], Yoga [ˈjoːgɑ]\"")
+        db.commit()
+        db_load = elara.exe("test.db")
+        recov_data = db.get("test_key")
+        self.assertEqual(recov_data,
+                         "test_data:\"ði ıntəˈnæʃənəl fəˈnɛtık əsoʊsiˈeıʃn Y [ˈʏpsilɔn], Yen [jɛn], Yoga [ˈjoːgɑ]\"")
+
+    def test_store_restore_data_secure(self):
+        db = elara.exe_secure("test_enc.db")
+        db.set("test_key", "test_data:\"ði ıntəˈnæʃənəl fəˈnɛtık əsoʊsiˈeıʃn Y [ˈʏpsilɔn], Yen [jɛn], Yoga [ˈjoːgɑ]\"")
+        db.commit()
+        db_load = elara.exe_secure("test_enc.db")
+        recov_data = db.get("test_key")
+        self.assertEqual(recov_data,
+                         "test_data:\"ði ıntəˈnæʃənəl fəˈnɛtık əsoʊsiˈeıʃn Y [ˈʏpsilɔn], Yen [jɛn], Yoga [ˈjoːgɑ]\"")
+
     def test_get(self):
         self.db.db["key"] = "test"
         res = self.db.db["key"]