Fix: read pkey as DER format, not PEM (#680)

* Fix: read pkey as DER fmt instead of PEM * Fix: update tests * Fix: clean up key-loading logic * Fix: please linter * Fix: missing variable assignment * provides valid DER key in test --------- Co-authored-by: Skynet <[email protected]> Co-authored-by: Marcin Rudolf <[email protected]>
dlt-hub · Oct 10, 2023 · 1436884 · 1436884
1 parent 65866fc
commit 1436884
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 13 deletions.
diff --git a/dlt/destinations/snowflake/configuration.py b/dlt/destinations/snowflake/configuration.py
@@ -23,22 +23,24 @@ def _read_private_key(private_key: str, password: Optional[str] = None) -> bytes
         from cryptography.hazmat.primitives.asymmetric import rsa
         from cryptography.hazmat.primitives.asymmetric import dsa
         from cryptography.hazmat.primitives import serialization
+        from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
     except ModuleNotFoundError as e:
         raise MissingDependencyException("SnowflakeCredentials with private key", dependencies=[f"{version.DLT_PKG_NAME}[snowflake]"]) from e
 
     try:
-        # decode base64 encoded private key
-        private_key_decoded = base64.b64decode(private_key).decode('utf-8')
-    except binascii.Error:
-        # private key is not in base64 -> assume it's already provided in plain-text
-        private_key_decoded = private_key
-    except UnicodeDecodeError:
-        # private key cannot be decoded into utf-8 -> assume it's already provided in plain-text
-        private_key_decoded = private_key
-
-    pkey = serialization.load_pem_private_key(
-        private_key_decoded.encode(), password.encode() if password is not None else None, backend=default_backend()
-    )
+        # load key from base64-encoded DER key
+        pkey = serialization.load_der_private_key(
+            base64.b64decode(private_key),
+            password=password.encode() if password is not None else None,
+            backend=default_backend(),
+        )
+    except Exception:
+        # loading base64-encoded DER key failed -> assume it's a plain-text PEM key
+        pkey = serialization.load_pem_private_key(
+            private_key.encode(encoding="ascii"),
+            password=password.encode() if password is not None else None,
+            backend=default_backend(),
+        )
 
     return pkey.private_bytes(
         encoding=serialization.Encoding.DER,

diff --git a/tests/common/cases/secrets/encrypted-private-key-base64 b/tests/common/cases/secrets/encrypted-private-key-base64
@@ -1 +1 @@
-LS0tLS1CRUdJTiBFTkNSWVBURUQgUFJJVkFURSBLRVktLS0tLQpNSUlGSERCT0Jna3Foa2lHOXcwQkJRMHdRVEFwQmdrcWhraUc5dzBCQlF3d0hBUUlWd3E3cFY5Vm93MENBZ2dBCk1Bd0dDQ3FHU0liM0RRSUpCUUF3RkFZSUtvWklodmNOQXdjRUNKUENXSXFNQXo1dkJJSUV5Q2EyNFpoeE9xSGEKUHJWeHF0Y1lzdTQvY1ZvUWZGMVlTcnB1eUh5WDlNcjZaVU5NYitRSkRFNTY4R0ZTVU82SndGRG1yc2U2K2V2RQovZUpDSUVYdHJ0OC90N0JrWU5MSlFVeS9LMWlxSEhFZENRYXZmYUkzU0E1bEZ1RElwTzlyM2w5LzM5czZWaFFoCnlMOGIwU2tYdGFnS2hqbGhJcFlXdHMxcFJTU2VreFpKTnk5N0pKVHRJRDRRUUtJRTF1TExBQ1JVQlRwczlENVYKdmNNN1d1dXBVbEdwSmZzclltSHBYZ2lPMzBLd2VZSVFuc0NMcHgyNmgzd1BhN3FLYzdTc2xNTitRWW51WEt5MQp6bFFGRW9YWnN1ZkM5bEE5bTVhbG40RDJvaGdob1J4ckRZQlMwVmFWb1dEM1ZUK1A1bG5FNHpyK2xjQi9Ud2FDCkZlWUJoT0x2Tm15eXloZTB1cGpGaVV6YVdqZGZyTFJnSXRpUUNjd3J2T1ZaUmp3MFVPTGlDN0E3VGR6cjZhVk0KNjYwUzZZaTlXRXVIN1RvZHBGMW15YjQrTDgwVjVEclVSd0lDTzcxNGQ3MFFtMTUvQ2Q5Y3hFVHlFb2g1UDdGQworMjBDQ2R6SW0vQTBsVTgwZjg1bDZyOVJIaEFpVWtSbjhaZXhqeU5QUUlpRUloYVlYSXFwNEhJS1BocmVmUTZQCjgwdEM4Rm5NRzVqaXhnbTZBNmt3RTh1NGVoaTlWeE1FQklTaS9WWW9Bdm9hSUFIWVJqOFZGM29MaWxhUjJ2Y0EKRkQ5UDlrZVdwWjFiOEVtc2N3ZnBWWGszK1U0ZWEyeE0zdXVBSEFnOG5FQ0wxNVdxNURFRWxtbmlPTkJMOEE0NAp5RDhFcmcwYWdTeWNKM3ZROFRLQTNiby81K0NFOWZLbWJlZ01JZkhpSXpKZGNIWGFDUlRWM2wrLzlNSHNIQmEvCjZURTdib2JTTFIxUkxVSm5Fd0VmNHJBWjZIbWxMR1ZJdGRsZlZHS3FSSXRKZWdrLzY0RjMvblozeGN2R2Fvc2YKQ0l1TzF5S3FiMUxKWm5ZNDcyeTVRbnpzMmF2L1E0ZG5YL3BCMDdQbUwwUFVsTFArK3BXUTNJM2NTOFQrVTZQRgpMSE9RY2pFY2doMUdpVk1zbU9DQlNKMHBHbVFhWXlaS0E5bkp1aFFHblFrRzFDQlhqNTZLV3JvcXpIWlVhdVI0CkpscUZuaUZXeUZPT2hsd2J6KzRDYlIxYkpFKzZra05xUFcvb003b05WT0ozS1pFMlVNcGxrWVRLR1U5MEdlUEgKNEdpcVYxbjA2QVRnWG53TjRXUE1EV09NY0NmY1RjcGE4VGk3VTNIbTViQU9NZ3lXektqUUxZTkRabEhFUDhTVgpIYWtqNkU0cmc1bkpQcllLWHgrMURhS3FQZS9DMEdvaDRNZ0VmSTJUSHhkYmVvdy9MU3NkZDV5dUlEcnNvd3hGCmEwNm1HdXB0Ym5JSXpRVkJjcHBCSysyTHhuL0RFN0xOb3ZKT2d3cWs2VzhLa3NQSnRUYU9JYndUcmM4N1dmSnQKVDFiWTAzWGJCak8rRkd4L3pESmxwT1l2NmlNZkJ6WlM2Q3lYZXR4Z2pwZGxEbGMyUjc3N3ZaY2lKQkJhazJqbApEUXZiV05OZVZGdGVDVEFEN3c4WWpCd2FrWEdZNUxKYy9icC9MZzZIUG5vMkd0bCtnRWUvbkhUcElBbHptNWI1CjBPZ083eXNxb0F3Rm83cVVhbStGQ21ZUHpoVW5ZS2FqSHVZb3g2bm1aaTFqSkVPNVc2Y1QxdVJBRmdRb2czM1MKTUpBNEM3S0xHejlPcmRTZkQ2bE9vMllBdmxIeE1lS3Y4OTJXUVh0Q2xTZ1hrOVBTcGpwbjJ1S2Q1UThsd2c5QQpiYy9jMVo3Nm9PL3FadW05TTdyMDQzbFprT2RkeEd1VksyNlNPMTE0M21XaEV5U2VGbEtySU84dHV3bVN2dHJRCkU5WE1DeVJkd0YrKzY4ZmpEbHRwTm5CMmpldm1wbjRIQ1c3WGd6eEVWeGUzN0JPRGkwRTM5VTJiTndwN05seFcKQ1gra2w5Z0pRY25iMk1tVWpMNXFRNXk4b1U1aWo4L3FZb0YySHBuZWJqYjlpd1RJRDBsMFJyZS9hSEFqU01vVgpQTmxPZWFSZTZZK0QrT2RjWlhPL1ZnPT0KLS0tLS1FTkQgRU5DUllQVEVEIFBSSVZBVEUgS0VZLS0tLS0=
+MIIFHDBOBgkqhkiG9w0BBQ0wQTApBgkqhkiG9w0BBQwwHAQIVwq7pV9Vow0CAggAMAwGCCqGSIb3DQIJBQAwFAYIKoZIhvcNAwcECJPCWIqMAz5vBIIEyCa24ZhxOqHaPrVxqtcYsu4/cVoQfF1YSrpuyHyX9Mr6ZUNMb+QJDE568GFSUO6JwFDmrse6+evE/eJCIEXtrt8/t7BkYNLJQUy/K1iqHHEdCQavfaI3SA5lFuDIpO9r3l9/39s6VhQhyL8b0SkXtagKhjlhIpYWts1pRSSekxZJNy97JJTtID4QQKIE1uLLACRUBTps9D5VvcM7WuupUlGpJfsrYmHpXgiO30KweYIQnsCLpx26h3wPa7qKc7SslMN+QYnuXKy1zlQFEoXZsufC9lA9m5aln4D2ohghoRxrDYBS0VaVoWD3VT+P5lnE4zr+lcB/TwaCFeYBhOLvNmyyyhe0upjFiUzaWjdfrLRgItiQCcwrvOVZRjw0UOLiC7A7Tdzr6aVM660S6Yi9WEuH7TodpF1myb4+L80V5DrURwICO714d70Qm15/Cd9cxETyEoh5P7FC+20CCdzIm/A0lU80f85l6r9RHhAiUkRn8ZexjyNPQIiEIhaYXIqp4HIKPhrefQ6P80tC8FnMG5jixgm6A6kwE8u4ehi9VxMEBISi/VYoAvoaIAHYRj8VF3oLilaR2vcAFD9P9keWpZ1b8EmscwfpVXk3+U4ea2xM3uuAHAg8nECL15Wq5DEElmniONBL8A44yD8Erg0agSycJ3vQ8TKA3bo/5+CE9fKmbegMIfHiIzJdcHXaCRTV3l+/9MHsHBa/6TE7bobSLR1RLUJnEwEf4rAZ6HmlLGVItdlfVGKqRItJegk/64F3/nZ3xcvGaosfCIuO1yKqb1LJZnY472y5Qnzs2av/Q4dnX/pB07PmL0PUlLP++pWQ3I3cS8T+U6PFLHOQcjEcgh1GiVMsmOCBSJ0pGmQaYyZKA9nJuhQGnQkG1CBXj56KWroqzHZUauR4JlqFniFWyFOOhlwbz+4CbR1bJE+6kkNqPW/oM7oNVOJ3KZE2UMplkYTKGU90GePH4GiqV1n06ATgXnwN4WPMDWOMcCfcTcpa8Ti7U3Hm5bAOMgyWzKjQLYNDZlHEP8SVHakj6E4rg5nJPrYKXx+1DaKqPe/C0Goh4MgEfI2THxdbeow/LSsdd5yuIDrsowxFa06mGuptbnIIzQVBcppBK+2Lxn/DE7LNovJOgwqk6W8KksPJtTaOIbwTrc87WfJtT1bY03XbBjO+FGx/zDJlpOYv6iMfBzZS6CyXetxgjpdlDlc2R777vZciJBBak2jlDQvbWNNeVFteCTAD7w8YjBwakXGY5LJc/bp/Lg6HPno2Gtl+gEe/nHTpIAlzm5b50OgO7ysqoAwFo7qUam+FCmYPzhUnYKajHuYox6nmZi1jJEO5W6cT1uRAFgQog33SMJA4C7KLGz9OrdSfD6lOo2YAvlHxMeKv892WQXtClSgXk9PSpjpn2uKd5Q8lwg9Abc/c1Z76oO/qZum9M7r043lZkOddxGuVK26SO1143mWhEySeFlKrIO8tuwmSvtrQE9XMCyRdwF++68fjDltpNnB2jevmpn4HCW7XgzxEVxe37BODi0E39U2bNwp7NlxWCX+kl9gJQcnb2MmUjL5qQ5y8oU5ij8/qYoF2Hpnebjb9iwTID0l0Rre/aHAjSMoVPNlOeaRe6Y+D+OdcZXO/Vg==
diff --git a/tests/load/snowflake/test_snowflake_configuration.py b/tests/load/snowflake/test_snowflake_configuration.py
@@ -36,6 +36,7 @@ def test_connection_string_with_all_params() -> None:
 
 
 def test_to_connector_params() -> None:
+    # PEM key
     pkey_str = Path('./tests/common/cases/secrets/encrypted-private-key').read_text('utf8')
 
     creds = SnowflakeCredentials()
@@ -61,6 +62,7 @@ def test_to_connector_params() -> None:
         role='role1',
     )
 
+    # base64 encoded DER key
     pkey_str = Path('./tests/common/cases/secrets/encrypted-private-key-base64').read_text('utf8')
 
     creds = SnowflakeCredentials()