Spaces:
Sleeping
Sleeping
File size: 3,762 Bytes
68e5370 69c9e26 68e5370 69c9e26 68e5370 69c9e26 68e5370 69c9e26 68e5370 fe83ebb 68e5370 69c9e26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import hashlib
from datetime import datetime, timedelta, timezone
import streamlit as st
from cryptography.fernet import Fernet
_COOKIE_MAX_AGE_DAYS = 400 # ~ the longest lifetime browsers allow anyway (Chrome caps at 400 days)
_HTTP_DATE_FORMAT = "%a, %d %b %Y %H:%M:%S GMT"
_OAUTH_STATE_COOKIE = "bgt_oauth_state"
def _fernet() -> Fernet:
return Fernet(st.secrets["cookies"]["encryption_key"])
def _cookie_name(email: str) -> str:
# cookie-name must be a valid HTTP token, so hash the email rather than use it raw.
return "bgt_rt_" + hashlib.sha256(email.encode()).hexdigest()[:16]
def _write_cookie(name: str, value: str, expires: datetime, *, reload: bool = False) -> None:
# st.context.cookies is read-only, so writing needs a JS-side assignment --
# same technique as analytics.py's Umami injection. Fernet's base64url
# alphabet (letters, digits, -, _, =) is all cookie-safe per RFC 6265, so
# no extra encoding is needed.
#
# st.context.cookies reflects whatever the browser sent in the HTTP
# request that opened the *current* session/WebSocket -- it does not
# magically pick up a cookie set via JS partway through that same live
# session (nothing forces a fresh HTTP request for that). So after saving
# a token we care about reading back later, force a real page reload
# (not st.rerun(), which reuses the existing session) so the very next
# request actually carries the new cookie.
expires_str = expires.strftime(_HTTP_DATE_FORMAT)
script = f'document.cookie = "{name}={value}; expires={expires_str}; path=/; SameSite=Lax";'
if reload:
script += "window.location.replace(window.location.pathname);"
st.html(f"<script>{script}</script>", unsafe_allow_javascript=True)
def save_refresh_token(email: str, refresh_token: str) -> None:
encrypted = _fernet().encrypt(refresh_token.encode()).decode()
expires = datetime.now(timezone.utc) + timedelta(days=_COOKIE_MAX_AGE_DAYS)
_write_cookie(_cookie_name(email), encrypted, expires, reload=True)
def get_refresh_token(email: str) -> str | None:
encrypted = st.context.cookies.get(_cookie_name(email))
if not encrypted:
return None
try:
return _fernet().decrypt(encrypted.encode()).decode()
except Exception:
st.warning(
"Found a stored Drive cookie but couldn't decrypt it (likely the "
"encryption key changed since it was saved). Treating it as absent."
)
return None
def delete_refresh_token(email: str) -> None:
# Expiring in the past is the standard way to delete a cookie via JS.
_write_cookie(_cookie_name(email), "", datetime(1970, 1, 1, tzinfo=timezone.utc))
def save_oauth_state(state: str) -> None:
"""Stashes the Connect-Drive CSRF state in a cookie rather than
st.session_state. st.link_button always opens the authorize URL in a new
tab (Streamlit's own documented behavior), and Google's redirect lands in
that new tab -- a completely separate session with its own empty
session_state. Cookies, unlike session_state, are shared browser-wide
across tabs, so this is the only place both tabs can agree on the value.
No reload needed here (unlike save_refresh_token): the *read* of this
cookie happens in the new tab after a genuinely fresh navigation, which
already carries whatever cookies exist at that point.
"""
expires = datetime.now(timezone.utc) + timedelta(minutes=10)
_write_cookie(_OAUTH_STATE_COOKIE, state, expires)
def get_oauth_state() -> str | None:
return st.context.cookies.get(_OAUTH_STATE_COOKIE)
def clear_oauth_state() -> None:
_write_cookie(_OAUTH_STATE_COOKIE, "", datetime(1970, 1, 1, tzinfo=timezone.utc))
|