From 5f59bd68ca700c55813c0928a2843c43b94d47f9 Mon Sep 17 00:00:00 2001
From: Craig Rueda <craig@craigrueda.com>
Date: Mon, 13 May 2024 09:29:52 -0700
Subject: [PATCH] chore(models): Adding encrypted field checks (#28436)

---
 superset/databases/ssh_tunnel/models.py       | 15 +++++------
 superset/utils/encrypt.py                     |  9 ++++++-
 .../integration_tests/utils/encrypt_tests.py  | 27 ++++++++++++++++++-
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/superset/databases/ssh_tunnel/models.py b/superset/databases/ssh_tunnel/models.py
index 3d1411cc80c01..5c1450cec090f 100644
--- a/superset/databases/ssh_tunnel/models.py
+++ b/superset/databases/ssh_tunnel/models.py
@@ -21,9 +21,10 @@
 from flask import current_app
 from flask_appbuilder import Model
 from sqlalchemy.orm import backref, relationship
-from sqlalchemy_utils import EncryptedType
+from sqlalchemy.types import Text
 
 from superset.constants import PASSWORD_MASK
+from superset.extensions import encrypted_field_factory
 from superset.models.core import Database
 from superset.models.helpers import (
     AuditMixinNullable,
@@ -53,19 +54,15 @@ class SSHTunnel(AuditMixinNullable, ExtraJSONMixin, ImportExportMixin, Model):
 
     server_address = sa.Column(sa.Text)
     server_port = sa.Column(sa.Integer)
-    username = sa.Column(EncryptedType(sa.String, app_config["SECRET_KEY"]))
+    username = sa.Column(encrypted_field_factory.create(Text))
 
     # basic authentication
-    password = sa.Column(
-        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
-    )
+    password = sa.Column(encrypted_field_factory.create(Text), nullable=True)
 
     # password protected pkey authentication
-    private_key = sa.Column(
-        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
-    )
+    private_key = sa.Column(encrypted_field_factory.create(Text), nullable=True)
     private_key_password = sa.Column(
-        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
+        encrypted_field_factory.create(Text), nullable=True
     )
 
     export_fields = [
diff --git a/superset/utils/encrypt.py b/superset/utils/encrypt.py
index e5bbd13825c6b..899b3391afa9e 100644
--- a/superset/utils/encrypt.py
+++ b/superset/utils/encrypt.py
@@ -24,6 +24,7 @@
 from sqlalchemy.engine import Connection, Dialect, Row
 from sqlalchemy_utils import EncryptedType
 
+ENC_ADAPTER_TAG_ATTR_NAME = "__created_by_enc_field_adapter__"
 logger = logging.getLogger(__name__)
 
 
@@ -70,12 +71,18 @@ def create(
         self, *args: list[Any], **kwargs: Optional[dict[str, Any]]
     ) -> TypeDecorator:
         if self._concrete_type_adapter:
-            return self._concrete_type_adapter.create(self._config, *args, **kwargs)
+            adapter = self._concrete_type_adapter.create(self._config, *args, **kwargs)
+            setattr(adapter, ENC_ADAPTER_TAG_ATTR_NAME, True)
+            return adapter
 
         raise Exception(  # pylint: disable=broad-exception-raised
             "App not initialized yet. Please call init_app first"
         )
 
+    @staticmethod
+    def created_by_enc_field_factory(field: TypeDecorator) -> bool:
+        return getattr(field, ENC_ADAPTER_TAG_ATTR_NAME, False)
+
 
 class SecretsMigrator:
     def __init__(self, previous_secret_key: str) -> None:
diff --git a/tests/integration_tests/utils/encrypt_tests.py b/tests/integration_tests/utils/encrypt_tests.py
index 0edca2d7f7cd5..cc882ee64b074 100644
--- a/tests/integration_tests/utils/encrypt_tests.py
+++ b/tests/integration_tests/utils/encrypt_tests.py
@@ -21,7 +21,11 @@
 from sqlalchemy_utils.types.encrypted.encrypted_type import StringEncryptedType
 
 from superset.extensions import encrypted_field_factory
-from superset.utils.encrypt import AbstractEncryptedFieldAdapter, SQLAlchemyUtilsAdapter
+from superset.utils.encrypt import (
+    AbstractEncryptedFieldAdapter,
+    SecretsMigrator,
+    SQLAlchemyUtilsAdapter,
+)
 from tests.integration_tests.base_tests import SupersetTestCase
 
 
@@ -60,4 +64,25 @@ def test_custom_adapter(self):
         field = encrypted_field_factory.create(String(1024))
         self.assertTrue(isinstance(field, StringEncryptedType))
         self.assertFalse(isinstance(field, EncryptedType))
+        self.assertTrue(getattr(field, "__created_by_enc_field_adapter__"))
         self.assertEqual(self.app.config["SECRET_KEY"], field.key)
+
+    def test_ensure_encrypted_field_factory_is_used(self):
+        """
+        Ensure that the EncryptedFieldFactory is used everywhere
+        that an encrypted field is needed.
+        :return:
+        """
+        from superset.extensions import encrypted_field_factory
+
+        migrator = SecretsMigrator("")
+        encrypted_fields = migrator.discover_encrypted_fields()
+        for table_name, cols in encrypted_fields.items():
+            for col_name, field in cols.items():
+                if not encrypted_field_factory.created_by_enc_field_factory(field):
+                    self.fail(
+                        f"The encrypted column [{col_name}]"
+                        f" in the table [{table_name}]"
+                        " was not created using the"
+                        " encrypted_field_factory"
+                    )