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

snowflake/connector/auth/_http_server.py

@@ -0,0 +1,220 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import logging
+import os
+import select
+import socket
+import time
+import urllib.parse
+from collections.abc import Callable
+from types import TracebackType
+
+from typing_extensions import Self
+
+from ..compat import IS_WINDOWS
+
+logger = logging.getLogger(__name__)
+
+
+def _use_msg_dont_wait() -> bool:
+    if os.getenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", "false").lower() != "true":
+        return False
+    if IS_WINDOWS:
+        logger.warning(
+            "Configuration SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT is not available in Windows. Ignoring."
+        )
+        return False
+    return True
+
+
+def _wrap_socket_recv() -> Callable[[socket.socket, int], bytes]:
+    dont_wait = _use_msg_dont_wait()
+    if dont_wait:
+        # WSL containerized environment sometimes causes socket_client.recv to hang indefinetly
+        #   To avoid this, passing the socket.MSG_DONTWAIT flag which raises BlockingIOError if
+        #   operation would block
+        logger.debug(
+            "Will call socket.recv with MSG_DONTWAIT flag due to SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT env var"
+        )
+    socket_recv = (
+        (lambda sock, buf_size: socket.socket.recv(sock, buf_size, socket.MSG_DONTWAIT))
+        if dont_wait
+        else (lambda sock, buf_size: socket.socket.recv(sock, buf_size))
+    )
+
+    def socket_recv_checked(sock: socket.socket, buf_size: int) -> bytes:
+        raw = socket_recv(sock, buf_size)
+        # when running in a containerized environment, socket_client.recv occasionally returns an empty byte array
+        #   an immediate successive call to socket_client.recv gets the actual data
+        if len(raw) == 0:
+            raw = socket_recv(sock, buf_size)
+        return raw
+
+    return socket_recv_checked
+
+
+class AuthHttpServer:
+    """Simple HTTP server to receive callbacks through for auth purposes."""
+
+    DEFAULT_MAX_ATTEMPTS = 15
+    DEFAULT_TIMEOUT = 30.0
+
+    PORT_BIND_MAX_ATTEMPTS = 10
+    PORT_BIND_TIMEOUT = 20.0
+
+    def __init__(
+        self,
+        uri: str,
+        buf_size: int = 16384,
+    ) -> None:
+        parsed_uri = urllib.parse.urlparse(uri)
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.buf_size = buf_size
+        if os.getenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "False").lower() == "true":
+            if IS_WINDOWS:
+                logger.warning(
+                    "Configuration SNOWFLAKE_AUTH_SOCKET_REUSE_PORT is not available in Windows. Ignoring."
+                )
+            else:
+                self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+
+        port = parsed_uri.port or 0
+        for attempt in range(1, self.DEFAULT_MAX_ATTEMPTS + 1):
+            try:
+                self._socket.bind(
+                    (
+                        parsed_uri.hostname,
+                        port,
+                    )
+                )
+                break
+            except socket.gaierror as ex:
+                logger.error(
+                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                )
+                raise
+            except OSError as ex:
+                if attempt == self.DEFAULT_MAX_ATTEMPTS:
+                    logger.error(
+                        f"Failed to bind authorization callback server to port {port}: {ex}"
+                    )
+                    raise
+                logger.warning(
+                    f"Attempt {attempt}/{self.DEFAULT_MAX_ATTEMPTS}. "
+                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                )
+                time.sleep(self.PORT_BIND_TIMEOUT / self.PORT_BIND_MAX_ATTEMPTS)
+        try:
+            self._socket.listen(0)  # no backlog
+        except Exception as ex:
+            logger.error(f"Failed to start listening for auth callback: {ex}")
+            self.close()
+            raise
+        port = self._socket.getsockname()[1]
+        self._uri = urllib.parse.ParseResult(
+            scheme=parsed_uri.scheme,
+            netloc=parsed_uri.hostname + ":" + str(port),
+            path=parsed_uri.path,
+            params=parsed_uri.params,
+            query=parsed_uri.query,
+            fragment=parsed_uri.fragment,
+        )
+
+    @property
+    def url(self) -> str:
+        return self._uri.geturl()
+
+    @property
+    def port(self) -> int:
+        return self._uri.port
+
+    @property
+    def hostname(self) -> str:
+        return self._uri.hostname
+
+    def _try_poll(
+        self, attempts: int, attempt_timeout: float | None
+    ) -> (socket.socket | None, int):
+        for attempt in range(attempts):
+            read_sockets = select.select([self._socket], [], [], attempt_timeout)[0]
+            if read_sockets and read_sockets[0] is not None:
+                return self._socket.accept()[0], attempt
+        return None, attempts
+
+    def _try_receive_block(
+        self, client_socket: socket.socket, attempts: int, attempt_timeout: float | None
+    ) -> bytes | None:
+        if attempt_timeout is not None:
+            client_socket.settimeout(attempt_timeout)
+        recv = _wrap_socket_recv()
+        for attempt in range(attempts):
+            try:
+                return recv(client_socket, self.buf_size)
+            except BlockingIOError:
+                if attempt < attempts - 1:
+                    cooldown = min(attempt_timeout, 0.25) if attempt_timeout else 0.25
+                    logger.debug(
+                        f"BlockingIOError raised from socket.recv on {1 + attempt}/{attempts} attempt."
+                        f"Waiting for {cooldown} seconds before trying again"
+                    )
+                    time.sleep(cooldown)
+            except socket.timeout:
+                logger.debug(
+                    f"socket.recv timed out on {1 + attempt}/{attempts} attempt."
+                )
+        return None
+
+    def receive_block(
+        self,
+        max_attempts: int = None,
+        timeout: float | int | None = None,
+    ) -> (list[str] | None, socket.socket | None):
+        if max_attempts is None:
+            max_attempts = self.DEFAULT_MAX_ATTEMPTS
+        if timeout is None:
+            timeout = self.DEFAULT_TIMEOUT
+        """Receive a message with a maximum attempt count and a timeout in seconds, blocking."""
+        if not self._socket:
+            raise RuntimeError(
+                "Operation is not supported, server was already shut down."
+            )
+        attempt_timeout = timeout / max_attempts if timeout else None
+        client_socket, poll_attempts = self._try_poll(max_attempts, attempt_timeout)
+        if client_socket is None:
+            return None, None
+        raw_block = self._try_receive_block(
+            client_socket, max_attempts - poll_attempts, attempt_timeout
+        )
+        if raw_block:
+            return raw_block.decode("utf-8").split("\r\n"), client_socket
+        try:
+            client_socket.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+        client_socket.close()
+        return None, None
+
+    def close(self) -> None:
+        """Closes the underlying socket.
+        After having close() being called the server object cannot be reused.
+        """
+        if self._socket:
+            self._socket.close()
+            self._socket = None
+
+    def __enter__(self) -> Self:
+        """Context manager."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Context manager with disposing underlying networking objects."""
+        self.close()

snowflake/connector/auth/_oauth_base.py

@@ -0,0 +1,400 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import urllib.parse
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+from urllib.error import HTTPError, URLError
+
+from ..errorcode import (
+    ER_FAILED_TO_REQUEST,
+    ER_IDP_CONNECTION_ERROR,
+    ER_NO_CLIENT_ID,
+    ER_NO_CLIENT_SECRET,
+)
+from ..errors import Error, ProgrammingError
+from ..network import OAUTH_AUTHENTICATOR
+from ..secret_detector import SecretDetector
+from ..token_cache import TokenCache, TokenKey, TokenType
+from ..vendored import urllib3
+from .by_plugin import AuthByPlugin, AuthType
+
+if TYPE_CHECKING:
+    from .. import SnowflakeConnection
+
+logger = logging.getLogger(__name__)
+
+
+class _OAuthTokensMixin:
+    def __init__(
+        self,
+        token_cache: TokenCache | None,
+        refresh_token_enabled: bool,
+        idp_host: str,
+    ) -> None:
+        self._access_token = None
+        self._refresh_token_enabled = refresh_token_enabled
+        if self._refresh_token_enabled:
+            self._refresh_token = None
+        self._token_cache = token_cache
+        if self._token_cache:
+            logger.debug("token cache is going to be used if needed")
+            self._idp_host = idp_host
+            self._access_token_key: TokenKey | None = None
+            if self._refresh_token_enabled:
+                self._refresh_token_key: TokenKey | None = None
+
+    def _update_cache_keys(self, user: str) -> None:
+        if self._token_cache:
+            self._user = user
+
+    def _get_access_token_cache_key(self) -> TokenKey | None:
+        return (
+            TokenKey(self._user, self._idp_host, TokenType.OAUTH_ACCESS_TOKEN)
+            if self._token_cache and self._user
+            else None
+        )
+
+    def _get_refresh_token_cache_key(self) -> TokenKey | None:
+        return (
+            TokenKey(self._user, self._idp_host, TokenType.OAUTH_REFRESH_TOKEN)
+            if self._refresh_token_enabled and self._token_cache and self._user
+            else None
+        )
+
+    def _pop_cached_token(self, key: TokenKey | None) -> str | None:
+        if self._token_cache is None or key is None:
+            return None
+        return self._token_cache.retrieve(key)
+
+    def _pop_cached_access_token(self) -> bool:
+        """Retrieves OAuth access token from the token cache if enabled"""
+        self._access_token = self._pop_cached_token(self._get_access_token_cache_key())
+        return self._access_token is not None
+
+    def _pop_cached_refresh_token(self) -> bool:
+        """Retrieves OAuth refresh token from the token cache if enabled"""
+        if self._refresh_token_enabled:
+            self._refresh_token = self._pop_cached_token(
+                self._get_refresh_token_cache_key()
+            )
+            return self._refresh_token is not None
+        return False
+
+    def _reset_cached_token(self, key: TokenKey | None, token: str | None) -> None:
+        if self._token_cache is None or key is None:
+            return
+        if token:
+            self._token_cache.store(key, token)
+        else:
+            self._token_cache.remove(key)
+
+    def _reset_access_token(self, access_token: str | None = None) -> None:
+        """Updates OAuth access token both in memory and in the token cache if enabled"""
+        logger.debug(
+            "resetting access token to %s",
+            "*" * len(access_token) if access_token else None,
+        )
+        self._access_token = access_token
+        self._reset_cached_token(self._get_access_token_cache_key(), self._access_token)
+
+    def _reset_refresh_token(self, refresh_token: str | None = None) -> None:
+        """Updates OAuth refresh token both in memory and in the token cache if necessary"""
+        if self._refresh_token_enabled:
+            logger.debug(
+                "resetting refresh token to %s",
+                "*" * len(refresh_token) if refresh_token else None,
+            )
+            self._refresh_token = refresh_token
+            self._reset_cached_token(
+                self._get_refresh_token_cache_key(), self._refresh_token
+            )
+
+    def _reset_temporary_state(self) -> None:
+        self._access_token = None
+        if self._refresh_token_enabled:
+            self._refresh_token = None
+        if self._token_cache:
+            self._user = None
+
+
+class AuthByOAuthBase(AuthByPlugin, _OAuthTokensMixin, ABC):
+    """A base abstract class for OAuth authenticators"""
+
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: str,
+        token_request_url: str,
+        scope: str,
+        token_cache: TokenCache | None,
+        refresh_token_enabled: bool,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        _OAuthTokensMixin.__init__(
+            self,
+            token_cache=token_cache,
+            refresh_token_enabled=refresh_token_enabled,
+            idp_host=urllib.parse.urlparse(token_request_url).hostname,
+        )
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._token_request_url = token_request_url
+        self._scope = scope
+        if refresh_token_enabled:
+            logger.debug("oauth refresh token is going to be used if needed")
+            self._scope += (" " if self._scope else "") + "offline_access"
+
+    @abstractmethod
+    def _request_tokens(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        password: str | None,
+        **kwargs: Any,
+    ) -> (str | None, str | None):
+        """Request new access and optionally refresh tokens from IdP.
+
+        This function should implement specific tokens querying flow.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_oauth_type_id(self) -> str:
+        """Get OAuth specific authenticator id to be passed to Snowflake.
+
+        This function should return a unique OAuth authenticator id.
+        """
+        raise NotImplementedError
+
+    def reset_secrets(self) -> None:
+        logger.debug("resetting secrets")
+        self._reset_temporary_state()
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.OAUTH
+
+    @property
+    def assertion_content(self) -> str:
+        """Returns the token."""
+        return self._access_token or ""
+
+    @staticmethod
+    def _validate_client_credentials_present(
+        client_id: str, client_secret: str, connection: SnowflakeConnection
+    ) -> tuple[str, str]:
+        if client_id is None or client_id == "":
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_id' is empty",
+                    "errno": ER_NO_CLIENT_ID,
+                },
+            )
+        if client_secret is None or client_secret == "":
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_secret' is empty",
+                    "errno": ER_NO_CLIENT_SECRET,
+                },
+            )
+
+        return client_id, client_secret
+
+    def reauthenticate(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        **kwargs: Any,
+    ) -> dict[str, bool]:
+        self._reset_access_token()
+        if self._pop_cached_refresh_token():
+            logger.debug(
+                "OAuth refresh token is available, try to use it and get a new access token"
+            )
+            self._do_refresh_token(conn=conn)
+        conn.authenticate_with_retry(self)
+        return {"success": True}
+
+    def prepare(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        **kwargs: Any,
+    ) -> None:
+        """Web Browser based Authentication."""
+        logger.debug("authenticating with OAuth authorization code flow")
+        self._update_cache_keys(user=user)
+        if self._pop_cached_access_token():
+            logger.info(
+                "OAuth access token is already available in cache, no need to authenticate."
+            )
+            return
+        access_token, refresh_token = self._request_tokens(
+            conn=conn,
+            authenticator=authenticator,
+            service_name=service_name,
+            account=account,
+            user=user,
+            **kwargs,
+        )
+        self._reset_access_token(access_token)
+        self._reset_refresh_token(refresh_token)
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Used by Auth to update the request that gets sent to /v1/login-request.
+
+        Args:
+            body: existing request dictionary
+        """
+        body["data"]["AUTHENTICATOR"] = OAUTH_AUTHENTICATOR
+        body["data"]["TOKEN"] = self._access_token
+        body["data"]["OAUTH_TYPE"] = self._get_oauth_type_id()
+
+    def _do_refresh_token(self, conn: SnowflakeConnection) -> None:
+        """If a refresh token is available exchanges it with a new access token.
+        Updates self as a side-effect. Needs at lest self._refresh_token and client_id set.
+        """
+        if not self._refresh_token_enabled:
+            logger.debug("refresh_token feature is disabled")
+            return
+
+        resp = self._get_refresh_token_response(conn)
+        if not resp:
+            logger.info(
+                "failed to exchange the refresh token on a new OAuth access token"
+            )
+            self._reset_refresh_token()
+            return
+
+        try:
+            json_resp = json.loads(resp.data.decode())
+            self._reset_access_token(json_resp["access_token"])
+            if "refresh_token" in json_resp:
+                self._reset_refresh_token(json_resp["refresh_token"])
+        except (
+            json.JSONDecodeError,
+            KeyError,
+        ):
+            logger.error(
+                "refresh token exchange response did not contain 'access_token'"
+            )
+            logger.debug(
+                "received the following response body when exchanging refresh token: %s",
+                SecretDetector.mask_secrets(str(resp.data)),
+            )
+            self._reset_refresh_token()
+
+    def _get_refresh_token_response(
+        self, conn: SnowflakeConnection
+    ) -> urllib3.BaseHTTPResponse | None:
+        fields = {
+            "grant_type": "refresh_token",
+            "refresh_token": self._refresh_token,
+        }
+        if self._scope:
+            fields["scope"] = self._scope
+        try:
+            return urllib3.PoolManager().request_encode_body(
+                # TODO: use network pool to gain use of proxy settings and so on
+                "POST",
+                self._token_request_url,
+                encode_multipart=False,
+                headers=self._create_token_request_headers(),
+                fields=fields,
+            )
+        except HTTPError as e:
+            self._handle_failure(
+                conn=conn,
+                ret={
+                    "code": ER_FAILED_TO_REQUEST,
+                    "message": f"Failed to request new OAuth access token with a refresh token,"
+                    f" url={e.url}, code={e.code}, reason={e.reason}",
+                },
+            )
+        except URLError as e:
+            self._handle_failure(
+                conn=conn,
+                ret={
+                    "code": ER_FAILED_TO_REQUEST,
+                    "message": f"Failed to request new OAuth access token with a refresh token, reason: {e.reason}",
+                },
+            )
+        except Exception:
+            self._handle_failure(
+                conn=conn,
+                ret={
+                    "code": ER_FAILED_TO_REQUEST,
+                    "message": "Failed to request new OAuth access token with a refresh token by unknown reason",
+                },
+            )
+        return None
+
+    def _get_request_token_response(
+        self,
+        connection: SnowflakeConnection,
+        fields: dict[str, str],
+    ) -> (str | None, str | None):
+        resp = urllib3.PoolManager().request_encode_body(
+            # TODO: use network pool to gain use of proxy settings and so on
+            "POST",
+            self._token_request_url,
+            headers=self._create_token_request_headers(),
+            encode_multipart=False,
+            fields=fields,
+        )
+        try:
+            logger.debug("OAuth IdP response received, try to parse it")
+            json_resp: dict = json.loads(resp.data)
+            access_token = json_resp["access_token"]
+            refresh_token = json_resp.get("refresh_token")
+            return access_token, refresh_token
+        except (
+            json.JSONDecodeError,
+            KeyError,
+        ):
+            logger.error("oauth response invalid, does not contain 'access_token'")
+            logger.debug(
+                "received the following response body when requesting oauth token: %s",
+                SecretDetector.mask_secrets(str(resp.data)),
+            )
+            self._handle_failure(
+                conn=connection,
+                ret={
+                    "code": ER_IDP_CONNECTION_ERROR,
+                    "message": "Invalid HTTP request from web browser. Idp "
+                    "authentication could have failed.",
+                },
+            )
+        return None, None
+
+    def _create_token_request_headers(self) -> dict[str, str]:
+        return {
+            "Authorization": "Basic "
+            + base64.b64encode(
+                f"{self._client_id}:{self._client_secret}".encode()
+            ).decode(),
+            "Accept": "application/json",
+            "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
+        }