Spaces:
Sleeping
Sleeping
Commit ·
68e5370
1
Parent(s): 40f139d
Replace Supabase Postgres refresh-token storage with an encrypted cookie
Browse filesSupabase connectivity didn't work out (this sandbox also can't reach
port 5432 directly, and the pooler swap didn't resolve it either).
Per user's direction, drop the Postgres dependency entirely and store
the Google refresh_token in a long-lived, Fernet-encrypted browser
cookie instead (extra-streamlit-components' CookieManager, since
Streamlit page code has no direct access to set real HttpOnly
cookies). Trade-off vs. the DB approach: this is per-browser/device,
not centrally keyed by account -- a new device still needs to click
Connect Drive once. The cookie only ever holds ciphertext; the
decryption key stays server-side in secrets and never reaches the
browser.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- src/cookie_token_store.py +46 -0
- src/token_store.py +0 -59
src/cookie_token_store.py
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import hashlib
|
| 2 |
+
from datetime import datetime, timedelta, timezone
|
| 3 |
+
|
| 4 |
+
import extra_streamlit_components as stx
|
| 5 |
+
import streamlit as st
|
| 6 |
+
from cryptography.fernet import Fernet
|
| 7 |
+
|
| 8 |
+
_COOKIE_MAX_AGE_DAYS = 400 # ~ the longest lifetime browsers allow anyway (Chrome caps at 400 days)
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
@st.cache_resource
|
| 12 |
+
def _cookie_manager() -> stx.CookieManager:
|
| 13 |
+
return stx.CookieManager(key="bgt_cookie_manager")
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def _fernet() -> Fernet:
|
| 17 |
+
return Fernet(st.secrets["cookies"]["encryption_key"])
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def _cookie_name(email: str) -> str:
|
| 21 |
+
# cookie-name must be a valid HTTP token, so hash the email rather than use it raw.
|
| 22 |
+
return "bgt_rt_" + hashlib.sha256(email.encode()).hexdigest()[:16]
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def save_refresh_token(email: str, refresh_token: str) -> None:
|
| 26 |
+
encrypted = _fernet().encrypt(refresh_token.encode()).decode()
|
| 27 |
+
expires_at = datetime.now(timezone.utc) + timedelta(days=_COOKIE_MAX_AGE_DAYS)
|
| 28 |
+
_cookie_manager().set(
|
| 29 |
+
_cookie_name(email), encrypted, expires_at=expires_at, key=f"set_{_cookie_name(email)}"
|
| 30 |
+
)
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def get_refresh_token(email: str) -> str | None:
|
| 34 |
+
encrypted = _cookie_manager().get(_cookie_name(email))
|
| 35 |
+
if not encrypted:
|
| 36 |
+
return None
|
| 37 |
+
try:
|
| 38 |
+
return _fernet().decrypt(encrypted.encode()).decode()
|
| 39 |
+
except Exception:
|
| 40 |
+
return None
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def delete_refresh_token(email: str) -> None:
|
| 44 |
+
name = _cookie_name(email)
|
| 45 |
+
if _cookie_manager().get(name) is not None:
|
| 46 |
+
_cookie_manager().delete(name, key=f"delete_{name}")
|
src/token_store.py
DELETED
|
@@ -1,59 +0,0 @@
|
|
| 1 |
-
import psycopg2
|
| 2 |
-
import streamlit as st
|
| 3 |
-
from cryptography.fernet import Fernet
|
| 4 |
-
|
| 5 |
-
_TABLE_DDL = """
|
| 6 |
-
CREATE TABLE IF NOT EXISTS user_refresh_tokens (
|
| 7 |
-
email TEXT PRIMARY KEY,
|
| 8 |
-
encrypted_refresh_token TEXT NOT NULL,
|
| 9 |
-
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
| 10 |
-
)
|
| 11 |
-
"""
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
@st.cache_resource
|
| 15 |
-
def _get_conn():
|
| 16 |
-
conn = psycopg2.connect(st.secrets["supabase"]["db_url"])
|
| 17 |
-
conn.autocommit = True
|
| 18 |
-
with conn.cursor() as cur:
|
| 19 |
-
cur.execute(_TABLE_DDL)
|
| 20 |
-
return conn
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
def _fernet() -> Fernet:
|
| 24 |
-
return Fernet(st.secrets["supabase"]["encryption_key"])
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
def save_refresh_token(email: str, refresh_token: str) -> None:
|
| 28 |
-
encrypted = _fernet().encrypt(refresh_token.encode()).decode()
|
| 29 |
-
conn = _get_conn()
|
| 30 |
-
with conn.cursor() as cur:
|
| 31 |
-
cur.execute(
|
| 32 |
-
"""
|
| 33 |
-
INSERT INTO user_refresh_tokens (email, encrypted_refresh_token, updated_at)
|
| 34 |
-
VALUES (%s, %s, now())
|
| 35 |
-
ON CONFLICT (email) DO UPDATE
|
| 36 |
-
SET encrypted_refresh_token = EXCLUDED.encrypted_refresh_token,
|
| 37 |
-
updated_at = now()
|
| 38 |
-
""",
|
| 39 |
-
(email, encrypted),
|
| 40 |
-
)
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
def get_refresh_token(email: str) -> str | None:
|
| 44 |
-
conn = _get_conn()
|
| 45 |
-
with conn.cursor() as cur:
|
| 46 |
-
cur.execute(
|
| 47 |
-
"SELECT encrypted_refresh_token FROM user_refresh_tokens WHERE email = %s",
|
| 48 |
-
(email,),
|
| 49 |
-
)
|
| 50 |
-
row = cur.fetchone()
|
| 51 |
-
if row is None:
|
| 52 |
-
return None
|
| 53 |
-
return _fernet().decrypt(row[0].encode()).decode()
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
def delete_refresh_token(email: str) -> None:
|
| 57 |
-
conn = _get_conn()
|
| 58 |
-
with conn.cursor() as cur:
|
| 59 |
-
cur.execute("DELETE FROM user_refresh_tokens WHERE email = %s", (email,))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|