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"", 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))