src/cookie_token_store.py

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