service-catalog/charts/kubeflow/tests/gh-actions/test_dex_login.py

#!/usr/bin/env python3

import re
import time
from urllib.parse import urlsplit, urlencode
import requests
import urllib3


class DexSessionManager:
    """
    This is a version of the KFPClientManager() which only generates the Dex session cookies.
    See https://www.kubeflow.org/docs/components/pipelines/user-guides/core-functions/connect-api/#kubeflow-platform---outside-the-cluster
    """

    def __init__(
        self,
        endpoint_url: str,
        dex_username: str,
        dex_password: str,
        dex_auth_type: str = "local",
        skip_tls_verify: bool = False,
    ):
        """
        Initialize the DexSessionManager

        :param endpoint_url: the Kubeflow Endpoint URL
        :param skip_tls_verify: if True, skip TLS verification
        :param dex_username: the Dex username
        :param dex_password: the Dex password
        :param dex_auth_type: the auth type to use if Dex has multiple enabled, one of: ['ldap', 'local']
        """
        self._endpoint_url = endpoint_url
        self._skip_tls_verify = skip_tls_verify
        self._dex_username = dex_username
        self._dex_password = dex_password
        self._dex_auth_type = dex_auth_type
        self._client = None

        # disable SSL verification, if requested
        if self._skip_tls_verify:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

        # ensure `dex_default_auth_type` is valid
        if self._dex_auth_type not in ["ldap", "local"]:
            raise ValueError(
                f"Invalid `dex_auth_type` '{self._dex_auth_type}', must be one of: ['ldap', 'local']"
            )

    def get_session_cookies(self) -> str:
        """
        Get the session cookies by authenticating against Dex
        :return: a string of session cookies in the form "key1=value1; key2=value2"
        """
        max_retries = 3
        retry_delay = 2

        # Use a persistent session for cookies
        session = requests.Session()

        for attempt in range(max_retries):
            try:
                # GET the endpoint_url, which should redirect to Dex
                response = session.get(
                    self._endpoint_url,
                    allow_redirects=True,
                    verify=not self._skip_tls_verify
                )
                if response.status_code == 200:
                    pass
                elif response.status_code == 403:
                    # if we get 403, we might be at the oauth2-proxy sign-in page
                    # the default path to start the sign-in flow is `/oauth2/start?rd=<url>`
                    url_object = urlsplit(response.url)
                    url_object = url_object._replace(
                        path="/oauth2/start",
                        query=urlencode({"rd": url_object.path})
                    )
                    response = session.get(
                        url_object.geturl(),
                        allow_redirects=True,
                        verify=not self._skip_tls_verify
                    )
                else:
                    raise RuntimeError(
                        f"HTTP status code '{response.status_code}' for GET against: {self._endpoint_url}"
                    )

                # if we were NOT redirected, then the endpoint is unsecured
                if len(response.history) == 0:
                    # No cookies are needed
                    return ""

                # if we are at `../auth` path, we need to select an auth type
                url_object = urlsplit(response.url)
                if re.search(r"/auth$", url_object.path):
                    url_object = url_object._replace(
                        path=re.sub(r"/auth$", f"/auth/{self._dex_auth_type}", url_object.path)
                    )

                # if we are at `../auth/xxxx/login` path, then we are at the login page
                if re.search(r"/auth/.*/login$", url_object.path):
                    dex_login_url = url_object.geturl()
                else:
                    # otherwise, we need to follow a redirect to the login page
                    response = session.get(
                        url_object.geturl(),
                        allow_redirects=True,
                        verify=not self._skip_tls_verify
                    )
                    if response.status_code != 200:
                        raise RuntimeError(
                            f"HTTP status code '{response.status_code}' for GET against: {url_object.geturl()}"
                        )
                    dex_login_url = response.url

                # attempt Dex login
                response = session.post(
                    dex_login_url,
                    data={"login": self._dex_username, "password": self._dex_password},
                    allow_redirects=True,
                    verify=not self._skip_tls_verify,
                )

                # Handle 403 specifically - might need to restart oauth flow
                if response.status_code == 403:
                    # Try one more approach - go through the oauth2 flow again
                    oauth_url = f"{urlsplit(self._endpoint_url).scheme}://{urlsplit(self._endpoint_url).netloc}/oauth2/start"
                    response = session.get(
                        oauth_url,
                        allow_redirects=True,
                        verify=not self._skip_tls_verify,
                    )
                    # Continue with normal flow after restart
                    if response.status_code == 200 and session.cookies:
                        return "; ".join([f"{c.name}={c.value}" for c in session.cookies])

                if response.status_code != 200:
                    raise RuntimeError(
                        f"HTTP status code '{response.status_code}' for POST against: {dex_login_url}"
                    )

                # if we were NOT redirected, then the login credentials were probably invalid
                if len(response.history) == 0:
                    raise RuntimeError(
                        f"Login credentials are probably invalid - "
                        f"No redirect after POST to: {dex_login_url}"
                    )

                # if we are at `../approval` path, we need to approve the login
                url_object = urlsplit(response.url)
                if re.search(r"/approval$", url_object.path):
                    dex_approval_url = url_object.geturl()
                    # Approve the login
                    response = session.post(
                        dex_approval_url,
                        data={"approval": "approve"},
                        allow_redirects=True,
                        verify=not self._skip_tls_verify,
                    )
                    if response.status_code != 200:
                        raise RuntimeError(
                            f"HTTP status code '{response.status_code}' for POST against: {url_object.geturl()}"
                        )

                return "; ".join([f"{c.name}={c.value}" for c in session.cookies])

            except Exception as e:
                if attempt == max_retries - 1:  # Last attempt
                    print(f"All {max_retries} attempts failed. Last error: {str(e)}")
                    raise
                print(f"Attempt {attempt + 1} failed: {str(e)}")
                print(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
                retry_delay *= 2  # Exponential backoff


KUBEFLOW_ENDPOINT = "http://localhost:8080"
KUBEFLOW_USERNAME = "user@example.com"
KUBEFLOW_PASSWORD = "12341234"

# initialize a DexSessionManager
dex_session_manager = DexSessionManager(
    endpoint_url=KUBEFLOW_ENDPOINT,
    skip_tls_verify=True,
    dex_username=KUBEFLOW_USERNAME,
    dex_password=KUBEFLOW_PASSWORD,
    dex_auth_type="local",
)

# try to get the session cookies
# NOTE: this will raise an exception if something goes wrong
session_cookies = dex_session_manager.get_session_cookies()