Connector metadata encryption/decryption on connector-instance model

backend/adapter_processor/migrations/0005_alter_adapterinstance_adapter_type.py

@@ -0,0 +1,27 @@
+# Generated by Django 4.2.1 on 2024-02-28 14:43
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("adapter_processor", "0004_alter_adapterinstance_adapter_type"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="adapterinstance",
+            name="adapter_type",
+            field=models.CharField(
+                choices=[
+                    ("UNKNOWN", "UNKNOWN"),
+                    ("LLM", "LLM"),
+                    ("EMBEDDING", "EMBEDDING"),
+                    ("VECTOR_DB", "VECTOR_DB"),
+                    ("OCR", "OCR"),
+                    ("X2TEXT", "X2TEXT"),
+                ],
+                db_comment="Type of adapter LLM/EMBEDDING/VECTOR_DB",
+            ),
+        ),
+    ]

backend/connector/migrations/0003_remove_connectorinstance_connector_metadata_and_more.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.1 on 2024-02-28 12:36
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("connector", "0002_connectorinstance_connector_metadata_b"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="connectorinstance",
+            name="connector_metadata",
+        ),
+        migrations.RemoveField(
+            model_name="connectorinstance",
+            name="connector_metadata_b",
+        ),
+    ]

backend/connector/migrations/0004_connectorinstance_connector_metadata.py

@@ -0,0 +1,26 @@
+# Generated by Django 4.2.1 on 2024-02-28 12:37
+
+import connector.fields
+import django_cryptography.fields
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        (
+            "connector",
+            "0003_remove_connectorinstance_connector_metadata_and_more",
+        ),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="connectorinstance",
+            name="connector_metadata",
+            field=django_cryptography.fields.encrypt(
+                connector.fields.ConnectorAuthJSONField(
+                    db_column="connector_metadata", default=dict
+                )
+            ),
+        ),
+    ]

backend/connector/models.py

@@ -6,6 +6,7 @@
 from connector_processor.connector_processor import ConnectorProcessor
 from connector_processor.constants import ConnectorKeys
 from django.db import models
+from django_cryptography.fields import encrypt
 from project.models import Project
 from utils.models.base_model import BaseModel
 from workflow_manager.workflow.models import Workflow
@@ -47,11 +48,14 @@ class ConnectorMode(models.IntegerChoices):
     connector_id = models.CharField(
         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 = encrypt(
+        ConnectorAuthJSONField(
+            db_column="connector_metadata",
+            null=False,
+            blank=False,
+            default=dict,
+        )
     )
-    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

backend/pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "social-auth-core==4.4.2", # For OAuth
     "unstract-sdk~=0.11.2",
     "unstract-adapters~=0.2.2",
+    "django-cryptography==1.1",
     # ! IMPORTANT!
     # Indirect local dependencies usually need to be added in their own projects
     # as: https://pdm-project.org/latest/usage/dependency/#local-dependencies.

backend/workflow_manager/endpoint/destination.py

@@ -6,9 +6,7 @@
 
 import fsspec
 import magic
-from account.models import EncryptionSecret
 from connector.models import ConnectorInstance
-from cryptography.fernet import Fernet
 from django.db import connection
 from fsspec.implementations.local import LocalFileSystem
 from unstract.sdk.constants import ToolExecKey
@@ -77,17 +75,6 @@ def _get_endpoint_for_workflow(
             workflow=workflow,
             endpoint_type=WorkflowEndpoint.EndpointType.DESTINATION,
         )
-        if endpoint.connector_instance:
-            encryption_secret: EncryptionSecret = EncryptionSecret.objects.get()
-            f: Fernet = Fernet(encryption_secret.key.encode("utf-8"))
-            endpoint.connector_instance.connector_metadata = json.loads(
-                f.decrypt(
-                    bytes(endpoint.connector_instance.connector_metadata_b
-                          ).decode(
-                        "utf-8"
-                    )
-                )
-            )
         return endpoint
 
     def validate(self) -> None:

backend/workflow_manager/endpoint/source.py

@@ -1,5 +1,4 @@
 import fnmatch
-import json
 import logging
 import os
 import shutil
@@ -8,10 +7,8 @@
 from typing import Any, Optional
 
 import fsspec
-from account.models import EncryptionSecret
 from connector.models import ConnectorInstance
 from connector_processor.constants import ConnectorKeys
-from cryptography.fernet import Fernet
 from django.core.files.uploadedfile import UploadedFile
 from django.db import connection
 from unstract.workflow_execution.enums import LogState
@@ -92,18 +89,6 @@ def _get_endpoint_for_workflow(
             workflow=workflow,
             endpoint_type=WorkflowEndpoint.EndpointType.SOURCE,
         )
-        if endpoint.connector_instance:
-            encryption_secret: EncryptionSecret = EncryptionSecret.objects.get()
-            f: Fernet = Fernet(encryption_secret.key.encode("utf-8"))
-            endpoint.connector_instance.connector_metadata = json.loads(
-                f.decrypt(
-                    bytes(endpoint.connector_instance.connector_metadata_b
-                          ).decode(
-                        "utf-8"
-                    )
-                )
-            )
-
         return endpoint
 
     def validate(self) -> None: