Merge branch 'security'

dext · Jul 10, 2024 · 76ca80c · 76ca80c
2 parents 048d28b + 020bf64
commit 76ca80c
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/target_snowflake/db_sync.py b/target_snowflake/db_sync.py
@@ -14,6 +14,9 @@
 from target_snowflake.upload_clients.s3_upload_client import S3UploadClient
 from target_snowflake.upload_clients.snowflake_upload_client import SnowflakeUploadClient
 
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+
 
 def validate_config(config):
     """Validate configuration"""
@@ -22,7 +25,6 @@ def validate_config(config):
         'account',
         'dbname',
         'user',
-        'password',
         'warehouse',
         's3_bucket',
         'stage',
@@ -33,7 +35,6 @@ def validate_config(config):
         'account',
         'dbname',
         'user',
-        'password',
         'warehouse',
         'file_format'
     ]
@@ -285,6 +286,23 @@ def __init__(self, connection_config, stream_schema_message=None, table_cache=No
         else:
             self.upload_client = SnowflakeUploadClient(connection_config, self)
 
+
+    def load_private_key(self):
+        if 'private_key' not in self.connection_config:
+            return
+
+        key = self.connection_config['private_key'].encode()
+        password=self.connection_config.get('private_key_passphrase').encode()
+        private_key = serialization.load_pem_private_key(key, password=password, backend=default_backend())
+
+        bytes = private_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption()
+        )
+
+        return bytes
+
     def open_connection(self):
         """Open snowflake connection"""
         stream = None
@@ -293,13 +311,13 @@ def open_connection(self):
 
         return snowflake.connector.connect(
             user=self.connection_config['user'],
-            password=self.connection_config['password'],
-            private_key=self.connection_config.get('private_key', None),
-            private_key_file_pwd=self.connection_config.get('private_key_password', None),
+            password=self.connection_config.get('password'),
+            private_key=self.load_private_key(),
+            authenticator=self.connection_config.get('authenticator'),
             account=self.connection_config['account'],
             database=self.connection_config['dbname'],
             warehouse=self.connection_config['warehouse'],
-            role=self.connection_config.get('role', None),
+            role=self.connection_config.get('role'),
             autocommit=True,
             session_parameters={
                 # Quoted identifiers should be case sensitive