adding support to JSONField

pgcrypto/fields.py

@@ -156,3 +156,15 @@ class TimePGPSymmetricKeyField(PGPSymmetricKeyFieldMixin, models.TimeField):
     """Float PGP symmetric key encrypted field for postgres."""
     encrypt_sql = PGP_SYM_ENCRYPT_SQL_WITH_NULLIF
     cast_type = 'TIME'
+
+
+class JSONPGPPublicKeyField(PGPPublicKeyFieldMixin, models.JSONField):
+    """JSON PGP public key encrypted field for postgres."""
+    encrypt_sql = PGP_PUB_ENCRYPT_SQL_WITH_NULLIF
+    cast_type = 'jsonb'
+
+
+class JSONPGPSymmetricKeyField(PGPSymmetricKeyFieldMixin, models.JSONField):
+    """Float PGP symmetric key encrypted field for postgres."""
+    encrypt_sql = PGP_SYM_ENCRYPT_SQL_WITH_NULLIF
+    cast_type = 'jsonb'

tests/factories.py

@@ -44,6 +44,9 @@ class EncryptedModelFactory(factory.django.DjangoModelFactory):
     datetime_pgp_sym_field = timezone.now()
     boolean_pgp_sym_field = False
 
+    json_pgp_pub_field = {"key": "value"}
+    json_pgp_sym_field = {"key": "value"}
+
     fk_model = factory.SubFactory(EncryptedFKModelFactory)
 
     class Meta:

tests/models.py

@@ -59,6 +59,9 @@ class EncryptedModel(models.Model):
     float_pgp_sym_field = fields.FloatPGPSymmetricKeyField(blank=True, null=True)
     boolean_pgp_sym_field = fields.BooleanPGPSymmetricKeyField(blank=True, null=True)
 
+    json_pgp_pub_field = fields.JSONPGPPublicKeyField(blank=True, null=True)
+    json_pgp_sym_field = fields.JSONPGPSymmetricKeyField(blank=True, null=True)
+
     fk_model = models.ForeignKey(
         EncryptedFKModel, blank=True, null=True, on_delete=models.CASCADE
     )

tests/test_fields.py

@@ -106,6 +106,8 @@ def test_fields(self):
             'float_pgp_sym_field',
             'boolean_pgp_pub_field',
             'boolean_pgp_sym_field',
+            'json_pgp_pub_field',
+            'json_pgp_sym_field',
             'fk_model',
         )
         self.assertCountEqual(fields, expected)
@@ -1617,3 +1619,117 @@ def test_write_to_diff_keys(self):
             instance.hmac_field,
             expected
         )
+
+    def test_json_pgp_pub_field(self):
+        """Test JSONPGPPublicKeyField."""
+        expected = {'a': 1, 'b': '2', 'c': [1, 2, 3]}
+        EncryptedModelFactory.create(json_pgp_pub_field=expected)
+
+        instance = EncryptedModel.objects.get()
+
+        self.assertIsInstance(
+            instance.json_pgp_pub_field,
+            dict
+        )
+
+        self.assertEqual(
+            instance.json_pgp_pub_field,
+            expected
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_pub_field__a=1)
+
+        self.assertEqual(
+            1,
+            len(items)
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_pub_field__c__contains=2)
+
+        self.assertEqual(
+            1,
+            len(items)
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_pub_field__b=2)
+
+        self.assertEqual(
+            0,
+            len(items)
+        )
+
+    def test_json_pgp_sym_field(self):
+        """Test JsonPGPSymmetricKeyField."""
+        expected = {'a': 1, 'b': '2', 'c': [1, 2, 3]}
+        EncryptedModelFactory.create(json_pgp_sym_field=expected)
+
+        instance = EncryptedModel.objects.get()
+
+        self.assertIsInstance(
+            instance.json_pgp_sym_field,
+            dict
+        )
+
+        self.assertEqual(
+            instance.json_pgp_sym_field,
+            expected
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_sym_field__a=1)
+
+        self.assertEqual(
+            1,
+            len(items)
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_sym_field__c__contains=2)
+
+        self.assertEqual(
+            1,
+            len(items)
+        )
+
+        items = EncryptedModel.objects.filter(json_pgp_sym_field__b=2)
+
+        self.assertEqual(
+            0,
+            len(items)
+        )
+
+    def test_pgp_public_key_json_form(self):
+        """Assert form field and widget for `JSONPGPSymmetricKeyField` field."""
+        expected = {'a': 1, 'b': '2', 'c': [1, 2, 3]}
+        instance = EncryptedModelFactory.create(json_pgp_pub_field=expected)
+
+        payload = {
+            'json_pgp_pub_field': expected
+        }
+
+        form = EncryptedForm(payload, instance=instance)
+        self.assertTrue(form.is_valid())
+
+        cleaned_data = form.cleaned_data
+
+        self.assertTrue(
+            cleaned_data['json_pgp_pub_field'],
+            expected
+        )
+
+    def test_pgp_symmetric_key_json_form(self):
+        """Assert form field and widget for `JSONPGPSymmetricKeyField` field."""
+        expected = {'a': 1, 'b': '2', 'c': [1, 2, 3]}
+        instance = EncryptedModelFactory.create(json_pgp_sym_field=expected)
+
+        payload = {
+            'json_pgp_sym_field': expected
+        }
+
+        form = EncryptedForm(payload, instance=instance)
+        self.assertTrue(form.is_valid())
+
+        cleaned_data = form.cleaned_data
+
+        self.assertTrue(
+            cleaned_data['json_pgp_sym_field'],
+            expected
+        )