Add authentication + security hardening
- Session-based login with bcrypt password hashing - First-time setup flow creates admin account - @login_required on all page/API routes - @admin_required on admin panel and lab routes - Rate limiting: 60 req/min global, 5 login attempts/min - Security headers: X-Frame-Options, XSS Protection, nosniff - Login page with dark theme matching main UI - Logout button in header - users table in PostgreSQL Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
parent
1711d33337
commit
189e8fb99b
251
llm_team_ui.py
251
llm_team_ui.py
@ -5,14 +5,238 @@ import json
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
import secrets
|
||||
import hashlib
|
||||
import requests
|
||||
import random
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
import bcrypt
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from flask import Flask, render_template_string, request, jsonify, Response
|
||||
from flask import Flask, render_template_string, request, jsonify, Response, redirect, url_for, session
|
||||
from functools import wraps
|
||||
|
||||
app = Flask(__name__)
|
||||
app.secret_key = os.environ.get("FLASK_SECRET", secrets.token_hex(32))
|
||||
|
||||
# ─── AUTH ─────────────────────────────────────────────────────
|
||||
|
||||
_rate_limit = {} # ip -> (count, window_start)
|
||||
RATE_LIMIT_WINDOW = 60 # seconds
|
||||
RATE_LIMIT_MAX = 60 # requests per window
|
||||
LOGIN_RATE_MAX = 5 # login attempts per window
|
||||
|
||||
|
||||
def rate_limited(ip, max_req=RATE_LIMIT_MAX):
|
||||
now = time.time()
|
||||
if ip not in _rate_limit or now - _rate_limit[ip][1] > RATE_LIMIT_WINDOW:
|
||||
_rate_limit[ip] = (1, now)
|
||||
return False
|
||||
count, start = _rate_limit[ip]
|
||||
if count >= max_req:
|
||||
return True
|
||||
_rate_limit[ip] = (count + 1, start)
|
||||
return False
|
||||
|
||||
|
||||
def login_required(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
if not session.get("user_id"):
|
||||
if request.path.startswith("/api/"):
|
||||
return jsonify({"error": "unauthorized"}), 401
|
||||
return redirect("/login")
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
|
||||
|
||||
def admin_required(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
if not session.get("user_id"):
|
||||
return redirect("/login")
|
||||
if session.get("role") != "admin":
|
||||
return "Forbidden", 403
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
|
||||
|
||||
@app.before_request
|
||||
def security_checks():
|
||||
ip = request.remote_addr
|
||||
# Rate limit
|
||||
if rate_limited(ip):
|
||||
return jsonify({"error": "rate limited"}), 429
|
||||
# Allow login/static without auth
|
||||
if request.path in ("/login", "/api/auth/login", "/api/auth/setup"):
|
||||
return
|
||||
if request.path.startswith("/static"):
|
||||
return
|
||||
|
||||
|
||||
@app.after_request
|
||||
def security_headers(response):
|
||||
response.headers["X-Content-Type-Options"] = "nosniff"
|
||||
response.headers["X-Frame-Options"] = "DENY"
|
||||
response.headers["X-XSS-Protection"] = "1; mode=block"
|
||||
response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
|
||||
if request.path.startswith("/api/"):
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
return response
|
||||
|
||||
|
||||
LOGIN_HTML = """
|
||||
<!DOCTYPE html>
|
||||
<html><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1.0">
|
||||
<title>LLM Team - Login</title>
|
||||
<style>
|
||||
:root{--bg:#0a0c10;--surface:#151820;--border:#272d3f;--text:#e4e4e7;--text2:#a1a1aa;--accent:#6366f1;--accent2:#818cf8;--red:#ef4444;--green:#22c55e}
|
||||
*{box-sizing:border-box;margin:0;padding:0}
|
||||
body{font-family:'Inter',-apple-system,sans-serif;background:var(--bg);color:var(--text);min-height:100vh;display:flex;align-items:center;justify-content:center}
|
||||
.login-box{background:var(--surface);border:1px solid var(--border);border-radius:12px;padding:32px;width:360px}
|
||||
.login-box h1{font-size:22px;margin-bottom:6px;font-weight:700}
|
||||
.login-box h1 span{background:linear-gradient(135deg,var(--accent2),#a78bfa);-webkit-background-clip:text;-webkit-text-fill-color:transparent}
|
||||
.login-box .sub{color:var(--text2);font-size:13px;margin-bottom:24px}
|
||||
.field{margin-bottom:14px}
|
||||
.field label{display:block;font-size:12px;color:var(--text2);margin-bottom:4px;font-weight:500}
|
||||
.field input{width:100%;background:var(--bg);border:1px solid var(--border);color:var(--text);border-radius:6px;padding:10px 12px;font-size:14px}
|
||||
.field input:focus{outline:none;border-color:var(--accent)}
|
||||
.btn{width:100%;padding:11px;background:linear-gradient(135deg,var(--accent),#7c3aed);color:white;border:none;border-radius:8px;font-size:14px;font-weight:600;cursor:pointer;transition:all .2s}
|
||||
.btn:hover{filter:brightness(1.15)}
|
||||
.error{color:var(--red);font-size:12px;margin-bottom:12px;display:none}
|
||||
.setup-note{color:var(--green);font-size:12px;margin-bottom:12px;display:none}
|
||||
</style>
|
||||
</head><body>
|
||||
<div class="login-box">
|
||||
<h1><span>LLM</span> Team</h1>
|
||||
<p class="sub" id="subtitle">Sign in to continue</p>
|
||||
<div class="error" id="error"></div>
|
||||
<div class="setup-note" id="setup-note"></div>
|
||||
<form id="login-form" onsubmit="return doLogin(event)">
|
||||
<div class="field"><label>Username</label><input id="username" autocomplete="username" required></div>
|
||||
<div class="field"><label>Password</label><input id="password" type="password" autocomplete="current-password" required></div>
|
||||
<div class="field" id="confirm-field" style="display:none"><label>Confirm Password</label><input id="confirm" type="password"></div>
|
||||
<button class="btn" type="submit" id="submit-btn">Sign In</button>
|
||||
</form>
|
||||
</div>
|
||||
<script>
|
||||
let isSetup = false;
|
||||
async function checkSetup() {
|
||||
const r = await fetch('/api/auth/setup');
|
||||
const d = await r.json();
|
||||
if (d.needs_setup) {
|
||||
isSetup = true;
|
||||
document.getElementById('subtitle').textContent = 'Create your admin account';
|
||||
document.getElementById('confirm-field').style.display = '';
|
||||
document.getElementById('submit-btn').textContent = 'Create Account';
|
||||
document.getElementById('setup-note').textContent = 'First time setup — this will be the admin account.';
|
||||
document.getElementById('setup-note').style.display = '';
|
||||
}
|
||||
}
|
||||
async function doLogin(e) {
|
||||
e.preventDefault();
|
||||
const user = document.getElementById('username').value;
|
||||
const pass = document.getElementById('password').value;
|
||||
const err = document.getElementById('error');
|
||||
err.style.display = 'none';
|
||||
if (isSetup) {
|
||||
const confirm = document.getElementById('confirm').value;
|
||||
if (pass !== confirm) { err.textContent = 'Passwords do not match'; err.style.display = ''; return false; }
|
||||
if (pass.length < 8) { err.textContent = 'Password must be at least 8 characters'; err.style.display = ''; return false; }
|
||||
}
|
||||
const r = await fetch('/api/auth/login', {method:'POST', headers:{'Content-Type':'application/json'},
|
||||
body: JSON.stringify({username: user, password: pass, setup: isSetup})});
|
||||
const d = await r.json();
|
||||
if (d.ok) { window.location.href = '/'; }
|
||||
else { err.textContent = d.error || 'Login failed'; err.style.display = ''; }
|
||||
return false;
|
||||
}
|
||||
checkSetup();
|
||||
</script>
|
||||
</body></html>
|
||||
"""
|
||||
|
||||
|
||||
@app.route("/login")
|
||||
def login_page():
|
||||
if session.get("user_id"):
|
||||
return redirect("/")
|
||||
return LOGIN_HTML
|
||||
|
||||
|
||||
@app.route("/api/auth/setup")
|
||||
def auth_setup():
|
||||
try:
|
||||
with get_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("SELECT COUNT(*) FROM users")
|
||||
count = cur.fetchone()[0]
|
||||
return jsonify({"needs_setup": count == 0})
|
||||
except Exception:
|
||||
return jsonify({"needs_setup": True})
|
||||
|
||||
|
||||
@app.route("/api/auth/login", methods=["POST"])
|
||||
def auth_login():
|
||||
ip = request.remote_addr
|
||||
if rate_limited(ip, LOGIN_RATE_MAX):
|
||||
return jsonify({"error": "Too many attempts. Wait a minute."}), 429
|
||||
|
||||
data = request.json or {}
|
||||
username = data.get("username", "").strip()
|
||||
password = data.get("password", "")
|
||||
is_setup = data.get("setup", False)
|
||||
|
||||
if not username or not password:
|
||||
return jsonify({"error": "Username and password required"}), 400
|
||||
|
||||
if len(password) < 8:
|
||||
return jsonify({"error": "Password must be at least 8 characters"}), 400
|
||||
|
||||
try:
|
||||
with get_db() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
||||
if is_setup:
|
||||
# First-time setup: create admin
|
||||
cur.execute("SELECT COUNT(*) as c FROM users")
|
||||
if cur.fetchone()["c"] > 0:
|
||||
return jsonify({"error": "Setup already completed"}), 400
|
||||
pw_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
|
||||
cur.execute("INSERT INTO users (username, password_hash, role) VALUES (%s, %s, 'admin') RETURNING id",
|
||||
(username, pw_hash))
|
||||
uid = cur.fetchone()["id"]
|
||||
conn.commit()
|
||||
session["user_id"] = uid
|
||||
session["username"] = username
|
||||
session["role"] = "admin"
|
||||
session.permanent = True
|
||||
return jsonify({"ok": True})
|
||||
|
||||
# Normal login
|
||||
cur.execute("SELECT * FROM users WHERE username = %s", (username,))
|
||||
user = cur.fetchone()
|
||||
if not user or not bcrypt.checkpw(password.encode(), user["password_hash"].encode()):
|
||||
return jsonify({"error": "Invalid credentials"}), 401
|
||||
|
||||
session["user_id"] = user["id"]
|
||||
session["username"] = user["username"]
|
||||
session["role"] = user["role"]
|
||||
session.permanent = True
|
||||
return jsonify({"ok": True})
|
||||
except Exception as e:
|
||||
return jsonify({"error": str(e)}), 500
|
||||
|
||||
|
||||
@app.route("/api/auth/logout", methods=["POST"])
|
||||
def auth_logout():
|
||||
session.clear()
|
||||
return jsonify({"ok": True})
|
||||
|
||||
|
||||
@app.route("/logout")
|
||||
def logout_page():
|
||||
session.clear()
|
||||
return redirect("/login")
|
||||
|
||||
CONFIG_PATH = "/root/llm_team_config.json"
|
||||
DEFAULT_CONFIG = {
|
||||
@ -237,6 +461,7 @@ HTML = r"""
|
||||
<button onclick="toggleHistory()" style="color:var(--text2);background:none;font-size:12px;padding:4px 10px;border:1px solid var(--border);border-radius:6px;cursor:pointer;">History</button>
|
||||
<a href="/lab" style="color:var(--green);text-decoration:none;font-size:12px;padding:4px 10px;border:1px solid rgba(34,197,94,0.3);border-radius:6px;">Lab</a>
|
||||
<a href="/admin" style="color:var(--text2);text-decoration:none;font-size:12px;padding:4px 10px;border:1px solid var(--border);border-radius:6px;">Admin</a>
|
||||
<a href="/logout" style="color:#ef4444;text-decoration:none;font-size:11px;padding:4px 8px;border:1px solid rgba(239,68,68,0.2);border-radius:6px;opacity:0.7">Logout</a>
|
||||
</div>
|
||||
</header>
|
||||
<div class="grid">
|
||||
@ -1923,11 +2148,13 @@ def parallel_query(models, prompt):
|
||||
# ─── ROUTES ────────────────────────────────────────────────────
|
||||
|
||||
@app.route("/")
|
||||
@login_required
|
||||
def index():
|
||||
return render_template_string(HTML)
|
||||
|
||||
|
||||
@app.route("/api/models")
|
||||
@login_required
|
||||
def get_models():
|
||||
SKIP = {"nomic-embed-text", "mxbai-embed-large", "all-minilm", "snowflake-arctic-embed"}
|
||||
cfg = load_config()
|
||||
@ -1968,11 +2195,13 @@ def get_models():
|
||||
# ─── ADMIN ROUTES ─────────────────────────────────────────────
|
||||
|
||||
@app.route("/admin")
|
||||
@admin_required
|
||||
def admin_page():
|
||||
return render_template_string(ADMIN_HTML)
|
||||
|
||||
|
||||
@app.route("/api/admin/config", methods=["GET"])
|
||||
@admin_required
|
||||
def admin_get_config():
|
||||
cfg = load_config()
|
||||
safe = json.loads(json.dumps(cfg))
|
||||
@ -1986,6 +2215,7 @@ def admin_get_config():
|
||||
|
||||
|
||||
@app.route("/api/admin/config", methods=["POST"])
|
||||
@admin_required
|
||||
def admin_save_config():
|
||||
data = request.json
|
||||
cfg = load_config()
|
||||
@ -2007,6 +2237,7 @@ def admin_save_config():
|
||||
|
||||
|
||||
@app.route("/api/admin/test-provider", methods=["POST"])
|
||||
@admin_required
|
||||
def admin_test_provider():
|
||||
data = request.json
|
||||
name = data.get("provider", "")
|
||||
@ -2043,6 +2274,7 @@ def admin_test_provider():
|
||||
_or_models_cache = {"data": None, "ts": 0}
|
||||
|
||||
@app.route("/api/admin/openrouter/models")
|
||||
@admin_required
|
||||
def admin_openrouter_models():
|
||||
import time
|
||||
now = time.time()
|
||||
@ -2067,6 +2299,7 @@ def admin_openrouter_models():
|
||||
|
||||
|
||||
@app.route("/api/admin/ollama-models")
|
||||
@admin_required
|
||||
def admin_ollama_models():
|
||||
cfg = load_config()
|
||||
base = cfg["providers"]["ollama"].get("base_url", "http://localhost:11434")
|
||||
@ -2092,6 +2325,7 @@ def admin_ollama_models():
|
||||
# ─── HISTORY ROUTES ────────────────────────────────────────────
|
||||
|
||||
@app.route("/api/runs")
|
||||
@login_required
|
||||
def get_runs():
|
||||
try:
|
||||
with get_db() as conn:
|
||||
@ -2106,6 +2340,7 @@ def get_runs():
|
||||
|
||||
|
||||
@app.route("/api/runs/<int:run_id>")
|
||||
@login_required
|
||||
def get_run(run_id):
|
||||
try:
|
||||
with get_db() as conn:
|
||||
@ -2121,6 +2356,7 @@ def get_run(run_id):
|
||||
|
||||
|
||||
@app.route("/api/runs/<int:run_id>", methods=["DELETE"])
|
||||
@login_required
|
||||
def delete_run(run_id):
|
||||
try:
|
||||
with get_db() as conn:
|
||||
@ -2133,6 +2369,7 @@ def delete_run(run_id):
|
||||
|
||||
|
||||
@app.route("/api/pipelines")
|
||||
@login_required
|
||||
def get_pipelines():
|
||||
try:
|
||||
with get_db() as conn:
|
||||
@ -2147,6 +2384,7 @@ def get_pipelines():
|
||||
|
||||
|
||||
@app.route("/api/pipelines/<int:pid>")
|
||||
@login_required
|
||||
def get_pipeline(pid):
|
||||
try:
|
||||
with get_db() as conn:
|
||||
@ -2363,11 +2601,13 @@ def _ratchet_loop(exp_id):
|
||||
# ─── LAB API ROUTES ───────────────────────────────────────────
|
||||
|
||||
@app.route("/lab")
|
||||
@admin_required
|
||||
def lab_page():
|
||||
return render_template_string(LAB_HTML)
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments", methods=["GET"])
|
||||
@admin_required
|
||||
def lab_list():
|
||||
with get_db() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
||||
@ -2379,6 +2619,7 @@ def lab_list():
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments", methods=["POST"])
|
||||
@admin_required
|
||||
def lab_create():
|
||||
d = request.json
|
||||
with get_db() as conn:
|
||||
@ -2398,6 +2639,7 @@ def lab_create():
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>", methods=["GET"])
|
||||
@admin_required
|
||||
def lab_get(eid):
|
||||
with get_db() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
||||
@ -2414,6 +2656,7 @@ def lab_get(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>", methods=["PUT"])
|
||||
@admin_required
|
||||
def lab_update(eid):
|
||||
d = request.json
|
||||
sets, vals = [], []
|
||||
@ -2439,6 +2682,7 @@ def lab_update(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>/start", methods=["POST"])
|
||||
@admin_required
|
||||
def lab_start(eid):
|
||||
with get_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
@ -2453,6 +2697,7 @@ def lab_start(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>/pause", methods=["POST"])
|
||||
@admin_required
|
||||
def lab_pause(eid):
|
||||
with get_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
@ -2462,6 +2707,7 @@ def lab_pause(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>/reset", methods=["POST"])
|
||||
@admin_required
|
||||
def lab_reset(eid):
|
||||
with get_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
@ -2472,6 +2718,7 @@ def lab_reset(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>/delete", methods=["DELETE"])
|
||||
@admin_required
|
||||
def lab_delete(eid):
|
||||
with get_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
@ -2481,6 +2728,7 @@ def lab_delete(eid):
|
||||
|
||||
|
||||
@app.route("/api/lab/experiments/<int:eid>/stream")
|
||||
@admin_required
|
||||
def lab_stream(eid):
|
||||
q = []
|
||||
_lab_streams.setdefault(eid, []).append(q)
|
||||
@ -2504,6 +2752,7 @@ def lab_stream(eid):
|
||||
# ─── TEAM ROUTES ──────────────────────────────────────────────
|
||||
|
||||
@app.route("/api/run", methods=["POST"])
|
||||
@login_required
|
||||
def run_team():
|
||||
config = request.json
|
||||
mode = config["mode"]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user