Optional SSH Support for DSR Processing (#3374)

Co-authored-by: Sean Preston <[email protected]>

CHANGELOG.md

@@ -19,9 +19,12 @@ The types of changes are:
 ## [Unreleased](https://github.com/ethyca/fides/compare/2.15.0...main)
 
 ### Added
+
+- Included optional env vars to have postgres or Redshift connected via bastion host [#3374](https://github.com/ethyca/fides/pull/3374/)
 - Support for acknowledge button for notice-only Privacy Notices and to disable toggling them off [#3546](https://github.com/ethyca/fides/pull/3546)
 
 ### Fixed
+
 - Fix race condition with consent modal link rendering [#3521](https://github.com/ethyca/fides/pull/3521)
 - Remove the `fides-js` banner from tab order when it is hidden and move the overlay components to the top of the tab order. [#3510](https://github.com/ethyca/fides/pull/3510)
 - Disable connector dropdown in integration tab on save [#3552](https://github.com/ethyca/fides/pull/3552)

dev-requirements.txt

@@ -14,6 +14,7 @@ pytest==7.2.2
 requests-mock==1.10.0
 setuptools>=64.0.2
 sqlalchemy-stubs
+types-paramiko==3.0.0.10
 types-PyYAML==6.0.11
 types-redis==4.3.4
 types-requests

docker-compose.yml

@@ -28,6 +28,9 @@ services:
       FIDES__DEV_MODE: "True"
       FIDES__USER__ANALYTICS_OPT_OUT: "True"
       FIDES__SECURITY__ALLOW_CUSTOM_CONNECTOR_FUNCTIONS: "True"
+      FIDES__SECURITY__BASTION_SERVER_HOST: ${FIDES__SECURITY__BASTION_SERVER_HOST-}
+      FIDES__SECURITY__BASTION_SERVER_SSH_USERNAME: ${FIDES__SECURITY__BASTION_SERVER_SSH_USERNAME-}
+      FIDES__SECURITY__BASTION_SERVER_SSH_PRIVATE_KEY: ${FIDES__SECURITY__BASTION_SERVER_SSH_PRIVATE_KEY-}
       VAULT_ADDR: ${VAULT_ADDR-}
       VAULT_NAMESPACE: ${VAULT_NAMESPACE-}
       VAULT_TOKEN: ${VAULT_TOKEN-}

requirements.txt

@@ -26,6 +26,7 @@ okta==2.7.0
 openpyxl==3.0.9
 packaging==23.0
 pandas==1.4.3
+paramiko==3.1.0
 passlib[bcrypt]==1.7.4
 plotly==5.13.1
 pyarrow==6.0.0
@@ -49,6 +50,7 @@ sqlalchemy-bigquery==1.4.4
 sqlalchemy-redshift==0.8.11
 sqlalchemy-stubs==0.4
 SQLAlchemy-Utils==0.38.3
+sshtunnel==0.4.0
 toml>=0.10.1
 twilio==7.15.0
 typing_extensions==4.5.0 # pinned to work around https://github.com/pydantic/pydantic/issues/5821

src/fides/api/common_exceptions.py

@@ -219,6 +219,10 @@ class NoSuchConnectionTypeSecretSchemaError(Exception):
     """Exception for when a connection type secret schema is not found."""
 
 
+class SSHTunnelConfigNotFoundException(Exception):
+    """Exception for when Fides is configured to use an SSH tunnel without config provided."""
+
+
 class AuthenticationError(HTTPException):
     """To be raised when attempting to fetch an access token using
     invalid credentials.

src/fides/api/schemas/connection_configuration/connection_secrets_postgres.py

@@ -17,6 +17,7 @@ class PostgreSQLSchema(ConnectionConfigSecretsSchema):
         str
     ] = None  # Either the entire "url" *OR* the "host" should be supplied.
     port: Optional[int] = None
+    ssh_required: bool = False
 
     _required_components: List[str] = ["host"]
 

src/fides/api/schemas/connection_configuration/connection_secrets_redshift.py

@@ -15,6 +15,7 @@ class RedshiftSchema(ConnectionConfigSecretsSchema):
     user: Optional[str] = None
     password: Optional[str] = None
     db_schema: Optional[str] = None
+    ssh_required: bool = False
 
     _required_components: List[str] = ["host", "user", "password"]
 

src/fides/api/service/connectors/sql_connector.py

@@ -1,6 +1,9 @@
+import io
 from abc import abstractmethod
 from typing import Any, Dict, List, Optional, Type
 
+import paramiko
+import sshtunnel  # type: ignore
 from loguru import logger
 from snowflake.sqlalchemy import URL as Snowflake_URL
 from sqlalchemy import Column, text
@@ -16,7 +19,10 @@
 from sqlalchemy.sql import Executable  # type: ignore
 from sqlalchemy.sql.elements import TextClause
 
-from fides.api.common_exceptions import ConnectionException
+from fides.api.common_exceptions import (
+    ConnectionException,
+    SSHTunnelConfigNotFoundException,
+)
 from fides.api.graph.traversal import TraversalNode
 from fides.api.models.connectionconfig import ConnectionConfig, ConnectionTestStatus
 from fides.api.models.policy import Policy
@@ -46,6 +52,12 @@
     SQLQueryConfig,
 )
 from fides.api.util.collection_util import Row
+from fides.core.config import get_config
+
+CONFIG = get_config()
+
+sshtunnel.SSH_TIMEOUT = CONFIG.security.bastion_server_ssh_timeout
+sshtunnel.TUNNEL_TIMEOUT = CONFIG.security.bastion_server_ssh_tunnel_timeout
 
 
 class SQLConnector(BaseConnector[Engine]):
@@ -61,6 +73,7 @@ def __init__(self, configuration: ConnectionConfig):
             raise NotImplementedError(
                 "SQL Connectors must define their secrets schema class"
             )
+        self.ssh_server: sshtunnel._ForwardServer = None
 
     @staticmethod
     def cursor_result_to_rows(results: CursorResult) -> List[Row]:
@@ -161,6 +174,8 @@ def close(self) -> None:
         if self.db_client:
             logger.debug(" disposing of {}", self.__class__)
             self.db_client.dispose()
+        if self.ssh_server:
+            self.ssh_server.stop()
 
     def create_client(self) -> Engine:
         """Returns a SQLAlchemy Engine that can be used to interact with a database"""
@@ -176,6 +191,29 @@ def set_schema(self, connection: Connection) -> None:
         """Optionally override to set the schema for a given database that
         persists through the entire session"""
 
+    def create_ssh_tunnel(self, host: Optional[str], port: Optional[int]) -> None:
+        """Creates an SSH Tunnel to forward ports as configured."""
+        if not CONFIG.security.bastion_server_ssh_private_key:
+            raise SSHTunnelConfigNotFoundException(
+                "Fides is configured to use an SSH tunnel without config provided."
+            )
+
+        with io.BytesIO(
+            CONFIG.security.bastion_server_ssh_private_key.encode("utf8")
+        ) as binary_file:
+            with io.TextIOWrapper(binary_file, encoding="utf8") as file_obj:
+                private_key = paramiko.RSAKey.from_private_key(file_obj=file_obj)
+
+        self.ssh_server = sshtunnel.SSHTunnelForwarder(
+            (CONFIG.security.bastion_server_host),
+            ssh_username=CONFIG.security.bastion_server_ssh_username,
+            ssh_pkey=private_key,
+            remote_bind_address=(
+                host,
+                port,
+            ),
+        )
+
 
 class PostgreSQLConnector(SQLConnector):
     """Connector specific to postgresql"""
@@ -197,6 +235,38 @@ def build_uri(self) -> str:
         dbname = f"/{config.dbname}" if config.dbname else ""
         return f"postgresql://{user_password}{netloc}{port}{dbname}"
 
+    def build_ssh_uri(self, local_address: tuple) -> str:
+        """Build URI of format postgresql://[user[:password]@][ssh_host][:ssh_port][/dbname]"""
+        config = self.secrets_schema(**self.configuration.secrets or {})
+
+        user_password = ""
+        if config.username:
+            user = config.username
+            password = f":{config.password}" if config.password else ""
+            user_password = f"{user}{password}@"
+
+        local_host, local_port = local_address
+        netloc = local_host
+        port = f":{local_port}" if local_port else ""
+        dbname = f"/{config.dbname}" if config.dbname else ""
+        return f"postgresql://{user_password}{netloc}{port}{dbname}"
+
+    # Overrides SQLConnector.create_client
+    def create_client(self) -> Engine:
+        """Returns a SQLAlchemy Engine that can be used to interact with a database"""
+        config = self.secrets_schema(**self.configuration.secrets or {})
+        if config.ssh_required and CONFIG.security.bastion_server_ssh_private_key:
+            self.create_ssh_tunnel(host=config.host, port=config.port)
+            self.ssh_server.start()
+            uri = self.build_ssh_uri(local_address=self.ssh_server.local_bind_address)
+        else:
+            uri = config.url or self.build_uri()
+        return create_engine(
+            uri,
+            hide_parameters=self.hide_parameters,
+            echo=not self.hide_parameters,
+        )
+
     def set_schema(self, connection: Connection) -> None:
         """Sets the schema for a postgres database if applicable"""
         config = self.secrets_schema(**self.configuration.secrets or {})
@@ -270,6 +340,19 @@ class RedshiftConnector(SQLConnector):
 
     secrets_schema = RedshiftSchema
 
+    def build_ssh_uri(self, local_address: tuple) -> str:
+        """Build SSH URI of format redshift+psycopg2://[user[:password]@][ssh_host][:ssh_port][/dbname]"""
+        config = self.secrets_schema(**self.configuration.secrets or {})
+
+        local_host, local_port = local_address
+
+        config = self.secrets_schema(**self.configuration.secrets or {})
+
+        port = f":{local_port}" if local_port else ""
+        database = f"/{config.database}" if config.database else ""
+        url = f"redshift+psycopg2://{config.user}:{config.password}@{local_host}{port}{database}"
+        return url
+
     # Overrides BaseConnector.build_uri
     def build_uri(self) -> str:
         """Build URI of format redshift+psycopg2://user:password@[host][:port][/database]"""
@@ -280,6 +363,22 @@ def build_uri(self) -> str:
         url = f"redshift+psycopg2://{config.user}:{config.password}@{config.host}{port}{database}"
         return url
 
+    # Overrides SQLConnector.create_client
+    def create_client(self) -> Engine:
+        """Returns a SQLAlchemy Engine that can be used to interact with a database"""
+        config = self.secrets_schema(**self.configuration.secrets or {})
+        if config.ssh_required and CONFIG.security.bastion_server_ssh_private_key:
+            self.create_ssh_tunnel(host=config.host, port=config.port)
+            self.ssh_server.start()
+            uri = self.build_ssh_uri(local_address=self.ssh_server.local_bind_address)
+        else:
+            uri = config.url or self.build_uri()
+        return create_engine(
+            uri,
+            hide_parameters=self.hide_parameters,
+            echo=not self.hide_parameters,
+        )
+
     def set_schema(self, connection: Connection) -> None:
         """Sets the search_path for the duration of the session"""
         config = self.secrets_schema(**self.configuration.secrets or {})

src/fides/core/config/security_settings.py

@@ -127,6 +127,26 @@ class SecuritySettings(FidesSettings):
         description="Either enables the collection of audit log resource data or bypasses the middleware",
     )
 
+    bastion_server_host: Optional[str] = Field(
+        default=None, description="An optional field to store the bastion server host"
+    )
+    bastion_server_ssh_username: Optional[str] = Field(
+        default=None,
+        description="An optional field to store the username used to access the bastion server",
+    )
+    bastion_server_ssh_private_key: Optional[str] = Field(
+        default=None,
+        description="An optional field to store the key used to SSH into the bastion server.",
+    )
+    bastion_server_ssh_timeout: float = Field(
+        default=0.1,
+        description="The timeout in seconds for the transport socket (``socket.settimeout``)",
+    )
+    bastion_server_ssh_tunnel_timeout: float = Field(
+        default=10,
+        description="The timeout in seconds for tunnel connection (open_channel timeout)",
+    )
+
     @validator("app_encryption_key")
     @classmethod
     def validate_encryption_key_length(

tests/ops/api/v1/endpoints/test_connection_config_endpoints.py

@@ -1341,7 +1341,11 @@ def test_put_connection_config_secrets(
     ) -> None:
         """Note: this test does not attempt to actually connect to the db, via use of verify query param."""
         auth_header = generate_auth_header(scopes=[CONNECTION_CREATE_OR_UPDATE])
-        payload = {"host": "localhost", "port": "1234", "dbname": "my_test_db"}
+        payload = {
+            "host": "localhost",
+            "port": "1234",
+            "dbname": "my_test_db",
+        }
         resp = api_client.put(
             url + "?verify=False",
             headers=auth_header,
@@ -1361,6 +1365,7 @@ def test_put_connection_config_secrets(
             "password": None,
             "url": None,
             "db_schema": None,
+            "ssh_required": False,
         }
 
         payload = {"url": "postgresql://test_user:test_pass@localhost:1234/my_test_db"}
@@ -1383,6 +1388,7 @@ def test_put_connection_config_secrets(
             "password": None,
             "url": payload["url"],
             "db_schema": None,
+            "ssh_required": False,
         }
         assert connection_config.last_test_timestamp is None
         assert connection_config.last_test_succeeded is None
@@ -1477,6 +1483,7 @@ def test_put_connection_config_redshift_secrets(
             "user": "awsuser",
             "password": "test_password",
             "db_schema": "test",
+            "ssh_required": False,
         }
         resp = api_client.put(
             url + "?verify=False",
@@ -1497,6 +1504,7 @@ def test_put_connection_config_redshift_secrets(
             "password": "test_password",
             "db_schema": "test",
             "url": None,
+            "ssh_required": False,
         }
         assert redshift_connection_config.last_test_timestamp is None
         assert redshift_connection_config.last_test_succeeded is None

tests/ops/integration_tests/test_connection_configuration_integration.py

@@ -69,6 +69,7 @@ def test_postgres_db_connection_incorrect_secrets(
             "password": None,
             "url": None,
             "db_schema": None,
+            "ssh_required": False,
         }
         assert connection_config.last_test_timestamp is not None
         assert connection_config.last_test_succeeded is False
@@ -114,6 +115,7 @@ def test_postgres_db_connection_connect_with_components(
             "password": "postgres",
             "url": None,
             "db_schema": None,
+            "ssh_required": False,
         }
         assert connection_config.last_test_timestamp is not None
         assert connection_config.last_test_succeeded is True
@@ -155,6 +157,7 @@ def test_postgres_db_connection_connect_with_url(
             "password": None,
             "url": payload["url"],
             "db_schema": None,
+            "ssh_required": False,
         }
         assert connection_config.last_test_timestamp is not None
         assert connection_config.last_test_succeeded is True