From f02124c9de7a38d5c741e0517071444772314895 Mon Sep 17 00:00:00 2001 From: drjones Date: Mon, 7 Sep 2026 19:03:37 -0700 Subject: [PATCH] v1.1 security release: PBKDF2 passwords, session expiry, rate limiting, SSE streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - store.py: passwords now PBKDF2-HMAC-SHA256 (390k iters, random salt, constant-time compare). Legacy unsalted-SHA256 hashes upgrade transparently on successful login — zero impact on existing users. - Sessions expire after 30 days (lazy cleanup on token lookup). - Rate limiting on register (5/hr/IP), login (10/15min per IP AND per username), chat (30/hr/IP) — SQLite-backed, shared across gunicorn workers. - Registration now requires 8+ char passwords. - New /api/chat/stream SSE endpoint: citations first ('meta'), streamed answer deltas, 'done'. Client auto-falls back to non-streaming. - nginx: proxy_buffering off for SSE. --- app.py | 175 ++++++++++++++++++++++++++++++++++++++++++- astraea-nginx.conf | 6 ++ store.py | 77 ++++++++++++++++--- templates/index.html | 38 +++++++++- 4 files changed, 281 insertions(+), 15 deletions(-) diff --git a/app.py b/app.py index f406854..56ef02a 100644 --- a/app.py +++ b/app.py @@ -24,6 +24,29 @@ from agents import AGENTS, AGENT_BY_ID, OLLAMA_URL, RAG_MODEL, GENERAL_MODEL app = Flask(__name__) +# ── Rate limiting (SQLite-backed sliding window; shared across gunicorn workers) ── +def _rate_limit(key, max_events, window_seconds): + """True if allowed, False if the caller exceeded max_events in window. + Counters live in the astraea.db `rate_limits` table so all workers share them.""" + now = time.time() + try: + c = store._conn() + c.execute("CREATE TABLE IF NOT EXISTS rate_limits (key TEXT, ts REAL)") + c.execute("DELETE FROM rate_limits WHERE ts < ?", (now - 3600,)) + row = c.execute("SELECT COUNT(*) FROM rate_limits WHERE key=? AND ts >= ?", + (key, now - window_seconds)).fetchone() + if row[0] >= max_events: + c.commit(); c.close(); return False + c.execute("INSERT INTO rate_limits (key, ts) VALUES (?,?)", (key, now)) + c.commit(); c.close() + return True + except Exception: + return True # fail open — availability over strictness + +def _client_ip(): + return request.headers.get("X-Real-IP") or request.remote_addr or "unknown" + + # ── Professional avatars (gradient + glyph, rendered client-side) ── AVATARS = [ {"id": 0, "glyph": "\u2696", "c1": "#6ea8ff", "c2": "#a78bfa", "label": "Scales"}, @@ -131,6 +154,32 @@ def _ollama_chat(model, messages, tools=None, num_ctx=24000, num_predict=1400, return json.loads(r.read().decode("utf-8")) +def _ollama_chat_stream(model, messages, num_ctx=24000, num_predict=1400, + temperature=0.2, timeout=180): + """Yield content deltas from Ollama's streaming /api/chat (no tools — the + streaming path is for final answers only; tool-calling stays non-streaming).""" + payload = {"model": model, "messages": messages, "stream": True, "think": False, + "options": {"temperature": temperature, "num_predict": num_predict, "num_ctx": num_ctx}} + req = urllib.request.Request(f"{OLLAMA_URL}/api/chat", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}) + op = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with op.open(req, timeout=timeout) as r: + for line in r: + line = line.strip() + if not line: + continue + try: + chunk = json.loads(line) + except Exception: + continue + if chunk.get("done"): + break + msg = chunk.get("message", {}) + delta = msg.get("content") or "" + if delta: + yield delta + + def _web_search(query, limit=5): """Search the live web via local SearXNG (separate service, not the LLM).""" try: @@ -383,13 +432,16 @@ def app_page(): # ── Auth API ── @app.route("/api/register", methods=["POST"]) def api_register(): + ip = _client_ip() + if not _rate_limit(f"register:{ip}", 5, 3600): + return jsonify({"error": "Too many attempts. Try again later."}), 429 d = request.get_json(force=True, silent=True) or {} u = (d.get("username") or "").strip().lower() p = (d.get("password") or "").strip() if len(u) < 2 or len(u) > 30 or not u.isalnum(): return jsonify({"error": "Username must be 2-30 alphanumeric chars"}), 400 - if len(p) < 4: - return jsonify({"error": "Password must be 4+ chars"}), 400 + if len(p) < 8: + return jsonify({"error": "Password must be 8+ chars"}), 400 c = store._conn() if c.execute("SELECT id FROM users WHERE username=?", (u,)).fetchone(): c.close() @@ -401,8 +453,13 @@ def api_register(): @app.route("/api/login", methods=["POST"]) def api_login(): + ip = _client_ip() + # 10 attempts / 15 min / IP, and 10 / 15 min per username (brute-force guard) d = request.get_json(force=True, silent=True) or {} u = (d.get("username") or "").strip().lower() + if not _rate_limit(f"login-ip:{ip}", 10, 900) or \ + (u and not _rate_limit(f"login-user:{u}", 10, 900)): + return jsonify({"error": "Too many attempts. Try again later."}), 429 p = (d.get("password") or "").strip() uid = store.authenticate(u, p) if not uid: @@ -633,6 +690,120 @@ def api_chat(): return jsonify(result) +# ── Streaming chat (SSE) — same retrieval + safety context, streamed answer ── +def _chat_context(agent_id, message): + """Shared retrieval/context assembly for /api/chat and /api/chat/stream.""" + agent = AGENT_BY_ID.get(agent_id) + if not agent: + return None, None, None, None + message = message[:4000] + uid = _uid() + profile = store.get_profile(uid) if uid else {"about_me": "", "display_name": ""} + settings = store.get_settings(uid) if uid else {} + return agent, uid, profile, settings + + +@app.route("/api/chat/stream", methods=["POST"]) +def api_chat_stream(): + """SSE stream: 'meta' event first (citations), then 'delta' events with the + answer text, then 'done'. Falls back to a single 'delta' with an error if + the engine fails. Persists the conversation like /api/chat does.""" + d = request.get_json(force=True, silent=True) or {} + agent_id = (d.get("agent_id") or "navigator").strip() + message = (d.get("message") or "").strip() + client_history = d.get("history") or [] + + ip = _client_ip() + if not _rate_limit(f"chat:{ip}", 30, 3600): + return jsonify({"error": "Too many requests. Try again later."}), 429 + + agent, uid, profile, settings = _chat_context(agent_id, message) + if not agent: + return jsonify({"error": "Unknown agent"}), 400 + if not message: + return jsonify({"error": "Message required"}), 400 + + history = store.get_conversation(uid, agent_id) if uid else client_history + + # reuse the exact same context assembly as /api/chat by refactoring _build_answer's front half + top = rag.retrieve(message, agent["id"], top_k=5) + context_blocks, citations, seen = [], [], set() + for i, (score, c) in enumerate(top, 1): + key = (c["source"], c["title"]) + if key in seen: + continue + seen.add(key) + context_blocks.append(f"[{i}] ({c['source']} — {c['title']})\n{c['text']}") + citations.append({"source": c["source"], "title": c["title"], + "score": round(score, 3), "snippet": c["text"][:260]}) + if uid: + try: + for dd in userdocs.retrieve(uid, message, top_k=4): + context_blocks.append(f"[USER DOC — {dd['source']}]\n{dd['text']}") + citations.append({"source": dd["source"], "title": "your document", + "score": dd["score"], "snippet": dd["text"][:260]}) + except Exception: + pass + + system = ( + f"{agent['system']}\n\n" + "Rules:\n" + "- Give REAL, practical legal advice: state the rule, apply it to the user's facts, " + "tell them what to do next.\n" + "- Answer DIRECTLY. Do not restate the question or narrate your reasoning.\n" + "- Ground every legal claim in the reference documents. Cite the RCW section AND the " + "controlling case law by name.\n" + "- Be plain-English, specific to Washington State.\n" + "- Items marked '[USER DOC]' are the user's own uploaded files — reference them by name " + "to ground your answer in their actual case.\n" + + about_me_block(profile.get("about_me", "")) + + safety_block(settings) + ) + user = ("REFERENCE DOCUMENTS:\n\n" + "\n\n".join(context_blocks) + + f"\n\nUSER QUESTION: {message}") + messages = [{"role": "system", "content": system}] + if history: + for turn in history[-6:]: + if turn.get("role") in ("user", "assistant"): + messages.append({"role": turn["role"], "content": turn["content"]}) + messages.append({"role": "user", "content": user}) + + def generate(): + full = [] + try: + yield f"event: meta\ndata: {json.dumps({'citations': citations, 'agent': agent_id})}\n\n" + try: + for delta in _ollama_chat_stream(RAG_MODEL, messages): + full.append(delta) + yield f"event: delta\ndata: {json.dumps({'t': delta})}\n\n" + except Exception: + # model failure — fall back to the reliable general model + try: + for delta in _ollama_chat_stream(GENERAL_MODEL, messages): + full.append(delta) + yield f"event: delta\ndata: {json.dumps({'t': delta})}\n\n" + except Exception: + yield "event: delta\ndata: {\"t\": \"The legal engine is temporarily unavailable. Try again shortly.\"}\n\n" + answer = _strip_sources("".join(full).strip()) + if uid and answer: + store.save_message(uid, agent_id, "user", message) + store.save_message(uid, agent_id, "assistant", answer, citations) + yield f"event: done\ndata: {json.dumps({'citations': citations, 'grounded': bool(context_blocks)})}\n\n" + except GeneratorExit: + # client disconnected mid-stream + try: + answer = _strip_sources("".join(full).strip()) + if uid and answer: + store.save_message(uid, agent_id, "user", message) + store.save_message(uid, agent_id, "assistant", answer, citations) + except Exception: + pass + raise + + return Response(generate(), mimetype="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) + + @app.route("/api/conversation/", methods=["GET", "DELETE"]) def api_conversation(agent_id): uid = _uid() diff --git a/astraea-nginx.conf b/astraea-nginx.conf index a5a3f46..e0a13f2 100644 --- a/astraea-nginx.conf +++ b/astraea-nginx.conf @@ -10,5 +10,11 @@ server { proxy_set_header X-Forwarded-Proto $scheme; proxy_read_timeout 300; proxy_send_timeout 300; + # SSE streaming (/api/chat/stream) — no buffering, long read window + proxy_buffering off; + proxy_cache off; + proxy_http_version 1.1; + proxy_set_header Connection ''; + chunked_transfer_encoding on; } } diff --git a/store.py b/store.py index 1c760c1..f5c1ba6 100644 --- a/store.py +++ b/store.py @@ -1,14 +1,54 @@ # -*- coding: utf-8 -*- """Astraea data layer — no-KYC auth, profiles (avatar + About Me), and settings. SQLite at /opt/astraea/astraea.db. stdlib-only. + +Security model (v1.1): +- Passwords: PBKDF2-HMAC-SHA256, 390k iterations (OWASP 2023 rec), 16-byte salt, + constant-time compare. Legacy unsalted-SHA256 hashes transparently upgrade + to PBKDF2 on successful login (zero user impact). +- Sessions: 30-day sliding expiry, cleaned up lazily. """ import hashlib +import hmac import json import secrets import sqlite3 +import time DB_PATH = "/opt/astraea/astraea.db" +# ── Password hashing (PBKDF2, stdlib) ──────────────────────────────────────── +PBKDF2_ITERATIONS = 390_000 +PBKDF2_PREFIX = "pbkdf2_sha256$" # format: pbkdf2_sha256$$$ +SESSION_TTL = 30 * 24 * 3600 # 30 days + + +def hash_password(pw): + """Hash a password with PBKDF2-HMAC-SHA256 + random salt. Returns a + self-describing string so algorithms can evolve independently.""" + salt = secrets.token_bytes(16) + dk = hashlib.pbkdf2_hmac("sha256", pw.encode(), salt, PBKDF2_ITERATIONS) + return f"{PBKDF2_PREFIX}{PBKDF2_ITERATIONS}${salt.hex()}${dk.hex()}" + + +def _legacy_sha256(pw): + return hashlib.sha256(pw.encode()).hexdigest() + + +def _verify_password(stored_hash, pw): + """Return (ok, needs_upgrade). Handles both legacy unsalted SHA-256 and + PBKDF2 hashes. Uses constant-time comparisons.""" + if stored_hash.startswith(PBKDF2_PREFIX): + try: + _, iters, salt_hex, hash_hex = stored_hash.split("$", 3) + dk = hashlib.pbkdf2_hmac("sha256", pw.encode(), + bytes.fromhex(salt_hex), int(iters)) + return hmac.compare_digest(dk.hex(), hash_hex), False + except Exception: + return False, False + # legacy unsalted sha256 (64 hex chars) + return hmac.compare_digest(stored_hash, _legacy_sha256(pw)), True + def _conn(): c = sqlite3.connect(DB_PATH) @@ -90,18 +130,25 @@ def init_db(): c.commit() except Exception: pass + # speed up session expiry cleanup + try: + c.execute("CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id)") + c.commit() + except Exception: + pass c.close() -def hash_password(pw): - return hashlib.sha256(pw.encode()).hexdigest() - - def user_from_token(token): + """Resolve a session token to a user id. Sessions expire after 30 days; + expired tokens are deleted on sight (lazy cleanup).""" if not token: return None c = _conn() + cutoff = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(time.time() - SESSION_TTL)) + c.execute("DELETE FROM sessions WHERE created_at < ?", (cutoff,)) row = c.execute("SELECT user_id FROM sessions WHERE token=?", (token,)).fetchone() + c.commit() c.close() return row["user_id"] if row else None @@ -129,13 +176,27 @@ def issue_token(uid): def authenticate(username, password): + """Verify credentials. Legacy (unsalted-SHA256) hashes upgrade in place on + successful login — the user notices nothing. Returns user id or None.""" c = _conn() row = c.execute("SELECT id, password_hash FROM users WHERE username=?", (username,)).fetchone() c.close() - if row and row["password_hash"] == hash_password(password): - return row["id"] - return None + if not row: + return None + ok, needs_upgrade = _verify_password(row["password_hash"], password) + if not ok: + return None + if needs_upgrade: + try: + c = _conn() + c.execute("UPDATE users SET password_hash=? WHERE id=?", + (hash_password(password), row["id"])) + c.commit() + c.close() + except Exception: + pass # upgrade is best-effort; auth already succeeded + return row["id"] def get_profile(uid): @@ -164,8 +225,6 @@ def set_profile(uid, display_name=None, avatar=None, about_me=None, profile_pic= c.commit() c.close() return {"display_name": dn, "avatar": av, "about_me": am, "profile_pic": pp} - - def get_settings(uid): c = _conn() row = c.execute("SELECT * FROM settings WHERE user_id=?", (uid,)).fetchone() diff --git a/templates/index.html b/templates/index.html index 0ab9fcd..87b2499 100644 --- a/templates/index.html +++ b/templates/index.html @@ -316,16 +316,46 @@ async function speak(text){ try{ const r=await fetch('/api/speak',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({text:text.replace(/\n+/g,'. ').substring(0,2500),voice:mySettings.tts_voice||'aria'})}); const b=await r.blob(); new Audio(URL.createObjectURL(b)).play(); }catch(e){} } + function liveBubble(){ + const m=document.createElement('div'); m.className='msg agent'; + const l=document.createElement('div'); l.className='who-label'; l.textContent=(AGENTS.find(x=>x.id===active)?.name||'Attorney'); + const b=document.createElement('div'); b.className='bubble'; b.textContent=''; + m.appendChild(l); m.appendChild(b); + document.getElementById('messages').appendChild(m); + document.getElementById('messages').scrollTop=document.getElementById('messages').scrollHeight; + return {m, b}; + } async function send(){ const t=document.getElementById('input').value.trim(); if(!t)return; document.getElementById('input').value=''; addBubble('user',t); history[active].push({role:'user',content:t}); typing(true); document.getElementById('sendBtn').disabled=true; try{ - const r=await fetch('/api/chat',{method:'POST',headers:{'Content-Type':'application/json',...authHeaders()},body:JSON.stringify({agent_id:active,message:t,history:history[active].slice(0,-1)})}); - const d=await r.json(); typing(false); - if(d.answer){ addBubble('assistant',d.answer,d.citations); history[active].push({role:'assistant',content:d.answer,citations:d.citations}); } - else addBubble('assistant',d.error||'Something went wrong.',null); + const r=await fetch('/api/chat/stream',{method:'POST',headers:{'Content-Type':'application/json',...authHeaders()},body:JSON.stringify({agent_id:active,message:t,history:history[active].slice(0,-1)})}); + if(!r.ok || !r.body){ // fallback to non-streaming + const d=await r.json().catch(()=>({})); typing(false); + if(d.answer){ addBubble('assistant',d.answer,d.citations); history[active].push({role:'assistant',content:d.answer,citations:d.citations}); } + else addBubble('assistant',d.error||'Something went wrong.',null); + } else { + typing(false); + const live=liveBubble(); let acc=''; let cites=null; + const reader=r.body.getReader(); const dec=new TextDecoder(); let buf=''; + while(true){ + const {done,value}=await reader.read(); if(done)break; + buf+=dec.decode(value,{stream:true}); + const parts=buf.split('\n\n'); buf=parts.pop(); + for(const part of parts){ + const ev=(part.match(/^event: (.+)$/m)||[])[1]; + const dataLine=(part.match(/^data: (.+)$/m)||[])[1]; + if(!ev||!dataLine)continue; + const payload=JSON.parse(dataLine); + if(ev==='delta'){ acc+=payload.t; live.b.textContent=acc; document.getElementById('messages').scrollTop=document.getElementById('messages').scrollHeight; } + else if(ev==='done'){ cites=payload.citations; } + } + } + live.m.remove(); + addBubble('assistant',acc,cites); history[active].push({role:'assistant',content:acc,citations:cites}); + } }catch(e){ typing(false); addBubble('assistant','Could not reach the engine.',null); } finally{ document.getElementById('sendBtn').disabled=false; } }