snowflake-connector-python 3.16.0__cp313-cp313-win_amd64.whl

snowflake/connector/auth/__init__.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from ._auth import Auth, get_public_key_fingerprint, get_token_from_private_key
+from .by_plugin import AuthByPlugin, AuthType
+from .default import AuthByDefault
+from .idtoken import AuthByIdToken
+from .keypair import AuthByKeyPair
+from .no_auth import AuthNoAuth
+from .oauth import AuthByOAuth
+from .oauth_code import AuthByOauthCode
+from .oauth_credentials import AuthByOauthCredentials
+from .okta import AuthByOkta
+from .pat import AuthByPAT
+from .usrpwdmfa import AuthByUsrPwdMfa
+from .webbrowser import AuthByWebBrowser
+from .workload_identity import AuthByWorkloadIdentity
+
+FIRST_PARTY_AUTHENTICATORS = frozenset(
+    (
+        AuthByDefault,
+        AuthByKeyPair,
+        AuthByOAuth,
+        AuthByOauthCode,
+        AuthByOauthCredentials,
+        AuthByOkta,
+        AuthByUsrPwdMfa,
+        AuthByWebBrowser,
+        AuthByIdToken,
+        AuthByPAT,
+        AuthByWorkloadIdentity,
+        AuthNoAuth,
+    )
+)
+
+__all__ = [
+    "AuthByPlugin",
+    "AuthByDefault",
+    "AuthByKeyPair",
+    "AuthByPAT",
+    "AuthByOAuth",
+    "AuthByOauthCode",
+    "AuthByOauthCredentials",
+    "AuthByOkta",
+    "AuthByUsrPwdMfa",
+    "AuthByWebBrowser",
+    "AuthByWorkloadIdentity",
+    "AuthNoAuth",
+    "Auth",
+    "AuthType",
+    "FIRST_PARTY_AUTHENTICATORS",
+    "get_public_key_fingerprint",
+    "get_token_from_private_key",
+]

snowflake/connector/auth/_auth.py

@@ -0,0 +1,586 @@
+from __future__ import annotations
+
+import copy
+import json
+import logging
+import uuid
+from datetime import datetime, timezone
+from threading import Thread
+from typing import TYPE_CHECKING, Any, Callable
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.serialization import (
+    Encoding,
+    NoEncryption,
+    PrivateFormat,
+    load_der_private_key,
+    load_pem_private_key,
+)
+
+from ..compat import urlencode
+from ..constants import (
+    DAY_IN_SECONDS,
+    HTTP_HEADER_ACCEPT,
+    HTTP_HEADER_CONTENT_TYPE,
+    HTTP_HEADER_SERVICE_NAME,
+    HTTP_HEADER_USER_AGENT,
+    PARAMETER_CLIENT_REQUEST_MFA_TOKEN,
+    PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL,
+)
+from ..description import (
+    COMPILER,
+    IMPLEMENTATION,
+    OPERATING_SYSTEM,
+    PLATFORM,
+    PYTHON_VERSION,
+)
+from ..errorcode import ER_FAILED_TO_CONNECT_TO_DB
+from ..errors import (
+    BadGatewayError,
+    DatabaseError,
+    Error,
+    ForbiddenError,
+    ProgrammingError,
+    ServiceUnavailableError,
+)
+from ..network import (
+    ACCEPT_TYPE_APPLICATION_SNOWFLAKE,
+    CONTENT_TYPE_APPLICATION_JSON,
+    ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE,
+    OAUTH_ACCESS_TOKEN_EXPIRED_GS_CODE,
+    PYTHON_CONNECTOR_USER_AGENT,
+    ReauthenticationRequest,
+)
+from ..sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
+from ..token_cache import TokenCache, TokenKey, TokenType
+from ..version import VERSION
+from .no_auth import AuthNoAuth
+
+if TYPE_CHECKING:
+    from . import AuthByPlugin
+
+logger = logging.getLogger(__name__)
+
+# keyring
+KEYRING_SERVICE_NAME = "net.snowflake.temporary_token"
+KEYRING_USER = "temp_token"
+KEYRING_DRIVER_NAME = "SNOWFLAKE-PYTHON-DRIVER"
+
+ID_TOKEN = "ID_TOKEN"
+MFA_TOKEN = "MFATOKEN"
+
+AUTHENTICATION_REQUEST_KEY_WHITELIST = {
+    "ACCOUNT_NAME",
+    "AUTHENTICATOR",
+    "CLIENT_APP_ID",
+    "CLIENT_APP_VERSION",
+    "CLIENT_ENVIRONMENT",
+    "EXT_AUTHN_DUO_METHOD",
+    "LOGIN_NAME",
+    "SESSION_PARAMETERS",
+    "SVN_REVISION",
+}
+
+
+class Auth:
+    """Snowflake Authenticator."""
+
+    def __init__(self, rest) -> None:
+        self._rest = rest
+        self._token_cache: TokenCache | None = None
+
+    @staticmethod
+    def base_auth_data(
+        user,
+        account,
+        application,
+        internal_application_name,
+        internal_application_version,
+        ocsp_mode,
+        login_timeout: int | None = None,
+        network_timeout: int | None = None,
+        socket_timeout: int | None = None,
+    ):
+        return {
+            "data": {
+                "CLIENT_APP_ID": internal_application_name,
+                "CLIENT_APP_VERSION": internal_application_version,
+                "SVN_REVISION": VERSION[3],
+                "ACCOUNT_NAME": account,
+                "LOGIN_NAME": user,
+                "CLIENT_ENVIRONMENT": {
+                    "APPLICATION": application,
+                    "OS": OPERATING_SYSTEM,
+                    "OS_VERSION": PLATFORM,
+                    "PYTHON_VERSION": PYTHON_VERSION,
+                    "PYTHON_RUNTIME": IMPLEMENTATION,
+                    "PYTHON_COMPILER": COMPILER,
+                    "OCSP_MODE": ocsp_mode.name,
+                    "TRACING": logger.getEffectiveLevel(),
+                    "LOGIN_TIMEOUT": login_timeout,
+                    "NETWORK_TIMEOUT": network_timeout,
+                    "SOCKET_TIMEOUT": socket_timeout,
+                },
+            },
+        }
+
+    def authenticate(
+        self,
+        auth_instance: AuthByPlugin,
+        account: str,
+        user: str,
+        database: str | None = None,
+        schema: str | None = None,
+        warehouse: str | None = None,
+        role: str | None = None,
+        passcode: str | None = None,
+        passcode_in_password: bool = False,
+        mfa_callback: Callable[[], None] | None = None,
+        password_callback: Callable[[], str] | None = None,
+        session_parameters: dict[Any, Any] | None = None,
+        # max time waiting for MFA response, currently unused
+        timeout: int | None = None,
+    ) -> dict[str, str | int | bool]:
+        logger.debug("authenticate")
+
+        # For no-auth connection, authentication is no-op, and we can return early here.
+        if isinstance(auth_instance, AuthNoAuth):
+            return {}
+
+        if timeout is None:
+            timeout = auth_instance.timeout
+
+        if session_parameters is None:
+            session_parameters = {}
+
+        request_id = str(uuid.uuid4())
+        headers = {
+            HTTP_HEADER_CONTENT_TYPE: CONTENT_TYPE_APPLICATION_JSON,
+            HTTP_HEADER_ACCEPT: ACCEPT_TYPE_APPLICATION_SNOWFLAKE,
+            HTTP_HEADER_USER_AGENT: PYTHON_CONNECTOR_USER_AGENT,
+        }
+        if HTTP_HEADER_SERVICE_NAME in session_parameters:
+            headers[HTTP_HEADER_SERVICE_NAME] = session_parameters[
+                HTTP_HEADER_SERVICE_NAME
+            ]
+        url = "/session/v1/login-request"
+
+        body_template = Auth.base_auth_data(
+            user,
+            account,
+            self._rest._connection.application,
+            self._rest._connection._internal_application_name,
+            self._rest._connection._internal_application_version,
+            self._rest._connection._ocsp_mode(),
+            self._rest._connection.login_timeout,
+            self._rest._connection._network_timeout,
+            self._rest._connection._socket_timeout,
+        )
+
+        body = copy.deepcopy(body_template)
+        # updating request body
+        auth_instance.update_body(body)
+
+        logger.debug(
+            "account=%s, user=%s, database=%s, schema=%s, "
+            "warehouse=%s, role=%s, request_id=%s",
+            account,
+            user,
+            database,
+            schema,
+            warehouse,
+            role,
+            request_id,
+        )
+        url_parameters = {"request_id": request_id}
+        if database is not None:
+            url_parameters["databaseName"] = database
+        if schema is not None:
+            url_parameters["schemaName"] = schema
+        if warehouse is not None:
+            url_parameters["warehouse"] = warehouse
+        if role is not None:
+            url_parameters["roleName"] = role
+
+        url = url + "?" + urlencode(url_parameters)
+
+        # first auth request
+        if passcode_in_password:
+            body["data"]["EXT_AUTHN_DUO_METHOD"] = "passcode"
+        elif passcode:
+            body["data"]["EXT_AUTHN_DUO_METHOD"] = "passcode"
+            body["data"]["PASSCODE"] = passcode
+
+        if session_parameters:
+            body["data"]["SESSION_PARAMETERS"] = session_parameters
+
+        logger.debug(
+            "body['data']: %s",
+            {
+                k: v if k in AUTHENTICATION_REQUEST_KEY_WHITELIST else "******"
+                for (k, v) in body["data"].items()
+            },
+        )
+
+        try:
+            ret = self._rest._post_request(
+                url,
+                headers,
+                json.dumps(body),
+                socket_timeout=auth_instance._socket_timeout,
+            )
+        except ForbiddenError as err:
+            # HTTP 403
+            raise err.__class__(
+                msg=(
+                    "Failed to connect to DB. "
+                    "Verify the account name is correct: {host}:{port}. "
+                    "{message}"
+                ).format(
+                    host=self._rest._host, port=self._rest._port, message=str(err)
+                ),
+                errno=ER_FAILED_TO_CONNECT_TO_DB,
+                sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+            )
+        except (ServiceUnavailableError, BadGatewayError) as err:
+            # HTTP 502/504
+            raise err.__class__(
+                msg=(
+                    "Failed to connect to DB. "
+                    "Service is unavailable: {host}:{port}. "
+                    "{message}"
+                ).format(
+                    host=self._rest._host, port=self._rest._port, message=str(err)
+                ),
+                errno=ER_FAILED_TO_CONNECT_TO_DB,
+                sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+            )
+
+        # waiting for MFA authentication
+        if ret["data"] and ret["data"].get("nextAction") in (
+            "EXT_AUTHN_DUO_ALL",
+            "EXT_AUTHN_DUO_PUSH_N_PASSCODE",
+        ):
+            body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+            body["data"]["EXT_AUTHN_DUO_METHOD"] = "push"
+            self.ret = {"message": "Timeout", "data": {}}
+
+            def post_request_wrapper(self, url, headers, body) -> None:
+                # get the MFA response
+                self.ret = self._rest._post_request(
+                    url,
+                    headers,
+                    body,
+                    socket_timeout=auth_instance._socket_timeout,
+                )
+
+            # send new request to wait until MFA is approved
+            t = Thread(
+                target=post_request_wrapper, args=[self, url, headers, json.dumps(body)]
+            )
+            t.daemon = True
+            t.start()
+            if callable(mfa_callback):
+                c = mfa_callback()
+                while not self.ret or self.ret.get("message") == "Timeout":
+                    next(c)
+            else:
+                # _post_request should already terminate on timeout, so this is just a safeguard
+                t.join(timeout=timeout)
+
+            ret = self.ret
+            if (
+                ret
+                and ret["data"]
+                and ret["data"].get("nextAction") == "EXT_AUTHN_SUCCESS"
+            ):
+                body = copy.deepcopy(body_template)
+                body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+                # final request to get tokens
+                ret = self._rest._post_request(
+                    url,
+                    headers,
+                    json.dumps(body),
+                    socket_timeout=auth_instance._socket_timeout,
+                )
+            elif not ret or not ret["data"] or not ret["data"].get("token"):
+                # not token is returned.
+                Error.errorhandler_wrapper(
+                    self._rest._connection,
+                    None,
+                    DatabaseError,
+                    {
+                        "msg": (
+                            "Failed to connect to DB. MFA "
+                            "authentication failed: {"
+                            "host}:{port}. {message}"
+                        ).format(
+                            host=self._rest._host,
+                            port=self._rest._port,
+                            message=ret["message"],
+                        ),
+                        "errno": ER_FAILED_TO_CONNECT_TO_DB,
+                        "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                    },
+                )
+                return session_parameters  # required for unit test
+
+        elif ret["data"] and ret["data"].get("nextAction") == "PWD_CHANGE":
+            if callable(password_callback):
+                body = copy.deepcopy(body_template)
+                body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+                body["data"]["LOGIN_NAME"] = user
+                body["data"]["PASSWORD"] = (
+                    auth_instance.password
+                    if hasattr(auth_instance, "password")
+                    else None
+                )
+                body["data"]["CHOSEN_NEW_PASSWORD"] = password_callback()
+                # New Password input
+                ret = self._rest._post_request(
+                    url,
+                    headers,
+                    json.dumps(body),
+                    socket_timeout=auth_instance._socket_timeout,
+                )
+
+        logger.debug("completed authentication")
+        if not ret["success"]:
+            errno = ret.get("code", ER_FAILED_TO_CONNECT_TO_DB)
+            if errno == ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE:
+                # clear stored id_token if failed to connect because of id_token
+                # raise an exception for reauth without id_token
+                self._rest.id_token = None
+                self._delete_temporary_credential(
+                    self._rest._host, user, TokenType.ID_TOKEN
+                )
+                raise ReauthenticationRequest(
+                    ProgrammingError(
+                        msg=ret["message"],
+                        errno=int(errno),
+                        sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                    )
+                )
+            elif errno == OAUTH_ACCESS_TOKEN_EXPIRED_GS_CODE:
+                raise ReauthenticationRequest(
+                    ProgrammingError(
+                        msg=ret["message"],
+                        errno=int(errno),
+                        sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                    )
+                )
+
+            from . import AuthByKeyPair
+
+            if isinstance(auth_instance, AuthByKeyPair):
+                logger.debug(
+                    "JWT Token authentication failed. "
+                    "Token expires at: %s. "
+                    "Current Time: %s",
+                    str(auth_instance._jwt_token_exp),
+                    str(datetime.now(timezone.utc).replace(tzinfo=None)),
+                )
+            from . import AuthByUsrPwdMfa
+
+            if isinstance(auth_instance, AuthByUsrPwdMfa):
+                self._delete_temporary_credential(
+                    self._rest._host, user, TokenType.MFA_TOKEN
+                )
+            Error.errorhandler_wrapper(
+                self._rest._connection,
+                None,
+                DatabaseError,
+                {
+                    "msg": (
+                        "Failed to connect to DB: {host}:{port}. " "{message}"
+                    ).format(
+                        host=self._rest._host,
+                        port=self._rest._port,
+                        message=ret["message"],
+                    ),
+                    "errno": ER_FAILED_TO_CONNECT_TO_DB,
+                    "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                },
+            )
+        else:
+            logger.debug(
+                "token = %s",
+                (
+                    "******"
+                    if ret["data"] and ret["data"].get("token") is not None
+                    else "NULL"
+                ),
+            )
+            logger.debug(
+                "master_token = %s",
+                (
+                    "******"
+                    if ret["data"] and ret["data"].get("masterToken") is not None
+                    else "NULL"
+                ),
+            )
+            logger.debug(
+                "id_token = %s",
+                (
+                    "******"
+                    if ret["data"] and ret["data"].get("idToken") is not None
+                    else "NULL"
+                ),
+            )
+            logger.debug(
+                "mfa_token = %s",
+                (
+                    "******"
+                    if ret["data"] and ret["data"].get("mfaToken") is not None
+                    else "NULL"
+                ),
+            )
+            if not ret["data"]:
+                Error.errorhandler_wrapper(
+                    None,
+                    None,
+                    Error,
+                    {
+                        "msg": "There is no data in the returning response, please retry the operation."
+                    },
+                )
+            self._rest.update_tokens(
+                ret["data"].get("token"),
+                ret["data"].get("masterToken"),
+                master_validity_in_seconds=ret["data"].get("masterValidityInSeconds"),
+                id_token=ret["data"].get("idToken"),
+                mfa_token=ret["data"].get("mfaToken"),
+            )
+            self.write_temporary_credentials(
+                self._rest._host, user, session_parameters, ret
+            )
+            if ret["data"] and "sessionId" in ret["data"]:
+                self._rest._connection._session_id = ret["data"].get("sessionId")
+            if ret["data"] and "sessionInfo" in ret["data"]:
+                session_info = ret["data"].get("sessionInfo")
+                self._rest._connection._database = session_info.get("databaseName")
+                self._rest._connection._schema = session_info.get("schemaName")
+                self._rest._connection._warehouse = session_info.get("warehouseName")
+                self._rest._connection._role = session_info.get("roleName")
+            if ret["data"] and "parameters" in ret["data"]:
+                session_parameters.update(
+                    {p["name"]: p["value"] for p in ret["data"].get("parameters")}
+                )
+            self._rest._connection._update_parameters(session_parameters)
+            return session_parameters
+
+    def _read_temporary_credential(
+        self,
+        host: str,
+        user: str,
+        cred_type: TokenType,
+    ) -> str | None:
+        return self.get_token_cache().retrieve(TokenKey(host, user, cred_type))
+
+    def read_temporary_credentials(
+        self,
+        host: str,
+        user: str,
+        session_parameters: dict[str, Any],
+    ) -> None:
+        if session_parameters.get(PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL, False):
+            self._rest.id_token = self._read_temporary_credential(
+                host,
+                user,
+                TokenType.ID_TOKEN,
+            )
+
+        if session_parameters.get(PARAMETER_CLIENT_REQUEST_MFA_TOKEN, False):
+            self._rest.mfa_token = self._read_temporary_credential(
+                host,
+                user,
+                TokenType.MFA_TOKEN,
+            )
+
+    def _write_temporary_credential(
+        self,
+        host: str,
+        user: str,
+        cred_type: TokenType,
+        cred: str | None,
+    ) -> None:
+        if not cred:
+            logger.debug(
+                "no credential is given when try to store temporary credential"
+            )
+            return
+        self.get_token_cache().store(TokenKey(host, user, cred_type), cred)
+
+    def write_temporary_credentials(
+        self,
+        host: str,
+        user: str,
+        session_parameters: dict[str, Any],
+        response: dict[str, Any],
+    ) -> None:
+        if (
+            self._rest._connection.auth_class.consent_cache_id_token
+            and session_parameters.get(
+                PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL, False
+            )
+        ):
+            self._write_temporary_credential(
+                host, user, TokenType.ID_TOKEN, response["data"].get("idToken")
+            )
+
+        if session_parameters.get(PARAMETER_CLIENT_REQUEST_MFA_TOKEN, False):
+            self._write_temporary_credential(
+                host, user, TokenType.MFA_TOKEN, response["data"].get("mfaToken")
+            )
+
+    def _delete_temporary_credential(
+        self, host: str, user: str, cred_type: TokenType
+    ) -> None:
+        self.get_token_cache().remove(TokenKey(host, user, cred_type))
+
+    def get_token_cache(self) -> TokenCache:
+        if self._token_cache is None:
+            self._token_cache = TokenCache.make()
+        return self._token_cache
+
+
+def get_token_from_private_key(
+    user: str, account: str, privatekey_path: str, key_password: str | None
+) -> str:
+    encoded_password = key_password.encode() if key_password is not None else None
+    with open(privatekey_path, "rb") as key:
+        p_key = load_pem_private_key(
+            key.read(), password=encoded_password, backend=default_backend()
+        )
+
+    private_key = p_key.private_bytes(
+        encoding=Encoding.DER,
+        format=PrivateFormat.PKCS8,
+        encryption_algorithm=NoEncryption(),
+    )
+    from . import AuthByKeyPair
+
+    auth_instance = AuthByKeyPair(
+        private_key,
+        DAY_IN_SECONDS,
+    )  # token valid for 24 hours
+    return auth_instance.prepare(account=account, user=user)
+
+
+def get_public_key_fingerprint(private_key_file: str, password: str) -> str:
+    """Helper function to generate the public key fingerprint from the private key file"""
+    with open(private_key_file, "rb") as key:
+        p_key = load_pem_private_key(
+            key.read(), password=password.encode(), backend=default_backend()
+        )
+    private_key = p_key.private_bytes(
+        encoding=Encoding.DER,
+        format=PrivateFormat.PKCS8,
+        encryption_algorithm=NoEncryption(),
+    )
+    private_key = load_der_private_key(
+        data=private_key, password=None, backend=default_backend()
+    )
+    from . import AuthByKeyPair
+
+    return AuthByKeyPair.calculate_public_key_fingerprint(private_key)