RandomCatLover Claude Sonnet 5 commited on
Commit
68e5370
·
1 Parent(s): 40f139d

Replace Supabase Postgres refresh-token storage with an encrypted cookie

Browse files

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

Files changed (2) hide show
  1. src/cookie_token_store.py +46 -0
  2. 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,))