feat: Adding encrypted field factory (#14109)

* First cut at adding enc type factory

* Finalized enc type factory

* Adding unit test

* PyLinting

* Adding license

* Apply suggestions from code review

Co-authored-by: Jesse Yang <jesse.yang@airbnb.com>

* Moving things from enc -> encrypt

* CI commit

* One more fix

* Tweaking config name

* Fixing broken test

* Fixing broken test again

Co-authored-by: Jesse Yang <jesse.yang@airbnb.com>

.pylintrc

@@ -81,7 +81,7 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=long-builtin,dict-view-method,intern-builtin,suppressed-message,no-absolute-import,unpacking-in-except,apply-builtin,delslice-method,indexing-exception,old-raise-syntax,print-statement,cmp-builtin,reduce-builtin,useless-suppression,coerce-method,input-builtin,cmp-method,raw_input-builtin,nonzero-method,backtick,basestring-builtin,setslice-method,reload-builtin,oct-method,map-builtin-not-iterating,execfile-builtin,old-octal-literal,zip-builtin-not-iterating,buffer-builtin,getslice-method,metaclass-assignment,xrange-builtin,long-suffix,round-builtin,range-builtin-not-iterating,next-method-called,parameter-unpacking,unicode-builtin,unichr-builtin,import-star-module-level,raising-string,filter-builtin-not-iterating,using-cmp-argument,coerce-builtin,file-builtin,old-division,hex-method,missing-docstring,too-many-lines,ungrouped-imports,import-outside-toplevel,raise-missing-from,super-with-arguments,bad-option-value
+disable=long-builtin,dict-view-method,intern-builtin,suppressed-message,no-absolute-import,unpacking-in-except,apply-builtin,delslice-method,indexing-exception,old-raise-syntax,print-statement,cmp-builtin,reduce-builtin,useless-suppression,coerce-method,input-builtin,cmp-method,raw_input-builtin,nonzero-method,backtick,basestring-builtin,setslice-method,reload-builtin,oct-method,map-builtin-not-iterating,execfile-builtin,old-octal-literal,zip-builtin-not-iterating,buffer-builtin,getslice-method,metaclass-assignment,xrange-builtin,long-suffix,round-builtin,range-builtin-not-iterating,next-method-called,parameter-unpacking,unicode-builtin,unichr-builtin,import-star-module-level,raising-string,filter-builtin-not-iterating,using-cmp-argument,coerce-builtin,file-builtin,old-division,hex-method,missing-docstring,too-many-lines,ungrouped-imports,import-outside-toplevel,raise-missing-from,super-with-arguments,bad-option-value,too-few-public-methods
 
 
 [REPORTS]

superset/app.py

@@ -35,6 +35,7 @@ from superset.extensions import (
     celery_app,
     csrf,
     db,
+    encrypted_field_factory,
     feature_flag_manager,
     machine_auth_provider_factory,
     manifest_processor,
@@ -75,6 +76,7 @@ class SupersetIndexView(IndexView):
         return redirect("/superset/welcome/")
 
 
+# pylint: disable=R0904
 class SupersetAppInitializer:
     def __init__(self, app: Flask) -> None:
         super().__init__()
@@ -551,6 +553,7 @@ class SupersetAppInitializer:
         self.pre_init()
         # Configuration of logging must be done first to apply the formatter properly
         self.configure_logging()
+        self.configure_db_encrypt()
         self.setup_db()
         self.configure_celery()
         self.setup_event_logger()
@@ -670,6 +673,9 @@ class SupersetAppInitializer:
             self.config, self.flask_app.debug
         )
 
+    def configure_db_encrypt(self) -> None:
+        encrypted_field_factory.init_app(self.flask_app)
+
     def setup_db(self) -> None:
         db.init_app(self.flask_app)
 

superset/config.py

@@ -45,6 +45,7 @@ from superset.jinja_context import (  # pylint: disable=unused-import
 from superset.stats_logger import DummyStatsLogger
 from superset.typing import CacheConfig
 from superset.utils.core import is_test
+from superset.utils.encrypt import SQLAlchemyUtilsAdapter
 from superset.utils.log import DBEventLogger
 from superset.utils.logging_configurator import DefaultLoggingConfigurator
 
@@ -172,6 +173,17 @@ SQLALCHEMY_DATABASE_URI = "sqlite:///" + os.path.join(DATA_DIR, "superset.db")
 # SQLALCHEMY_CUSTOM_PASSWORD_STORE = lookup_password
 SQLALCHEMY_CUSTOM_PASSWORD_STORE = None
 
+#
+# The EncryptedFieldTypeAdapter is used whenever we're building SqlAlchemy models
+# which include sensitive fields that should be app-encrypted BEFORE sending
+# to the DB.
+#
+# Note: the default impl leverages SqlAlchemyUtils' EncryptedType, which defaults
+#  to AES-128 under the covers using the app's SECRET_KEY as key material.
+#
+# pylint: disable=C0103
+SQLALCHEMY_ENCRYPTED_FIELD_TYPE_ADAPTER = SQLAlchemyUtilsAdapter
+
 # The limit of queries fetched for query search
 QUERY_SEARCH_LIMIT = 1000
 

superset/connectors/druid/models.py

@@ -47,12 +47,12 @@ from sqlalchemy import (
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import backref, relationship, Session
 from sqlalchemy.sql import expression
-from sqlalchemy_utils import EncryptedType
 
-from superset import conf, db, is_feature_enabled, security_manager
+from superset import conf, db, security_manager
 from superset.connectors.base.models import BaseColumn, BaseDatasource, BaseMetric
 from superset.constants import NULL_STRING
 from superset.exceptions import SupersetException
+from superset.extensions import encrypted_field_factory
 from superset.models.core import Database
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin, QueryResult
 from superset.typing import FilterValues, Granularity, Metric, QueryObjectDict
@@ -138,7 +138,7 @@ class DruidCluster(Model, AuditMixinNullable, ImportExportMixin):
     metadata_last_refreshed = Column(DateTime)
     cache_timeout = Column(Integer)
     broker_user = Column(String(255))
-    broker_pass = Column(EncryptedType(String(255), conf.get("SECRET_KEY")))
+    broker_pass = Column(encrypted_field_factory.create(String(255)))
 
     export_fields = [
         "cluster_name",

superset/extensions.py

@@ -29,6 +29,7 @@ from werkzeug.local import LocalProxy
 
 from superset.utils.async_query_manager import AsyncQueryManager
 from superset.utils.cache_manager import CacheManager
+from superset.utils.encrypt import EncryptedFieldFactory
 from superset.utils.feature_flag_manager import FeatureFlagManager
 from superset.utils.machine_auth import MachineAuthProviderFactory
 
@@ -104,6 +105,7 @@ celery_app = celery.Celery()
 csrf = CSRFProtect()
 db = SQLA()
 _event_logger: Dict[str, Any] = {}
+encrypted_field_factory = EncryptedFieldFactory()
 event_logger = LocalProxy(lambda: _event_logger.get("event_logger"))
 feature_flag_manager = FeatureFlagManager()
 machine_auth_provider_factory = MachineAuthProviderFactory()

superset/migrations/versions/289ce07647b_add_encrypted_password_field.py

@@ -24,7 +24,6 @@ Create Date: 2015-11-21 11:18:00.650587
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy_utils import EncryptedType
 
 # revision identifiers, used by Alembic.
 revision = "289ce07647b"
@@ -32,9 +31,7 @@ down_revision = "2929af7925ed"
 
 
 def upgrade():
-    op.add_column(
+    op.add_column("dbs", sa.Column("password", sa.LargeBinary(), nullable=True))
-        "dbs", sa.Column("password", EncryptedType(sa.String(1024)), nullable=True)
-    )
 
 
 def downgrade():

superset/migrations/versions/b5998378c225_add_certificate_to_dbs.py

@@ -30,15 +30,13 @@ from typing import Dict
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy_utils import EncryptedType
 
 
 def upgrade():
     kwargs: Dict[str, str] = {}
     bind = op.get_bind()
     op.add_column(
-        "dbs",
+        "dbs", sa.Column("server_cert", sa.LargeBinary(), nullable=True, **kwargs),
-        sa.Column("server_cert", EncryptedType(sa.Text()), nullable=True, **kwargs),
     )
 
 

superset/migrations/versions/c2acd2cf3df2_alter_type_of_dbs_encrypted_extra.py

@@ -39,7 +39,7 @@ def upgrade():
             batch_op.alter_column(
                 "encrypted_extra",
                 existing_type=sa.Text(),
-                type_=EncryptedType(sa.Text()),
+                type_=sa.LargeBinary(),
                 postgresql_using="encrypted_extra::bytea",
                 existing_nullable=True,
             )
@@ -49,7 +49,7 @@ def upgrade():
                 "dbs",
                 "encrypted_extra",
                 existing_type=sa.Text(),
-                type_=EncryptedType(sa.Text()),
+                type_=sa.LargeBinary(),
                 existing_nullable=True,
             )
 
@@ -58,7 +58,7 @@ def downgrade():
     with op.batch_alter_table("dbs") as batch_op:
         batch_op.alter_column(
             "encrypted_extra",
-            existing_type=EncryptedType(sa.Text()),
+            existing_type=sa.LargeBinary(),
             type_=sa.Text(),
             existing_nullable=True,
         )

superset/migrations/versions/e553e78e90c5_add_druid_auth_py_py.py

@@ -33,7 +33,7 @@ from sqlalchemy_utils import EncryptedType
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column("clusters", sa.Column("broker_pass", EncryptedType(), nullable=True))
+    op.add_column("clusters", sa.Column("broker_pass", sa.LargeBinary(), nullable=True))
     op.add_column(
         "clusters", sa.Column("broker_user", sa.String(length=255), nullable=True)
     )

superset/models/core.py

@@ -52,11 +52,10 @@ from sqlalchemy.orm import relationship
 from sqlalchemy.pool import NullPool
 from sqlalchemy.schema import UniqueConstraint
 from sqlalchemy.sql import expression, Select
-from sqlalchemy_utils import EncryptedType
 
 from superset import app, db_engine_specs, is_feature_enabled
 from superset.db_engine_specs.base import TimeGrain
-from superset.extensions import cache_manager, security_manager
+from superset.extensions import cache_manager, encrypted_field_factory, security_manager
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
 from superset.models.tags import FavStarUpdater
 from superset.result_set import SupersetResultSet
@@ -115,7 +114,7 @@ class Database(
     # short unique name, used in permissions
     database_name = Column(String(250), unique=True, nullable=False)
     sqlalchemy_uri = Column(String(1024), nullable=False)
-    password = Column(EncryptedType(String(1024), config["SECRET_KEY"]))
+    password = Column(encrypted_field_factory.create(String(1024)))
     cache_timeout = Column(Integer)
     select_as_create_table_as = Column(Boolean, default=False)
     expose_in_sqllab = Column(Boolean, default=True)
@@ -141,9 +140,9 @@ class Database(
     """
         ),
     )
-    encrypted_extra = Column(EncryptedType(Text, config["SECRET_KEY"]), nullable=True)
+    encrypted_extra = Column(encrypted_field_factory.create(Text), nullable=True)
     impersonate_user = Column(Boolean, default=False)
-    server_cert = Column(EncryptedType(Text, config["SECRET_KEY"]), nullable=True)
+    server_cert = Column(encrypted_field_factory.create(Text), nullable=True)
     export_fields = [
         "database_name",
         "sqlalchemy_uri",

superset/utils/encrypt.py

@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from flask import Flask
+from sqlalchemy import TypeDecorator
+from sqlalchemy_utils import EncryptedType
+
+
+class AbstractEncryptedFieldAdapter(ABC):
+    @abstractmethod
+    def create(
+        self,
+        app_config: Optional[Dict[str, Any]],
+        *args: List[Any],
+        **kwargs: Optional[Dict[str, Any]]
+    ) -> TypeDecorator:
+        pass
+
+
+class SQLAlchemyUtilsAdapter(AbstractEncryptedFieldAdapter):
+    def create(
+        self,
+        app_config: Optional[Dict[str, Any]],
+        *args: List[Any],
+        **kwargs: Optional[Dict[str, Any]]
+    ) -> TypeDecorator:
+        if app_config:
+            return EncryptedType(*args, app_config["SECRET_KEY"], **kwargs)
+
+        raise Exception("Missing app_config kwarg")
+
+
+class EncryptedFieldFactory:
+    def __init__(self) -> None:
+        self._concrete_type_adapter: Optional[AbstractEncryptedFieldAdapter] = None
+        self._config: Optional[Dict[str, Any]] = None
+
+    def init_app(self, app: Flask) -> None:
+        self._config = app.config
+        self._concrete_type_adapter = self._config[
+            "SQLALCHEMY_ENCRYPTED_FIELD_TYPE_ADAPTER"
+        ]()
+
+    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)
+
+        raise Exception("App not initialized yet. Please call init_app first")

tests/utils/encrypt_tests.py

@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from typing import Any, Dict, List, Optional
+
+from sqlalchemy import String, TypeDecorator
+from sqlalchemy_utils import EncryptedType
+from sqlalchemy_utils.types.encrypted.encrypted_type import StringEncryptedType
+
+from superset.extensions import encrypted_field_factory
+from superset.utils.encrypt import AbstractEncryptedFieldAdapter, SQLAlchemyUtilsAdapter
+from tests.base_tests import SupersetTestCase
+
+
+class CustomEncFieldAdapter(AbstractEncryptedFieldAdapter):
+    def create(
+        self,
+        app_config: Optional[Dict[str, Any]],
+        *args: List[Any],
+        **kwargs: Optional[Dict[str, Any]]
+    ) -> TypeDecorator:
+        if app_config:
+            return StringEncryptedType(*args, app_config["SECRET_KEY"], **kwargs)
+        else:
+            raise Exception("Missing app_config kwarg")
+
+
+class EncryptedFieldTest(SupersetTestCase):
+    def setUp(self) -> None:
+        self.app.config[
+            "SQLALCHEMY_ENCRYPTED_FIELD_TYPE_ADAPTER"
+        ] = SQLAlchemyUtilsAdapter
+        encrypted_field_factory.init_app(self.app)
+
+        super().setUp()
+
+    def test_create_field(self):
+        field = encrypted_field_factory.create(String(1024))
+        self.assertTrue(isinstance(field, EncryptedType))
+        self.assertEqual(self.app.config["SECRET_KEY"], field.key)
+
+    def test_custom_adapter(self):
+        self.app.config[
+            "SQLALCHEMY_ENCRYPTED_FIELD_TYPE_ADAPTER"
+        ] = CustomEncFieldAdapter
+        encrypted_field_factory.init_app(self.app)
+        field = encrypted_field_factory.create(String(1024))
+        self.assertTrue(isinstance(field, StringEncryptedType))
+        self.assertFalse(isinstance(field, EncryptedType))
+        self.assertEqual(self.app.config["SECRET_KEY"], field.key)