diff --git a/scripts/api.py b/scripts/api.py index b8caca7..9c86727 100644 --- a/scripts/api.py +++ b/scripts/api.py @@ -102,7 +102,15 @@ re-brief on context that's already in memory or documents.""" # Auth configuration import os SESSION_PASSWORD = os.getenv("AARON_AI_PASSWORD", "changeme") -SESSIONS: set = set() # In-memory session store +SESSIONS_DB = str(Path.home() / "aaronai" / "sessions.db") + +def _init_sessions(): + conn = sqlite3.connect(SESSIONS_DB) + conn.execute("CREATE TABLE IF NOT EXISTS sessions (token TEXT PRIMARY KEY, created_at TEXT)") + conn.commit() + conn.close() + +_init_sessions() def make_session_token() -> str: return secrets.token_urlsafe(32) @@ -110,12 +118,30 @@ def make_session_token() -> str: def hash_password(password: str) -> str: return hashlib.sha256(password.encode()).hexdigest() +def save_session(token: str): + conn = sqlite3.connect(SESSIONS_DB) + conn.execute("INSERT OR REPLACE INTO sessions VALUES (?, ?)", (token, datetime.now().isoformat())) + conn.commit() + conn.close() + +def delete_session(token: str): + conn = sqlite3.connect(SESSIONS_DB) + conn.execute("DELETE FROM sessions WHERE token = ?", (token,)) + conn.commit() + conn.close() + +def session_exists(token: str) -> bool: + conn = sqlite3.connect(SESSIONS_DB) + row = conn.execute("SELECT 1 FROM sessions WHERE token = ?", (token,)).fetchone() + conn.close() + return row is not None + def get_session(request: Request) -> str | None: return request.cookies.get("aaronai_session") def require_auth(request: Request): token = get_session(request) - if not token or token not in SESSIONS: + if not token or not session_exists(token): raise HTTPException(status_code=401, detail="Not authenticated") return token @@ -321,7 +347,7 @@ async def login(request: Request, response: Response): if hash_password(password) != hash_password(SESSION_PASSWORD): raise HTTPException(status_code=401, detail="Invalid password") token = make_session_token() - SESSIONS.add(token) + save_session(token) response.set_cookie( key="aaronai_session", value=token, @@ -339,7 +365,7 @@ async def login(request: Request, response: Response): async def logout(request: Request, response: Response): token = get_session(request) if token: - SESSIONS.discard(token) + delete_session(token) response.delete_cookie("aaronai_session") return JSONResponse({"ok": True})