API Reference

def init_auth(
    user_roles: dict = None,
    tenant_id: str = None,
    client_id: str = None,
    email_suffix: str = None,
    init_obo_process: bool = False,
    client_secret: str = None,
    downstream_scope=None,
    retry_times: int = 5,
) -> None:
    """
    Initializes the authentication process for streamlit applications. This
    function also supports obo (on-behalf-of) token acquisition process, which
    can be enabled by setting the `init_obo_process=True`. Note that
    `client_secret`, `downstream_scope`, and `retry_times` are only used when
    `init_obo_process=True`. The sign-in button is by default on the st.sidebar

    This function initializes the authentication process by checking the user's
    account and role. If the user's account is not valid or does not have the
    required role, the function stops and displays an error message. If the
    user's account is valid and has the required role, the function sets the
    necessary information to the session state. The user sign information is
    stored in `st.session_state.auth_data`, the user name is stored in
    `st.session_state.username`, and the user roles are stored in
    `st.session_state.roles`.

    The obo process can be triggered by setting the `init_obo_process=True`.
    The obo token information is stored in `st.session_state.obo_info`, this
    is a dictionary containing the `access_token`, `refresh_token`, and the
    `expires_at`. The `access_token` will be stored in system environment
    variables as `OBO_TOKEN`, and also `st.session_state.obo_token` for easy
    access.

    Args:
        user_roles (dict): A dictionary containing the required roles. The keys
            are the role names and the values are the role descriptions. The
            default value is None. This corresponds to the user roles defined
            in the Azure App registration. The value of each key should be the
            value defined in the Azure App registration.
        tenant_id (str): The tenant ID in Microsoft Azure. If the `TENANT_ID`
            is already set in the environment variables, this argument is not
            required. If no `TENANT_ID` is found in the environment variables,
            and no passed value is provided, streamlit will stop.
        client_id (str): The client ID from Microsoft Azure App registration.
            If the `CLIENT_ID` is already set in the environment variables,
            this argument is not required. If no `CLIENT_ID` is found in the
            environment variables, and no passed value is provided, streamlit
            will stop.
        email_suffix (str): The email suffix to match with the enterprise email
            suffix. The default value is None. If it is None, the function will
            not check the email suffix. The email suffix such as
            `@microsoft.com` is used to check whether the login user's email
            has valid enterprise email.
        init_obo_process (bool): A boolean value to enable the obo token
            acquisition process. The default value is False. If set to True,
            the function will acquire the obo token.
        client_secret (str): The client secret for the `cliend_id` above. If
            the `CLIENT_SECRET` is already set in the environment variables,
            this argument is not required. If no `CLIENT_SECRET` is found in
            the environment variables, and no passed value is provided, error
            will be raised.
        downstream_scope (str): The scope for the obo token acquisition
            process. This is required when `init_obo_process=True`. Note that
            the scope here is the downstream API scope, not the scope of the
            provided `client_id` above.
        retry_times (int): The number of times to retry the token acquisition
            process. The default value is 5.

    Examples:
        >>> user_roles = {
        ...     "Admin": "Admin role",
        ...     "User": "User role",
        ... }
        >>> tenant_id = "your_tenant_id"
        >>> client_id = "your_client_id"
        >>> email_suffix = "@microsoft.com"
        >>> init_auth(user_roles, tenant_id, client_id, email_suffix)
    """
    client_id = client_id or os.getenv("CLIENT_ID", None)
    tenant_id = tenant_id or os.getenv("TENANT_ID", None)

    if tenant_id is None or client_id is None:
        st.warning(
            "Tenant ID and client ID cannot be None! "
            f"Your input: tenant_id={tenant_id}, client_id={client_id}. \n"
            "Please set the TENANT_ID environment variable"
        )
        st.stop()

    with st.sidebar:
        auth_data = Msal.initialize_ui(
            client_id=client_id,
            authority=f"https://login.microsoftonline.com/{tenant_id}",
            scopes=[],
            connecting_label="Connecting",
            disconnected_label="Disconnected",
            sign_in_label="Sign in",
            sign_out_label="Sign out",
        )

        if not auth_data:
            # not valid enterprise account
            st.warning("Invalid account! No permission found for you")
            st.stop()

        if email_suffix is not None:
            # in case to match with enterprise email suffix
            # e.g., @microsoft.com
            if email_suffix not in auth_data["account"]["username"]:
                st.warning("Invalid account! No permission found for you")
                st.stop()

        if _check_role(auth_data["idTokenClaims"], user_roles) is False:
            # check the login user has the required roles defined
            # in the Microsoft Azure App registration
            st.warning(
                f"Hello, {auth_data['account']['name']}, "
                "No permission found for you"
            )
            st.stop()

        if "auth_data" not in st.session_state:
            # set the names to the session state
            st.session_state.auth_data = auth_data
            st.session_state.username = auth_data["account"]["name"]

            st.session_state.roles = auth_data["idTokenClaims"]["roles"]

        if init_obo_process:
            obo_token = _acquire_access_token_obo(
                auth_data["idToken"],
                tenant_id=tenant_id,
                client_id=client_id,
                client_secret=client_secret,
                downstream_scope=downstream_scope,
                retry_times=retry_times,
            )
            if not obo_token:
                st.write(
                    "Failed acquiring token, refresh the page and login again!"
                )
                st.stop()
            st.session_state.obo_info = {
                "access_token": obo_token.get("access_token"),
                "refresh_token": obo_token.get("refresh_token"),
                "expires_at": datetime.datetime.now()
                + datetime.timedelta(seconds=obo_token.get("expires_in")),
            }
            os.environ["OBO_TOKEN"] = obo_token.get("access_token")
            st.session_state.obo_token = obo_token.get("access_token")

def refresh_obo_token(
    tenant_id: str = None,
    client_id: str = None,
    client_secret: str = None,
    downstream_scope=None,
) -> str:
    """
    Checks if the user is logged in and refreshes the obo access token if
    necessary.

    This function will store the full obo token information in the
    `st.session_state.obo_info` and the current access token in the
    `st.session_state.obo_token`. This function can be only run after
    the `init_auth` function is called, and the `init_obo_process` is set to
    `True`, in the background, it uses the `st.session_state.obo_info` in
    user sign in process.

    Returns:
        str: The refreshed access token.

    Raises:
        None
    """
    tokens = st.session_state.obo_info

    if not tokens:
        st.write("User not logged in")
        return None

    if tokens.get("expires_at") < datetime.datetime.now():
        logger.info("Refreshing access token...")

        new_tokens = _refresh_access_token(
            tenant_id=tenant_id,
            client_id=client_id,
            client_secret=client_secret,
            downstream_scope=downstream_scope,
            refresh_token=tokens.get("refresh_token"),
        )
        tokens["access_token"] = new_tokens.get("access_token")
        tokens["refresh_token"] = new_tokens.get("refresh_token")
        tokens["expires_at"] = datetime.datetime.now() + datetime.timedelta(
            seconds=new_tokens.get("expires_in")
        )
        st.session_state.obo_info = tokens
        os.environ["OBO_TOKEN"] = tokens["access_token"]
        st.session_state.obo_token = tokens["access_token"]
    return st.session_state.obo_info