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 2026-03-25 03:14:51 -05:00
parent 1711d33337
commit 189e8fb99b

View File

@ -5,14 +5,238 @@ import json
import os import os
import time import time
import threading import threading
import secrets
import hashlib
import requests import requests
import random import random
import psycopg2 import psycopg2
import psycopg2.extras import psycopg2.extras
import bcrypt
from concurrent.futures import ThreadPoolExecutor, as_completed 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 = 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" CONFIG_PATH = "/root/llm_team_config.json"
DEFAULT_CONFIG = { 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> <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="/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="/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> </div>
</header> </header>
<div class="grid"> <div class="grid">
@ -1923,11 +2148,13 @@ def parallel_query(models, prompt):
# ─── ROUTES ──────────────────────────────────────────────────── # ─── ROUTES ────────────────────────────────────────────────────
@app.route("/") @app.route("/")
@login_required
def index(): def index():
return render_template_string(HTML) return render_template_string(HTML)
@app.route("/api/models") @app.route("/api/models")
@login_required
def get_models(): def get_models():
SKIP = {"nomic-embed-text", "mxbai-embed-large", "all-minilm", "snowflake-arctic-embed"} SKIP = {"nomic-embed-text", "mxbai-embed-large", "all-minilm", "snowflake-arctic-embed"}
cfg = load_config() cfg = load_config()
@ -1968,11 +2195,13 @@ def get_models():
# ─── ADMIN ROUTES ───────────────────────────────────────────── # ─── ADMIN ROUTES ─────────────────────────────────────────────
@app.route("/admin") @app.route("/admin")
@admin_required
def admin_page(): def admin_page():
return render_template_string(ADMIN_HTML) return render_template_string(ADMIN_HTML)
@app.route("/api/admin/config", methods=["GET"]) @app.route("/api/admin/config", methods=["GET"])
@admin_required
def admin_get_config(): def admin_get_config():
cfg = load_config() cfg = load_config()
safe = json.loads(json.dumps(cfg)) safe = json.loads(json.dumps(cfg))
@ -1986,6 +2215,7 @@ def admin_get_config():
@app.route("/api/admin/config", methods=["POST"]) @app.route("/api/admin/config", methods=["POST"])
@admin_required
def admin_save_config(): def admin_save_config():
data = request.json data = request.json
cfg = load_config() cfg = load_config()
@ -2007,6 +2237,7 @@ def admin_save_config():
@app.route("/api/admin/test-provider", methods=["POST"]) @app.route("/api/admin/test-provider", methods=["POST"])
@admin_required
def admin_test_provider(): def admin_test_provider():
data = request.json data = request.json
name = data.get("provider", "") name = data.get("provider", "")
@ -2043,6 +2274,7 @@ def admin_test_provider():
_or_models_cache = {"data": None, "ts": 0} _or_models_cache = {"data": None, "ts": 0}
@app.route("/api/admin/openrouter/models") @app.route("/api/admin/openrouter/models")
@admin_required
def admin_openrouter_models(): def admin_openrouter_models():
import time import time
now = time.time() now = time.time()
@ -2067,6 +2299,7 @@ def admin_openrouter_models():
@app.route("/api/admin/ollama-models") @app.route("/api/admin/ollama-models")
@admin_required
def admin_ollama_models(): def admin_ollama_models():
cfg = load_config() cfg = load_config()
base = cfg["providers"]["ollama"].get("base_url", "http://localhost:11434") base = cfg["providers"]["ollama"].get("base_url", "http://localhost:11434")
@ -2092,6 +2325,7 @@ def admin_ollama_models():
# ─── HISTORY ROUTES ──────────────────────────────────────────── # ─── HISTORY ROUTES ────────────────────────────────────────────
@app.route("/api/runs") @app.route("/api/runs")
@login_required
def get_runs(): def get_runs():
try: try:
with get_db() as conn: with get_db() as conn:
@ -2106,6 +2340,7 @@ def get_runs():
@app.route("/api/runs/<int:run_id>") @app.route("/api/runs/<int:run_id>")
@login_required
def get_run(run_id): def get_run(run_id):
try: try:
with get_db() as conn: with get_db() as conn:
@ -2121,6 +2356,7 @@ def get_run(run_id):
@app.route("/api/runs/<int:run_id>", methods=["DELETE"]) @app.route("/api/runs/<int:run_id>", methods=["DELETE"])
@login_required
def delete_run(run_id): def delete_run(run_id):
try: try:
with get_db() as conn: with get_db() as conn:
@ -2133,6 +2369,7 @@ def delete_run(run_id):
@app.route("/api/pipelines") @app.route("/api/pipelines")
@login_required
def get_pipelines(): def get_pipelines():
try: try:
with get_db() as conn: with get_db() as conn:
@ -2147,6 +2384,7 @@ def get_pipelines():
@app.route("/api/pipelines/<int:pid>") @app.route("/api/pipelines/<int:pid>")
@login_required
def get_pipeline(pid): def get_pipeline(pid):
try: try:
with get_db() as conn: with get_db() as conn:
@ -2363,11 +2601,13 @@ def _ratchet_loop(exp_id):
# ─── LAB API ROUTES ─────────────────────────────────────────── # ─── LAB API ROUTES ───────────────────────────────────────────
@app.route("/lab") @app.route("/lab")
@admin_required
def lab_page(): def lab_page():
return render_template_string(LAB_HTML) return render_template_string(LAB_HTML)
@app.route("/api/lab/experiments", methods=["GET"]) @app.route("/api/lab/experiments", methods=["GET"])
@admin_required
def lab_list(): def lab_list():
with get_db() as conn: with get_db() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
@ -2379,6 +2619,7 @@ def lab_list():
@app.route("/api/lab/experiments", methods=["POST"]) @app.route("/api/lab/experiments", methods=["POST"])
@admin_required
def lab_create(): def lab_create():
d = request.json d = request.json
with get_db() as conn: with get_db() as conn:
@ -2398,6 +2639,7 @@ def lab_create():
@app.route("/api/lab/experiments/<int:eid>", methods=["GET"]) @app.route("/api/lab/experiments/<int:eid>", methods=["GET"])
@admin_required
def lab_get(eid): def lab_get(eid):
with get_db() as conn: with get_db() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: 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"]) @app.route("/api/lab/experiments/<int:eid>", methods=["PUT"])
@admin_required
def lab_update(eid): def lab_update(eid):
d = request.json d = request.json
sets, vals = [], [] sets, vals = [], []
@ -2439,6 +2682,7 @@ def lab_update(eid):
@app.route("/api/lab/experiments/<int:eid>/start", methods=["POST"]) @app.route("/api/lab/experiments/<int:eid>/start", methods=["POST"])
@admin_required
def lab_start(eid): def lab_start(eid):
with get_db() as conn: with get_db() as conn:
with conn.cursor() as cur: with conn.cursor() as cur:
@ -2453,6 +2697,7 @@ def lab_start(eid):
@app.route("/api/lab/experiments/<int:eid>/pause", methods=["POST"]) @app.route("/api/lab/experiments/<int:eid>/pause", methods=["POST"])
@admin_required
def lab_pause(eid): def lab_pause(eid):
with get_db() as conn: with get_db() as conn:
with conn.cursor() as cur: with conn.cursor() as cur:
@ -2462,6 +2707,7 @@ def lab_pause(eid):
@app.route("/api/lab/experiments/<int:eid>/reset", methods=["POST"]) @app.route("/api/lab/experiments/<int:eid>/reset", methods=["POST"])
@admin_required
def lab_reset(eid): def lab_reset(eid):
with get_db() as conn: with get_db() as conn:
with conn.cursor() as cur: with conn.cursor() as cur:
@ -2472,6 +2718,7 @@ def lab_reset(eid):
@app.route("/api/lab/experiments/<int:eid>/delete", methods=["DELETE"]) @app.route("/api/lab/experiments/<int:eid>/delete", methods=["DELETE"])
@admin_required
def lab_delete(eid): def lab_delete(eid):
with get_db() as conn: with get_db() as conn:
with conn.cursor() as cur: with conn.cursor() as cur:
@ -2481,6 +2728,7 @@ def lab_delete(eid):
@app.route("/api/lab/experiments/<int:eid>/stream") @app.route("/api/lab/experiments/<int:eid>/stream")
@admin_required
def lab_stream(eid): def lab_stream(eid):
q = [] q = []
_lab_streams.setdefault(eid, []).append(q) _lab_streams.setdefault(eid, []).append(q)
@ -2504,6 +2752,7 @@ def lab_stream(eid):
# ─── TEAM ROUTES ────────────────────────────────────────────── # ─── TEAM ROUTES ──────────────────────────────────────────────
@app.route("/api/run", methods=["POST"]) @app.route("/api/run", methods=["POST"])
@login_required
def run_team(): def run_team():
config = request.json config = request.json
mode = config["mode"] mode = config["mode"]