Files
doczyai-pipelines/streamlit/security.py
T
2024-04-15 14:24:59 -05:00

73 lines
2.0 KiB
Python

import streamlit as st
import msal
import requests
import boto3
from botocore.exceptions import ClientError
# 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://171.29.20.126:8501'
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
)
try:
get_secret_value_response = client.get_secret_value(
SecretId=secret_name
)
except ClientError as e:
raise e
secret = get_secret_value_response['SecretString']
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