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

snowflake/connector/auth/okta.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import json
+import logging
+import time
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable
+
+from ..compat import unescape, urlencode, urlsplit
+from ..constants import (
+    HTTP_HEADER_ACCEPT,
+    HTTP_HEADER_CONTENT_TYPE,
+    HTTP_HEADER_SERVICE_NAME,
+    HTTP_HEADER_USER_AGENT,
+)
+from ..errorcode import ER_IDP_CONNECTION_ERROR, ER_INCORRECT_DESTINATION
+from ..errors import DatabaseError, Error, RefreshTokenError
+from ..network import CONTENT_TYPE_APPLICATION_JSON, PYTHON_CONNECTOR_USER_AGENT
+from ..sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
+from . import Auth
+from .by_plugin import AuthByPlugin, AuthType
+
+if TYPE_CHECKING:
+    from .. import SnowflakeConnection
+
+logger = logging.getLogger(__name__)
+
+
+def _is_prefix_equal(url1, url2):
+    """Checks if URL prefixes are identical.
+
+    The scheme, hostname and port number are compared. If the port number is not specified and the scheme is https,
+    the port number is assumed to be 443.
+    """
+    parsed_url1 = urlsplit(url1)
+    parsed_url2 = urlsplit(url2)
+
+    port1 = parsed_url1.port
+    if not port1 and parsed_url1.scheme == "https":
+        port1 = "443"
+    port2 = parsed_url1.port
+    if not port2 and parsed_url2.scheme == "https":
+        port2 = "443"
+
+    return (
+        parsed_url1.hostname == parsed_url2.hostname
+        and port1 == port2
+        and parsed_url1.scheme == parsed_url2.scheme
+    )
+
+
+def _get_post_back_url_from_html(html):
+    """Gets the post back URL.
+
+    Since the HTML is not well-formed, minidom cannot be used to convert to
+    DOM. The first discovered form is assumed to be the form to post back
+    and the URL is taken from action attributes.
+    """
+    logger.debug(html)
+
+    idx = html.find("<form")
+    start_idx = html.find('action="', idx)
+    end_idx = html.find('"', start_idx + 8)
+    return unescape(html[start_idx + 8 : end_idx])
+
+
+class AuthByOkta(AuthByPlugin):
+    """Authenticate user by OKTA."""
+
+    def __init__(self, application: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._saml_response = None
+        self._application = application
+
+    def reset_secrets(self) -> None:
+        pass
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.OKTA
+
+    @property
+    def assertion_content(self) -> str:
+        return self._saml_response
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        body["data"]["RAW_SAML_RESPONSE"] = self._saml_response
+
+    def prepare(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        password: str,
+        **kwargs: Any,
+    ) -> None:
+        """SAML Authentication.
+
+        Steps are:
+        1.  query GS to obtain IDP token and SSO url
+        2.  IMPORTANT Client side validation:
+            validate both token url and sso url contains same prefix
+            (protocol + host + port) as the given authenticator url.
+            Explanation:
+            This provides a way for the user to 'authenticate' the IDP it is
+            sending his/her credentials to.  Without such a check, the user could
+            be coerced to provide credentials to an IDP impersonator.
+        3.  query IDP token url to authenticate and retrieve access token
+        4.  given access token, query IDP URL snowflake app to get SAML response
+        5.  IMPORTANT Client side validation:
+            validate the post back url come back with the SAML response
+            contains the same prefix as the Snowflake's server url, which is the
+            intended destination url to Snowflake.
+        Explanation:
+            This emulates the behavior of IDP initiated login flow in the user
+            browser where the IDP instructs the browser to POST the SAML
+            assertion to the specific SP endpoint.  This is critical in
+            preventing a SAML assertion issued to one SP from being sent to
+            another SP.
+        """
+        logger.debug("authenticating by SAML")
+        headers, sso_url, token_url = self._step1(
+            conn,
+            authenticator,
+            service_name,
+            account,
+            user,
+        )
+        self._step2(conn, authenticator, sso_url, token_url)
+        response_html = self._step4(
+            conn,
+            partial(self._step3, conn, headers, token_url, user, password),
+            sso_url,
+        )
+        self._step5(conn, response_html)
+
+    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return {"success": False}
+
+    def _step1(
+        self,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+    ) -> tuple[dict[str, str], str, str]:
+        logger.debug("step 1: query GS to obtain IDP token and SSO url")
+
+        headers = {
+            HTTP_HEADER_CONTENT_TYPE: CONTENT_TYPE_APPLICATION_JSON,
+            HTTP_HEADER_ACCEPT: CONTENT_TYPE_APPLICATION_JSON,
+            HTTP_HEADER_USER_AGENT: PYTHON_CONNECTOR_USER_AGENT,
+        }
+        if service_name:
+            headers[HTTP_HEADER_SERVICE_NAME] = service_name
+        url = "/session/authenticator-request"
+        body = Auth.base_auth_data(
+            user,
+            account,
+            conn.application,
+            conn._internal_application_name,
+            conn._internal_application_version,
+            conn._ocsp_mode(),
+            conn.login_timeout,
+            conn._network_timeout,
+        )
+
+        body["data"]["AUTHENTICATOR"] = authenticator
+        logger.debug(
+            "account=%s, authenticator=%s",
+            account,
+            authenticator,
+        )
+        ret = conn._rest._post_request(
+            url,
+            headers,
+            json.dumps(body),
+            timeout=conn._rest._connection.login_timeout,
+            socket_timeout=conn._rest._connection.login_timeout,
+        )
+
+        if not ret["success"]:
+            self._handle_failure(conn=conn, ret=ret)
+
+        data = ret["data"]
+        token_url = data["tokenUrl"]
+        sso_url = data["ssoUrl"]
+        return headers, sso_url, token_url
+
+    def _step2(
+        self,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        sso_url: str,
+        token_url: str,
+    ) -> None:
+        logger.debug(
+            "step 2: validate Token and SSO URL has the same prefix as authenticator"
+        )
+        if not _is_prefix_equal(authenticator, token_url) or not _is_prefix_equal(
+            authenticator, sso_url
+        ):
+            Error.errorhandler_wrapper(
+                conn._rest._connection,
+                None,
+                DatabaseError,
+                {
+                    "msg": (
+                        "The specified authenticator is not supported: "
+                        f"{authenticator}, token_url: {token_url}, "
+                        f"sso_url: {sso_url}"
+                    ),
+                    "errno": ER_IDP_CONNECTION_ERROR,
+                    "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                },
+            )
+
+    @staticmethod
+    def _step3(
+        conn: SnowflakeConnection,
+        headers: dict[str, str],
+        token_url: str,
+        user: str,
+        password: str,
+    ) -> str:
+        logger.debug(
+            "step 3: query IDP token url to authenticate and " "retrieve access token"
+        )
+        data = {
+            "username": user,
+            "password": password,
+        }
+        ret = conn._rest.fetch(
+            "post",
+            token_url,
+            headers,
+            data=json.dumps(data),
+            timeout=conn._rest._connection.login_timeout,
+            socket_timeout=conn._rest._connection.login_timeout,
+            catch_okta_unauthorized_error=True,
+        )
+        one_time_token = ret.get("sessionToken", ret.get("cookieToken"))
+        if not one_time_token:
+            Error.errorhandler_wrapper(
+                conn._rest._connection,
+                None,
+                DatabaseError,
+                {
+                    "msg": (
+                        "The authentication failed for {user} "
+                        "by {token_url}.".format(
+                            token_url=token_url,
+                            user=user,
+                        )
+                    ),
+                    "errno": ER_IDP_CONNECTION_ERROR,
+                    "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                },
+            )
+        return one_time_token
+
+    @staticmethod
+    def _step4(
+        conn: SnowflakeConnection,
+        generate_one_time_token: Callable,
+        sso_url: str,
+    ) -> dict[Any, Any]:
+        logger.debug("step 4: query IDP URL snowflake app to get SAML " "response")
+        timeout_time = time.time() + conn.login_timeout if conn.login_timeout else None
+        response_html = {}
+        origin_sso_url = sso_url
+        while timeout_time is None or time.time() < timeout_time:
+            try:
+                url_parameters = {
+                    "RelayState": "/some/deep/link",
+                    "onetimetoken": generate_one_time_token(),
+                }
+                sso_url = origin_sso_url + "?" + urlencode(url_parameters)
+                headers = {
+                    HTTP_HEADER_ACCEPT: "*/*",
+                }
+                remaining_timeout = timeout_time - time.time() if timeout_time else None
+                response_html = conn._rest.fetch(
+                    "get",
+                    sso_url,
+                    headers,
+                    timeout=remaining_timeout,
+                    socket_timeout=remaining_timeout,
+                    is_raw_text=True,
+                    is_okta_authentication=True,
+                )
+                break
+            except RefreshTokenError:
+                logger.debug("step4: refresh token for re-authentication")
+        return response_html
+
+    def _step5(
+        self,
+        conn: SnowflakeConnection,
+        response_html: str,
+    ) -> None:
+        logger.debug("step 5: validate post_back_url matches Snowflake URL")
+        post_back_url = _get_post_back_url_from_html(response_html)
+        full_url = "{protocol}://{host}:{port}".format(
+            protocol=conn._rest._protocol,
+            host=conn._rest._host,
+            port=conn._rest._port,
+        )
+        if not getattr(conn, "_disable_saml_url_check", False) and not _is_prefix_equal(
+            post_back_url, full_url
+        ):
+            Error.errorhandler_wrapper(
+                conn._rest._connection,
+                None,
+                DatabaseError,
+                {
+                    "msg": (
+                        "The specified authenticator and destination "
+                        "URL in the SAML assertion do not match: "
+                        "expected: {url}, "
+                        "post back: {post_back_url}".format(
+                            url=full_url,
+                            post_back_url=post_back_url,
+                        )
+                    ),
+                    "errno": ER_INCORRECT_DESTINATION,
+                    "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                },
+            )
+        self._saml_response = response_html

snowflake/connector/auth/pat.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import typing
+
+from snowflake.connector.network import PROGRAMMATIC_ACCESS_TOKEN
+
+from .by_plugin import AuthByPlugin, AuthType
+
+
+class AuthByPAT(AuthByPlugin):
+
+    def __init__(self, pat_token: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._pat_token: str | None = pat_token
+
+    def type_(self) -> AuthType:
+        return AuthType.PAT
+
+    def reset_secrets(self) -> None:
+        self._pat_token = None
+
+    def update_body(self, body: dict[typing.Any, typing.Any]) -> None:
+        body["data"]["AUTHENTICATOR"] = PROGRAMMATIC_ACCESS_TOKEN
+        body["data"]["TOKEN"] = self._pat_token
+
+    def prepare(
+        self,
+        **kwargs: typing.Any,
+    ) -> None:
+        """Nothing to do here, token should be obtained outside the driver."""
+        pass
+
+    def reauthenticate(self, **kwargs: typing.Any) -> dict[str, bool]:
+        return {"success": False}
+
+    @property
+    def assertion_content(self) -> str | None:
+        """Returns the token."""
+        return self._pat_token

snowflake/connector/auth/usrpwdmfa.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from ..errorcode import ER_NO_PASSWORD
+from ..errors import ProgrammingError
+from .by_plugin import AuthByPlugin, AuthType
+
+if TYPE_CHECKING:
+    from .. import SnowflakeConnection
+
+logger = logging.getLogger(__name__)
+
+MFA_TOKEN = "MFATOKEN"
+
+
+class AuthByUsrPwdMfa(AuthByPlugin):
+    """Username & password & mfa authenticator."""
+
+    @property
+    def assertion_content(self) -> str:
+        return "*********"
+
+    def __init__(
+        self,
+        password: str,
+        mfa_token: str | None = None,
+        **kwargs,
+    ) -> None:
+        """Initializes and instance with a password and a mfa token."""
+        super().__init__(**kwargs)
+        self._password: str | None = password
+        self._mfa_token: str | None = mfa_token
+
+    def reset_secrets(self) -> None:
+        self._password = None
+        self._mfa_token = None
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.USR_PWD_MFA
+
+    def prepare(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        **kwargs: Any,
+    ) -> None:
+        if conn._rest and conn._rest.mfa_token:
+            self._mfa_token = conn._rest.mfa_token
+
+    def reauthenticate(self, **kwargs) -> dict[str, bool]:
+        return {"success": False}
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Sets the password and mfa_token if available.
+
+        Don't set body['data']['AUTHENTICATOR'], since this is still snowflake default authenticator.
+        """
+        if not self._password:
+            raise ProgrammingError(
+                msg="Password for username password authenticator is empty.",
+                errno=ER_NO_PASSWORD,
+            )
+        body["data"]["PASSWORD"] = self._password
+        if self._mfa_token:
+            body["data"]["TOKEN"] = self._mfa_token