Merge pull request #2 from jaedsonpys/encrypt-improve

Improved data encryption

CHANGELOG.md

@@ -245,4 +245,17 @@
 ## Improves
 
 - [`577ef81`](https://github.com/jaedsonpys/cookiedb/commit/577ef81): Store length of list instead of calculating every iteration;
-- [`3b861ed`](https://github.com/jaedsonpys/cookiedb/commit/3b861ed): Update `cryptography` library version.
+- [`3b861ed`](https://github.com/jaedsonpys/cookiedb/commit/3b861ed): Update `cryptography` library version.
+
+# 7.0.0
+
+- [CookieDB 7.0.0 in PyPi](https://pypi.org/project/cookiedb/7.0.0/)
+- [CookieDB 7.0.0 in GitHub Release](https://github.com/jaedsonpys/cookiedb/releases/tag/v7.0.0)
+
+## Features
+
+- [`32bfd2e`](https://github.com/jaedsonpys/cookiedb/commit/32bfd2e): Create `Cryptography` class;
+- [`94b8e9a`](https://github.com/jaedsonpys/cookiedb/commit/94b8e9a): Raise custom `InvalidTokenError` and `InvalidSignatureError` exceptions;
+- [`b319915`](https://github.com/jaedsonpys/cookiedb/commit/b319915): Use `Cryptography` in `_document` module;
+- [`190167b`](https://github.com/jaedsonpys/cookiedb/commit/190167b): Remove key base64 encode from `CookieDB.__init__`;
+- [`190167b`](https://github.com/jaedsonpys/cookiedb/commit/190167b): Remove key base64 encode from `CookieDB.__init__`;

README.md

@@ -17,22 +17,12 @@ pip install cookiedb
 from cookiedb import CookieDB
 
 database = CookieDB(key='secret')
-database.create_database('MyDatabase')
 database.open('MyDatabase')
 
-database.add('languages', {
-    'python': {
-        'name': 'Python',
-        'ext': '.py'
-    },
-    'cpp': {
-        'name': 'C++',
-        'ext': '.cpp'
-    }
-})
-
-languages = database.get('languages')
-print(f'All languages: {languages}')
+database.add('languages/python', {'name': 'Python', 'ext': '.py'})
+python_ext = database.get('languages/python/ext')
+
+print(f'Python extension: {python_ext}')
 ```
 
 ## License

cookiedb/__init__.py

@@ -9,4 +9,4 @@
 from .cookiedb import CookieDB
 from . import exceptions
 
-__version__ = '6.0.1'
+__version__ = '7.0.0'

cookiedb/_document.py

@@ -7,20 +7,18 @@
 # http://www.apache.org/licenses/LICENSE-2.0
 
 import os
+import pickle
 from datetime import datetime
 from typing import Union
 
-import pickle
-from cryptography import fernet
-
 from . import exceptions
+from ._encrypt import Cryptography
 
 
 class Document:
     def __init__(self, key: bytes, database_local: str) -> None:
-        self._fernet = fernet.Fernet(key)
+        self._crypt = Cryptography(key)
         self._document_local = database_local
-        self._key = key.decode()
 
     @staticmethod
     def _save_file(file_content: str, filepath: str) -> None:
@@ -33,11 +31,11 @@ def exists_document(self, database: str) -> bool:
 
     def _encrypt(self, obj: dict) -> str:
         pickle_file = pickle.dumps(obj)
-        encrypted_data = self._fernet.encrypt(pickle_file)
+        encrypted_data = self._crypt.encrypt(pickle_file)
         return encrypted_data
 
     def _decrypt(self, encrypted: bytes) -> dict:
-        decrypted_data = self._fernet.decrypt(encrypted)
+        decrypted_data = self._crypt.decrypt(encrypted)
         data = pickle.loads(decrypted_data)
         return data
 

cookiedb/_encrypt.py

@@ -0,0 +1,104 @@
+import hashlib
+import secrets
+
+from Crypto.Cipher import AES
+from Crypto.Util import Padding
+from Crypto.Hash import HMAC, SHA256
+
+from . import exceptions
+
+
+class Cryptography:
+    def __init__(self, key: str) -> None:
+        """initialize a Cryptography instance.
+
+        :param key: Key to encrypt and decrypt
+        :type key: str
+        """
+
+        hash_key = hashlib.sha256(key.encode()).digest()
+        self._encryption_key = hash_key[:128]
+        self._signature_key = hash_key[128:]
+
+    def _get_hmac(self, data: bytes) -> bytes:
+        hmac = HMAC.new(self._signature_key, digestmod=SHA256)
+        hmac.update(data)
+        return hmac.digest()
+
+    def _valid_hmac(self, mac: bytes, data: bytes) -> bool:
+        hmac = HMAC.new(self._signature_key, digestmod=SHA256)
+        hmac.update(data)
+
+        try:
+            hmac.verify(mac)
+        except ValueError:
+            return False
+        else:
+            return True
+
+    def encrypt(self, data: bytes) -> bytes:
+        """Encrypt a data in bytes.
+
+        The result will be a string in bytes containing
+        the length of the encrypted data, initialization
+        vector (IV), encrypted data and a MAC hash.
+
+        :param data: Any data in bytes
+        :type data: bytes
+        :return: Encrypted data
+        :rtype: bytes
+        """
+
+        random_iv = secrets.token_bytes(16)
+        padding_data = Padding.pad(data, AES.block_size)
+
+        cipher = AES.new(self._encryption_key, AES.MODE_CBC, iv=random_iv)
+        encrypted_data = cipher.encrypt(padding_data)
+
+        result = (
+            len(data).to_bytes(4, 'big')
+            + random_iv
+            + encrypted_data
+        )
+
+        hmac = self._get_hmac(result)
+        return (result + hmac)
+
+    def decrypt(self, token: bytes) -> bytes:
+        """Decrypt a token in bytes.
+
+        :param token: Encrypted token
+        :type token: bytes
+        :raises Exception: If token is invalid
+        :raises Exception: If token has a invalid signature
+        :return: Decrypted data
+        :rtype: bytes
+        """
+
+        random_iv = token[4:20]
+        mac = token[-32:]
+        encrypted_data = token[20:-32]
+
+        cipher = AES.new(self._encryption_key, AES.MODE_CBC, iv=random_iv)
+
+        if self._valid_hmac(mac, token[:-32]):
+            try:
+                decrypted_data = cipher.decrypt(encrypted_data)
+                unpad_data = Padding.unpad(decrypted_data, AES.block_size)
+            except ValueError:
+                raise exceptions.InvalidTokenError('The token is invalid') from None
+            return unpad_data
+        else:
+            raise exceptions.InvalidSignatureError('Token signature don\'t match')
+
+    def get_data_size(self, token: bytes) -> int:
+        """Return the encrypted data size.
+
+        :param token: Encrypted token
+        :type token: bytes
+        :return: Data size in bytes
+        :rtype: int
+        """
+
+        size = int.from_bytes(token[:4], 'big')
+        return size

cookiedb/cookiedb.py

@@ -6,8 +6,6 @@
 
 # http://www.apache.org/licenses/LICENSE-2.0
 
-import hashlib
-from base64 import urlsafe_b64encode
 from functools import wraps
 from typing import Any
 
@@ -26,12 +24,6 @@ def decorator(ref, *args, **kwargs):
     return decorator
 
 
-def _generate_fernet_key(text: str):
-    key_hash = hashlib.md5(text.encode()).hexdigest()
-    key = urlsafe_b64encode(key_hash.encode())
-    return key
-
-
 class CookieDB:
     def __init__(
         self,
@@ -51,11 +43,9 @@ def __init__(
 
         if not key or type(key) != str:
             raise exceptions.InvalidKeyError(f'Argument "key" must be of type "str", not "{type(key)}"')
-        else:
-            b64_key = _generate_fernet_key(key)
 
         self._open_database = None
-        self._document = document.Document(b64_key, database_local)
+        self._document = document.Document(key, database_local)
 
     def checkout(self) -> str:
         """Return opened databsase name
@@ -66,27 +56,27 @@ def checkout(self) -> str:
 
         return self._open_database
 
-    def open(self, database_name: str) -> None:
+    def open(self, database: str) -> None:
         """
         Stores the name of the database if it exists,
         otherwise an exception `DatabaseNotFoundError`
         is thrown.
 
-        :param database_name: Database name;
+        :param database: Database name;
         :return: None.
         """
 
-        database_exists = self._document.exists_document(database_name)
+        database_exists = self._document.exists_document(database)
 
         if not database_exists:
-            raise exceptions.DatabaseNotFoundError(f'Database {database_name} not found.')
+            raise exceptions.DatabaseNotFoundError(f'Database {database} not found.')
         else:
-            self._open_database = database_name
+            self._open_database = database
 
         try:
-            self._document.get_document(database_name)
-        except document.fernet.InvalidToken:
-            raise exceptions.InvalidDatabaseKeyError('Invalid database key')
+            self._document.get_document(database)
+        except (exceptions.InvalidTokenError, exceptions.InvalidSignatureError):
+            raise exceptions.InvalidDatabaseKeyError('Invalid database encryption key')
 
     def close(self) -> None:
         """Close a open database.
@@ -100,23 +90,23 @@ def close(self) -> None:
         else:
             raise exceptions.NoOpenDatabaseError('No open database.')
 
-    def create_database(self, database_name, if_not_exists: bool = False) -> None:
+    def create_database(self, name: str, if_not_exists: bool = False) -> None:
         """
         Create a database at the location specified
         in **database local** in the `CookieDB`
         class instance.
 
-        :param database_name: Database name;
+        :param name: Database name;
         :param if_not_exists: If "True", exceptions will
         not be thrown if you are trying to create a
         database that already exists;
         :return: None.
         """
 
-        if not self._document.exists_document(database_name):
-            self._document.create_document(database_name)
+        if not self._document.exists_document(name):
+            self._document.create_document(name)
         elif not if_not_exists:
-            raise exceptions.DatabaseExistsError(f'Database {database_name} already exists.')
+            raise exceptions.DatabaseExistsError(f'Database {name} already exists.')
 
     def _get_database_items(self):
         try:

cookiedb/exceptions.py

@@ -34,3 +34,13 @@ def __init__(self, message):
 class InvalidKeyError(Exception):
     def __init__(self, message):
         super().__init__(message)
+
+
+class InvalidTokenError(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+
+
+class InvalidSignatureError(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)