Connector metadata encryption/decryption on connector-instance model - 2

backend/connector/migrations/0003_alter_connectorinstance_connector_metadata.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.2.1 on 2024-02-28 05:13
+
+import connector.models
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("connector", "0002_connectorinstance_connector_metadata_b"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="connectorinstance",
+            name="connector_metadata",
+            field=connector.models.EncryptedField(null=True),
+        ),
+    ]

backend/connector/models.py

@@ -1,10 +1,12 @@
+import json
 import uuid
+from typing import Any
 
-from account.models import User
-from connector.fields import ConnectorAuthJSONField
+from account.models import EncryptionSecret, User
 from connector_auth.models import ConnectorAuth
 from connector_processor.connector_processor import ConnectorProcessor
 from connector_processor.constants import ConnectorKeys
+from cryptography.fernet import Fernet
 from django.db import models
 from project.models import Project
 from utils.models.base_model import BaseModel
@@ -15,6 +17,32 @@
 CONNECTOR_NAME_SIZE = 128
 VERSION_NAME_SIZE = 64
 
+encryption_secret: EncryptionSecret = EncryptionSecret.objects.get()
+cipher_suite: Fernet = Fernet(encryption_secret.key.encode("utf-8"))
+
+
+class EncryptedField(models.Field):
+    def from_db_value(self, value) -> Any:  # type: ignore
+        if value is not None:
+            decrypted_value = cipher_suite.decrypt(value.encode()).decode()
+            return json.loads(decrypted_value)
+        return value
+
+    def to_python(self, value) -> Any:  # type: ignore
+        if isinstance(value, str):
+            decrypted_value = cipher_suite.decrypt(value.encode()).decode()
+            return json.loads(decrypted_value)
+        return value
+
+    def get_prep_value(self, value) -> Any:  # type: ignore
+        if value is not None:
+            serialized_value = json.dumps(value)
+            encrypted_value = cipher_suite.encrypt(
+                serialized_value.encode()
+            ).decode()
+            return encrypted_value
+        return value
+
 
 class ConnectorInstance(BaseModel):
     class ConnectorType(models.TextChoices):
@@ -48,9 +76,10 @@ class ConnectorMode(models.IntegerChoices):
         max_length=FLC.CONNECTOR_ID_LENGTH, default=""
     )
     # TODO Required to be removed
-    connector_metadata = ConnectorAuthJSONField(
-        db_column="connector_metadata", null=False, blank=False, default=dict
-    )
+    # connector_metadata = ConnectorAuthJSONField(
+    #     db_column="connector_metadata", null=False, blank=False, default=dict
+    # )
+    connector_metadata = EncryptedField(null=True)
     connector_metadata_b = models.BinaryField(null=True)
     connector_version = models.CharField(
         max_length=VERSION_NAME_SIZE, default=""

backend/connector/serializers.py

@@ -1,16 +1,13 @@
-import json
 import logging
 from collections import OrderedDict
 from typing import Any, Optional
 
-from account.models import EncryptionSecret
 from connector.constants import ConnectorInstanceKey as CIKey
 from connector_auth.models import ConnectorAuth
 from connector_auth.pipeline.common import ConnectorAuthHelper
 from connector_processor.connector_processor import ConnectorProcessor
 from connector_processor.constants import ConnectorKeys
 from connector_processor.exceptions import OAuthTimeOut
-from cryptography.fernet import Fernet
 from utils.serializer_utils import SerializerUtils
 
 from backend.serializers import AuditSerializer
@@ -54,17 +51,6 @@ def save(self, **kwargs):  # type: ignore
             connector_id, CIKey.CONNECTOR_MODE
         )
         kwargs[CIKey.CONNECTOR_MODE] = connector_mode.value
-
-        encryption_secret: EncryptionSecret = EncryptionSecret.objects.get()
-        f: Fernet = Fernet(encryption_secret.key.encode("utf-8"))
-        json_string: str = json.dumps(kwargs.pop(CIKey.CONNECTOR_METADATA))
-        if self.validated_data:
-            self.validated_data.pop(CIKey.CONNECTOR_METADATA)
-
-        kwargs[CIKey.CONNECTOR_METADATA_B] = f.encrypt(
-            json_string.encode("utf-8")
-        )
-
         instance = super().save(**kwargs)
         return instance
 
@@ -81,13 +67,5 @@ def to_representation(self, instance: ConnectorInstance) -> dict[str, str]:
             ] = ConnectorProcessor.get_connector_data_with_key(
                 instance.connector_id, ConnectorKeys.ICON
             )
-        encryption_secret: EncryptionSecret = EncryptionSecret.objects.get()
-        f: Fernet = Fernet(encryption_secret.key.encode("utf-8"))
-
-        rep.pop(CIKey.CONNECTOR_METADATA_B)
-        if instance.connector_metadata_b:
-            adapter_metadata = json.loads(
-                f.decrypt(bytes(instance.connector_metadata_b).decode("utf-8"))
-            )
-            rep[CIKey.CONNECTOR_METADATA] = adapter_metadata
+            rep[CIKey.CONNECTOR_METADATA] = instance.connector_metadata
         return rep