Hi, I implemented a custom discord oauth flow (works similarly with any sso) using mostly only python and js only to create a custom login button.
This is working in a cloud deployed app (streamlig 1.37) using GCP and it does not require additional packages. It should be compatible with older Streamlit versions as well by removing the newest decorators.
The logic is as follows:
- a login/logout button in rendered on depending on the absence/presence of user data
- if the user clicks on the login button, the auth flow is started and the user is redirected to the discord login page
- once the login completes, the user is reditected to the oauth2-callback page of the streamlit app where the auth code is exchanged for an access token
- when the operation is completed, the user is redirected to the homepage
- when the user is logged in, user data is filled and the login button is transformed into a div with the user image and logout button. You may want to improve the html code for this, as a base i would suggest to add the following into the index.html head:
.profileImage {
width: 50px;
height: 50px;
border-radius: 50%;
}
Of course you can customize the redirect link to redirect the user wherever you want, maybe based on where he/she was before loggin in. For more information check the auth flow guide of your SSO provider.
The code is structured as follows:
app.py
pages/oauth2-callback.py
src/utils/frontend.py
src/utils/discord.py
Helper functions for frontend (frontend.py)
def create_link_button(key: str, url: str, label: str | None = None, icon: str | None = None) -> None:
st.html(
f'''
<style>
.{key} {{
display: inline-flex;
align-items: center;
justify-content: center;
font-weight: 400;
padding: 0.25rem 0.75rem;
border-radius: 0.5rem;
min-height: 38.4px;
margin: 0px;
line-height: 1.6;
text-decoration: none;
width: 100%;
user-select: none;
background-color: rgb(43, 44, 54);
color: rgb(250, 250, 250) !important;
border: 1px solid rgba(250, 250, 250, 0.2);
}}
.{key}:hover {{
border-color: rgb(255, 75, 75);
color: rgb(255, 75, 75);
}}
.{key} p, .{key} img {{
margin: 0;
}}
</style>
<div class="row-widget stLinkButton" data-testid="stLinkButton">
<a class="{key}" data-testid="baseLinkButton-secondary"
kind="secondary" href="{url}">
<div data-testid="stMarkdownContainer">
{f'<img src="{icon}" width="50px" height="50px">' if icon else ''}
{f'<p>{label}</p>' if label else ''}
</div>
</a>
</div>
'''
)
def logout() -> None:
revoke_token(st.session_state.access_token)
st.session_state.pop('logged_user_data')
st.session_state.pop('access_token')
st.session_state.logout_requested = True
def show_user_icon() -> None:
st.html(f'''
<div>
<img class="profileImage" src="{st.session_state.logged_user_data.get('avatar_url')}"/>
</div>
''')
@st.fragment
def show_user_data() -> None:
if 'logged_user_data' not in st.session_state:
st.write('Login to restart your past chats at any time')
create_link_button(
"login_btn",
get_login_url(),
icon="https://cdn3.emoji.gg/emojis/7561-discord-clyde.png"
)
if st.session_state.get('logout_requested'):
st.toast('Logged out')
st.session_state.logout_requested = False
if st.session_state.get('start_chat'):
reset_app()
st.switch_page(PagesEnum.HOME_PAGE)
else:
st.write(
f"Logged in as {st.session_state.logged_user_data.get('username')}")
show_user_icon()
if st.session_state.get('login_requested'):
st.toast(
f"Welcome {st.session_state.logged_user_data.get('username')}!")
st.session_state.login_requested = False
st.button(
key="logout_btn",
label="Logout",
on_click=logout,
disabled=st.session_state.logout_requested
)
SSO functions (discord.py)
from os import getenv
from requests import post, get
if not getenv('ENVIRONMENT'):
from dotenv import load_dotenv
load_dotenv(override=True)
DISCORD_API_ENDPOINT = 'https://discord.com/api'
DISCORD_CDN_URL = 'https://cdn.discordapp.com'
REDIRECT_URI = f"{getenv('BASE_URI')}/oauth2-callback"
DISCORD_CLIENT_ID = getenv('DISCORD_CLIENT_ID')
DISCORD_CLIENT_SECRET = getenv('DISCORD_CLIENT_SECRET')
LOGIN_URL = f'https://discord.com/oauth2/authorize?client_id={DISCORD_CLIENT_ID}&response_type=code&redirect_uri={REDIRECT_URI}&scope=identify&prompt=none'
LOGOUT_URL = f'{DISCORD_API_ENDPOINT}/oauth2/token/revoke'
def get_login_url() -> str:
return LOGIN_URL
def exchange_code(code) -> str:
data = {
'grant_type': 'authorization_code',
'code': code,
'redirect_uri': REDIRECT_URI
}
headers = {
'Content-Type': 'application/x-www-form-urlencoded'
}
r = post(f'{DISCORD_API_ENDPOINT}/oauth2/token', data=data,
params={'state': '1234'},
headers=headers, auth=(DISCORD_CLIENT_ID, DISCORD_CLIENT_SECRET))
access_token = r.json()['access_token']
return access_token
def get_user_data(token) -> dict[str, str]:
headers = {
'Authorization': f'Bearer {token}'
}
r = get(f'{DISCORD_API_ENDPOINT}/users/@me', headers=headers)
r.raise_for_status()
user_data = r.json()
if user_data.get('avatar'):
avatar_url = f"{DISCORD_CDN_URL}/avatars/{user_data.get('id')}/{user_data.get('avatar')}.png"
else:
avatar_url = f"{DISCORD_CDN_URL}/embed/avatars/{(int(user_data.get('id')) >> 22) % 6}.png"
user_data['avatar_url'] = avatar_url
user_data['provider'] = 'discord'
user_data.pop('avatar')
return user_data
def revoke_token(access_token) -> None:
data = {
'token': access_token,
'token_type_hint': 'access_token'
}
headers = {
'Content-Type': 'application/x-www-form-urlencoded'
}
r = post(
LOGOUT_URL,
headers=headers,
data=data,
auth=(DISCORD_CLIENT_ID, DISCORD_CLIENT_SECRET)
)
r.raise_for_status()
Oauth callback page to manage the auth flow redirect (oauth2-callback.py)
import streamlit as st
from src.utils.discord import exchange_code, get_user_data
from src.utils.frontend import reset_app
from src.utils.pages import HOME_PAGE
reset_app()
st.session_state.access_token = exchange_code(st.query_params['code'])
st.session_state.logged_user_data = get_user_data(
st.session_state.access_token)
st.session_state.login_requested = True
st.session_state.has_new_chat = True
st.switch_page(HOME_PAGE)
).