diff --git a/llm_team_ui.py b/llm_team_ui.py index 9c69331..958b44b 100644 --- a/llm_team_ui.py +++ b/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 = """ + + +LLM Team - Login + + +
+

LLM Team

+

Sign in to continue

+
+
+
+
+
+ + +
+
+ + +""" + + +@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""" Lab Admin + Logout
@@ -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/") +@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/", 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/") +@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/", 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/", 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//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//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//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//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//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"]