SNOW-1825621 OAuth code flow PKCE support (#2137)

Author: sfc-gh-mkeller

src/snowflake/connector/auth/oauth_code.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import base64
+import hashlib
 import json
 import logging
 import secrets
@@ -55,6 +56,7 @@ def __init__(
         token_request_url: str,
         redirect_uri: str,
         scope: str,
+        pkce: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -72,6 +74,10 @@ def __init__(
         logger.debug("chose oauth state: %s", "".join("*" for _ in self._state))
         self._oauth_token = None
         self._protocol = "http"
+        self.pkce = pkce
+        if pkce:
+            logger.debug("oauth pkce is going to be used")
+        self._verifier: str | None = None
 
     def reset_secrets(self) -> None:
         self._oauth_token = None
@@ -104,6 +110,18 @@ def construct_url(self) -> str:
         }
         if self.scope:
             params["scope"] = self.scope
+        if self.pkce:
+            self._verifier = secrets.token_urlsafe(43)
+            # calculate challenge and verifier
+            challenge = (
+                base64.urlsafe_b64encode(
+                    hashlib.sha256(self._verifier.encode("utf-8")).digest()
+                )
+                .decode("utf-8")
+                .rstrip("=")
+            )
+            params["code_challenge"] = challenge
+            params["code_challenge_method"] = "S256"
         url_params = urllib.parse.urlencode(params)
         url = f"{self.authentication_url}?{url_params}"
         return url
@@ -186,6 +204,10 @@ def prepare(
         }
         if self.client_secret:
             fields["client_secret"] = self.client_secret
+        if self.pkce:
+            assert self._verifier is not None
+            fields["code_verifier"] = self._verifier
+
         resp = urllib3.PoolManager().request_encode_body(  # TODO: use network pool to gain use of proxy settings and so on
             "POST",
             self.token_request_url,

src/snowflake/connector/connection.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import atexit
+import collections.abc
 import logging
 import os
 import pathlib
@@ -334,6 +335,11 @@ def _get_private_bytes_from_file(
         str,
         # SNOW-1825621: OAUTH implementation
     ),
+    "oauth_security_features": (
+        ("pkce",),
+        collections.abc.Iterable,  # of strings
+        # SNOW-1825621: OAUTH PKCE
+    ),
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")
@@ -1088,7 +1094,7 @@ def __open_connection(self):
                     self.auth_class = AuthByWebBrowser(
                         application=self.application,
                         protocol=self._protocol,
-                        host=self.host,
+                        host=self.host,  # TODO: delete this?
                         port=self.port,
                         timeout=self.login_timeout,
                         backoff_generator=self._backoff_generator,
@@ -1125,6 +1131,8 @@ def __open_connection(self):
                     backoff_generator=self._backoff_generator,
                 )
             elif self._authenticator == OAUTH_AUTHORIZATION_CODE:
+                pkce = "pkce" in map(lambda e: e.lower(), self._oauth_security_features)
+
                 if self._oauth_client_id is None:
                     Error.errorhandler_wrapper(
                         self,
@@ -1150,6 +1158,7 @@ def __open_connection(self):
                     ),
                     redirect_uri="http://127.0.0.1:{port}/",
                     scope=self._oauth_scope,
+                    pkce=pkce,
                 )
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (