doczyai-pipelines/streamlit/security.py

import streamlit as st
import msal
import requests
import boto3
from botocore.exceptions import ClientError
import json

# Replace with your own values
CLIENT_ID = "effafe90-7ed7-43a3-ab03-19a0be2f1758"
# CLIENT_SECRET = 'bjQ8Q~lpR2uBcGI34VDu16t73doz8Crj0YY_~dgD'
# TENANT_ID = ''

AUTHORITY = "https://login.microsoftonline.com/organizations/"
SCOPE = ["User.Read"]
REDIRECT_URI = "https://172.29.20.102:8501"

# Initialize boto3 client to interact with AWS Secrets Manager


# URL Masking
def clear_url():
    js_code = """
        window.history.replaceState({}, document.title, window.location.pathname);
    """
    st.components.v1.html(f"<script>{js_code}</script>", height=0, width=0)


def get_secret():

    secret_name = "doczy-sso-azure-app-key"
    region_name = "us-east-2"

    # Create a Secrets Manager client
    # session = boto3.session.Session()
    # client = session.client(
    #     service_name='secretsmanager',
    #     region_name=region_name
    # )
    client = boto3.client("secretsmanager", region_name=region_name)
    try:
        get_secret_value_response = client.get_secret_value(SecretId=secret_name)
    except ClientError as e:
        raise e

    secret = get_secret_value_response["SecretString"]
    secret = json.loads(secret)["CLIENT_SECRET"]

    return secret


CLIENT_SECRET = get_secret()

app = msal.ConfidentialClientApplication(
    CLIENT_ID, authority=AUTHORITY, client_credential=CLIENT_SECRET
)


def get_auth_url(REDIRECT_URI):
    auth_url = app.get_authorization_request_url(SCOPE, redirect_uri=REDIRECT_URI)
    return auth_url


@st.cache_data
def get_token_from_code(auth_code, REDIRECT_URI):
    app = msal.ConfidentialClientApplication(
        CLIENT_ID, authority=AUTHORITY, client_credential=CLIENT_SECRET
    )
    result = app.acquire_token_by_authorization_code(
        auth_code, scopes=SCOPE, redirect_uri=REDIRECT_URI
    )
    return result["access_token"]


def get_user_info(access_token):
    headers = {"Authorization": f"Bearer {access_token}"}
    response = requests.get("https://graph.microsoft.com/v1.0/me", headers=headers)
    return response.json()


def handle_redirect(REDIRECT_URI):
    if not st.session_state.get("access_token"):
        code = st.query_params.get("code")
        if code:
            access_token = get_token_from_code(code, REDIRECT_URI)
            st.session_state["access_token"] = access_token
            st.session_state


# URL Masking Pt.2
if "access_token" in st.session_state:
    clear_url()