From f9214a722305bbf8130f07c68089616b9e07afcd Mon Sep 17 00:00:00 2001
From: Mark Keller <mark.keller@snowflake.com>
Date: Thu, 16 Jan 2025 10:44:50 -0800
Subject: [PATCH 1/2] implement PKCE

---
 DESCRIPTION.md                             |  1 +
 src/snowflake/connector/auth/oauth_code.py | 24 ++++++++++++++++++++++
 src/snowflake/connector/connection.py      |  8 ++++++++
 3 files changed, 33 insertions(+)

diff --git a/DESCRIPTION.md b/DESCRIPTION.md
index 5b911a918..67f48890d 100644
--- a/DESCRIPTION.md
+++ b/DESCRIPTION.md
@@ -16,6 +16,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added a feature to verify if the connection is still good enough to send queries over.
   - Added support for base64-encoded DER private key strings in the `private_key` authentication type.
   - Added support for OAuth authorization code flow.
+  - Added support for PKCE on top of OAuth authorization flow.
 
 - v3.12.4(December 3,2024)
   - Fixed a bug where multipart uploads to Azure would be missing their MD5 hashes.
diff --git a/src/snowflake/connector/auth/oauth_code.py b/src/snowflake/connector/auth/oauth_code.py
index 49de0f2a6..eba8da672 100644
--- a/src/snowflake/connector/auth/oauth_code.py
+++ b/src/snowflake/connector/auth/oauth_code.py
@@ -5,8 +5,10 @@
 from __future__ import annotations
 
 import base64
+import hashlib
 import json
 import logging
+import re
 import secrets
 import socket
 import time
@@ -54,6 +56,7 @@ def __init__(
         token_request_url: str,
         redirect_uri: str,
         scope: str,
+        pkce: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -71,6 +74,10 @@ def __init__(
         logger.debug("chose oauth state: %s", 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
@@ -102,6 +109,19 @@ def construct_url(self) -> str:
         }
         if self.scope:
             params["scope"] = self.scope
+        if self.pkce:
+            self._verifier = secrets.token_urlsafe(43)
+            self._verifier = re.sub("[^a-zA-Z0-9]+", "", self._verifier)
+            # calculate challenge and verifier
+            challenge = (
+                base64.urlsafe_b64encode(
+                    hashlib.sha256(self._verifier.encode("utf-8")).digest()
+                )
+                .decode("utf-8")
+                .replace("=", "")
+            )
+            params["code_challenge"] = challenge
+            params["code_challenge_method"] = "S256"
         url_params = urllib.parse.urlencode(params)
         url = f"{self.authentication_url}?{url_params}"
         return url
@@ -184,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,
diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py
index 670aa1003..34b872623 100644
--- a/src/snowflake/connector/connection.py
+++ b/src/snowflake/connector/connection.py
@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import atexit
+import collections.abc
 import logging
 import os
 import pathlib
@@ -333,6 +334,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_]+")
@@ -1117,6 +1123,7 @@ 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,
@@ -1142,6 +1149,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] = (

From 81627afbfe4e601f8667b8a56646e95310a9ed17 Mon Sep 17 00:00:00 2001
From: Mark Keller <mark.keller@snowflake.com>
Date: Thu, 30 Jan 2025 10:26:56 -0800
Subject: [PATCH 2/2] review feedback

---
 src/snowflake/connector/auth/oauth_code.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/snowflake/connector/auth/oauth_code.py b/src/snowflake/connector/auth/oauth_code.py
index eba8da672..7012f899d 100644
--- a/src/snowflake/connector/auth/oauth_code.py
+++ b/src/snowflake/connector/auth/oauth_code.py
@@ -8,7 +8,6 @@
 import hashlib
 import json
 import logging
-import re
 import secrets
 import socket
 import time
@@ -111,14 +110,13 @@ def construct_url(self) -> str:
             params["scope"] = self.scope
         if self.pkce:
             self._verifier = secrets.token_urlsafe(43)
-            self._verifier = re.sub("[^a-zA-Z0-9]+", "", self._verifier)
             # calculate challenge and verifier
             challenge = (
                 base64.urlsafe_b64encode(
                     hashlib.sha256(self._verifier.encode("utf-8")).digest()
                 )
                 .decode("utf-8")
-                .replace("=", "")
+                .rstrip("=")
             )
             params["code_challenge"] = challenge
             params["code_challenge_method"] = "S256"