v1.1 security release: PBKDF2 passwords, session expiry, rate limiting, SSE streaming
- 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.
This commit is contained in:
drjones
175
app.py
175
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/<agent_id>", methods=["GET", "DELETE"])
|
||||
def api_conversation(agent_id):
|
||||
uid = _uid()
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
75
store.py
75
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$<iters>$<salt_hex>$<hash_hex>
|
||||
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"]
|
||||
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()
|
||||
|
||||
@@ -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);
|
||||
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; }
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user