feat: add EncryptedJSONField

Author: tcitry

mirage/fields.py

@@ -1,7 +1,9 @@
+import json
 from django.core import exceptions
 from django.db import models
 from .crypto import Crypto
 from .exceptions import EncryptedFieldException
+from django.db.models.fields.json import KeyTransform
 
 
 class EncryptedMixin(models.Field):
@@ -36,6 +38,24 @@ def get_internal_type(self):
 class EncryptedTextField(EncryptedMixin, models.TextField):
     internal_type = "TextField"
 
+class EncryptedJSONField(EncryptedMixin, models.JSONField):
+    internal_type = "TextField"
+
+    def from_db_value(self, value, expression, connection):
+        if value is None:
+            return value
+        if isinstance(expression, KeyTransform) and not isinstance(value, str):
+            return value
+        try:
+            return json.loads(self.crypto.decrypt(value), cls=self.decoder)
+        except json.JSONDecodeError as e:
+            return self.crypto.decrypt(value)
+
+    def get_prep_value(self, value):
+        if value is None:
+            return value
+        return json.dumps(self.crypto.encrypt(json.dumps(value, cls=self.encoder)), cls=self.encoder)
+
 
 class EncryptedCharField(EncryptedMixin, models.CharField):
     prepared_max_length = 255

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.3.0',
+    version='1.4.0',
     install_requires=[
         "cryptography",
         "tqdm",

tests/apps/migrations/0004_testmodel_json.py

@@ -0,0 +1,19 @@
+# Generated by Django 4.0 on 2022-03-14 02:41
+
+from django.db import migrations
+import mirage.fields
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('apps', '0003_testmodel_url'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='testmodel',
+            name='json',
+            field=mirage.fields.EncryptedJSONField(default={}),
+        ),
+    ]

tests/apps/models.py

@@ -1,3 +1,4 @@
+from email.policy import default
 from django.db import models
 from mirage import fields
 
@@ -9,3 +10,4 @@ class TestModel(models.Model):
     integer = fields.EncryptedIntegerField(blank=True, null=True)
     email = fields.EncryptedEmailField(blank=True, null=True)
     url = fields.EncryptedURLField(blank=True, null=True)
+    json = fields.EncryptedJSONField(default={})

tests/test_field.py

@@ -1,3 +1,4 @@
+import json
 from django.test import TestCase
 from apps.models import TestModel
 
@@ -8,6 +9,7 @@ class TestField(TestCase):
     INTEGER = 1234567890
     EMAIL = '[email protected]'
     URL = 'https://yindongliang.com'
+    JSON = {"hello": "world", "foo": "bar"}
 
     @classmethod
     def setUpTestData(cls):
@@ -17,6 +19,7 @@ def setUpTestData(cls):
         obj.integer = cls.INTEGER
         obj.email = cls.EMAIL
         obj.url = cls.URL
+        obj.json = cls.JSON
         obj.save()
 
     def setUp(self):
@@ -41,3 +44,7 @@ def test_email_field(self):
     def test_url_field(self):
         self.assertEqual(self.obj.url, self.URL)
         self.assertEqual(type(self.obj.url), str)
+
+    def test_json_field(self):
+        self.assertEqual(self.obj.json, self.JSON)
+        self.assertEqual(type(self.obj.json), dict)