Merge pull request luojilab#16 from luojilab/develop

feat: add CBC mode

Author: tcitry

.gitignore

@@ -8,3 +8,4 @@ build
 client/out
 client/client.iml
 client/java/META-INF
+__pycache__

README.md

@@ -11,7 +11,7 @@ A Django model field that encrypt your data when save to and decrypt when get fr
 * Use settings.SECRET_KEY as secret key default or anyelse which length >= 32
 * Support CharField、TextField、IntegerField、EmailField
 * Support Django ORM's `get()`、`filter()` query method
-* Use AES-256-ECB algorithm
+* Support AES-256-ECB/AES-256-CBC
 * Support PostgreSQL and MySQL database
 * Support Django model field `db_index` and `unique` attributes
 
@@ -52,10 +52,25 @@ c.decrypt('-bYijegsEDrmS1s7ilnspA==')   # some_address
 
 ### Settings
 
+#### MIRAGE_SECRET_KEY
+
 You can use the `settings.SECRET_KEY` as default key, if you want custom another key for mirage, set the `MIRAGE_SECRET_KEY` in settings.
 
 Mirage will get the `settings.MIRAGE_SECRET_KEY` first, if not set, mirage will get the `settings.SECRET_KEY`.
 
+#### MIRAGE_CIPHER_MODE
+
+`MIRAGE_CIPHER_MODE` is optional, choices are:
+
+- `ECB`
+- `CBC`
+
+If you don't set, default is `ECB`.
+
+#### MIRAGE_CIPHER_IV
+
+`MIRAGE_CIPHER_IV` is optional, if you don't set, it will use a default: "1234567890abcdef", it's length must be 16.
+
 
 ## Model Fields
 
@@ -113,6 +128,8 @@ from mirage import exceptions
 
 ## Performance
 
+### With ECB mode
+
 Migrate data: 6000,000 columns takes 40 minutes, Average 1 column/2.5ms
 
 Only encrypt/decrypt: Average 1 value/ms

mirage/crypto.py

@@ -5,44 +5,79 @@
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 from django.conf import settings
-from django.utils.encoding import force_bytes, force_text
+from django.utils.encoding import force_bytes, force_str
 
 
+class BaseCipher:
+    def __init__(self, key, iv) -> None:
+        self.key = key
+        self.iv = iv
+
+    def encrypt(self, text) -> str:
+        return text
+
+    def decrypt(self, encrypted) -> str:
+        return encrypted
+
+
+class ECBCipher(BaseCipher):
+
+    def encrypt(self, text):
+        encryptor = Cipher(algorithms.AES(self.key), modes.ECB(), default_backend()).encryptor()
+        padder = padding.PKCS7(algorithms.AES(self.key).block_size).padder()
+        padded_data = padder.update(force_bytes(text)) + padder.finalize()
+        encrypted_text = encryptor.update(padded_data) + encryptor.finalize()
+        return force_str(base64.urlsafe_b64encode(encrypted_text))
+
+    def decrypt(self, encrypted):
+        decryptor = Cipher(algorithms.AES(self.key), modes.ECB(), default_backend()).decryptor()
+        padder = padding.PKCS7(algorithms.AES(self.key).block_size).unpadder()
+        decrypted_text = decryptor.update(base64.urlsafe_b64decode(encrypted))
+        unpadded_text = padder.update(decrypted_text) + padder.finalize()
+        return force_str(unpadded_text)
+
+class CBCCipher(BaseCipher):
+
+    def encrypt(self, text) -> str:
+        encryptor = Cipher(algorithms.AES(self.key), modes.CBC(self.iv), default_backend()).encryptor()
+        padder = padding.PKCS7(algorithms.AES(self.key).block_size).padder()
+        padded_data = padder.update(force_bytes(text)) + padder.finalize()
+        encrypted_text = encryptor.update(padded_data) + encryptor.finalize()
+        return force_str(base64.urlsafe_b64encode(encrypted_text))
+
+    def decrypt(self, encrypted) -> str:
+        decryptor = Cipher(algorithms.AES(self.key), modes.CBC(self.iv), default_backend()).decryptor()
+        padder = padding.PKCS7(algorithms.AES(self.key).block_size).unpadder()
+        decrypted_text = decryptor.update(base64.urlsafe_b64decode(encrypted))
+        unpadded_text = padder.update(decrypted_text) + padder.finalize()
+        return force_str(unpadded_text)
+
 class Crypto:
 
-    def __init__(self, key=None):
+    def __init__(self, key=None, mode=None, iv=None):
         if key is None:
             key = getattr(settings, "MIRAGE_SECRET_KEY", None) or getattr(settings, "SECRET_KEY")
         assert len(key) >= 32, "mirage key length must more than 32!"
-        self.key = base64.urlsafe_b64encode(force_bytes(key))[:32]
+        key = base64.urlsafe_b64encode(force_bytes(key))[:32]
+        if mode is None:
+            mode = getattr(settings, "MIRAGE_CIPHER_MODE", "ECB")
+        if iv is None:
+            iv=getattr(settings, "MIRAGE_CIPHER_IV", "1234567890abcdef")
+        self.cipher = eval(f"{mode}Cipher")(key=key, iv=force_bytes(iv))
 
     def encrypt(self, text):
         if text is None:
             return None
         try:
-            self.try_decrypt(text)
+            self.cipher.decrypt(text)
             return text
         except Exception:
-            return self.try_encrypt(text)
-
-    def try_encrypt(self, text):
-        encryptor = Cipher(algorithms.AES(self.key), modes.ECB(), default_backend()).encryptor()
-        padder = padding.PKCS7(algorithms.AES(self.key).block_size).padder()
-        padded_data = padder.update(force_bytes(text)) + padder.finalize()
-        encrypted_text = encryptor.update(padded_data) + encryptor.finalize()
-        return force_text(base64.urlsafe_b64encode(encrypted_text))
-
-    def try_decrypt(self, encrypted_text):
-        decryptor = Cipher(algorithms.AES(self.key), modes.ECB(), default_backend()).decryptor()
-        padder = padding.PKCS7(algorithms.AES(self.key).block_size).unpadder()
-        decrypted_text = decryptor.update(base64.urlsafe_b64decode(encrypted_text))
-        unpadded_text = padder.update(decrypted_text) + padder.finalize()
-        return force_text(unpadded_text)
+            return self.cipher.encrypt(text)
 
-    def decrypt(self, encrypted_text):
-        if encrypted_text is None:
+    def decrypt(self, encrypted):
+        if encrypted is None:
             return None
         try:
-            return self.try_decrypt(encrypted_text)
+            return self.cipher.decrypt(encrypted)
         except Exception:
-            return encrypted_text
+            return encrypted

setup.py

@@ -7,7 +7,7 @@
 os.chdir(os.path.normpath(os.path.join(os.path.abspath(__file__), os.pardir)))
 setup(
     name='django-mirage-field',
-    version='1.1.8',
+    version='2.0.0.beta1',
     install_requires=[
         "cryptography",
         "tqdm",

tests/settings/base.py

@@ -2,6 +2,8 @@
 SECRET_KEY = "ZGphbmdvLW1pcmFnZS1maWVsZCBoYXZlIGZ1bg=="
 
 MIRAGE_SECRET_KEY = "MIRAGE_SECRET_KEYMIRAGE_SECRET_KEY"
+MIRAGE_CIPHER_MODE = "CBC"
+MIRAGE_CIPHER_IV = "1234567890abcdef"
 
 INSTALLED_APPS = [
     'django.contrib.auth',

tests/test_crypto.py

@@ -1,20 +1,30 @@
 from django.test import TestCase
-from django.conf import settings
 from mirage.crypto import Crypto
 
 
-class TestCrypto(TestCase):
+class TestCryptoECB(TestCase):
 
     def setUp(self):
-        self.crypto = Crypto()
+        self.crypto = Crypto(mode='ECB')
         self.value = 'hello,text'
-        if getattr(settings, "MIRAGE_SECRET_KEY", None):
-            self.encrypted = "4DIIbNsZPqO1DuXX1GjpkQ=="
-        else:
-            self.encrypted = 'pyy1FL2ftjBjUrJlGjgl3g=='
+        self.encrypted = "4DIIbNsZPqO1DuXX1GjpkQ=="
 
-    def test_encrypt(self):
+    def test_ecb_encrypt(self):
         self.assertEqual(self.crypto.encrypt(self.value), self.encrypted)
 
-    def test_decrypt(self):
+    def test_ecb_decrypt(self):
+        self.assertEqual(self.crypto.decrypt(self.encrypted), self.value)
+
+
+class TestCryptoCBC(TestCase):
+
+    def setUp(self) -> None:
+        self.crypto = Crypto(mode='CBC')
+        self.value = 'hello,text'
+        self.encrypted = "E_RFOSafjW-FQ-PDkXkv5g=="
+
+    def test_cbc_encrypt(self):
+        self.assertEqual(self.crypto.encrypt(self.value), self.encrypted)
+
+    def test_cbc_decrypt(self):
         self.assertEqual(self.crypto.decrypt(self.encrypted), self.value)