diff --git a/mirage/fields.py b/mirage/fields.py index c6db2e4..28b4f6f 100755 --- a/mirage/fields.py +++ b/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 diff --git a/setup.py b/setup.py index a18780c..a267c06 100755 --- a/setup.py +++ b/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", diff --git a/tests/apps/migrations/0004_testmodel_json.py b/tests/apps/migrations/0004_testmodel_json.py new file mode 100644 index 0000000..953abdf --- /dev/null +++ b/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={}), + ), + ] diff --git a/tests/apps/models.py b/tests/apps/models.py index 0bf87e4..68621d2 100755 --- a/tests/apps/models.py +++ b/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={}) diff --git a/tests/test_field.py b/tests/test_field.py index 89308ea..c6ffd7a 100755 --- a/tests/test_field.py +++ b/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 = 'hello@email.com' 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)