233 lines
6.7 KiB
Python
233 lines
6.7 KiB
Python
import json
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from werkzeug.security import check_password_hash, generate_password_hash
|
|
|
|
_db_path: Path | None = None
|
|
_local = threading.local()
|
|
|
|
_SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id TEXT PRIMARY KEY,
|
|
username TEXT NOT NULL UNIQUE,
|
|
password_hash TEXT NOT NULL,
|
|
role TEXT NOT NULL DEFAULT 'user',
|
|
created_at REAL NOT NULL,
|
|
last_login REAL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
urls TEXT NOT NULL,
|
|
options TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
output TEXT NOT NULL DEFAULT '[]',
|
|
command TEXT,
|
|
return_code INTEGER,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL,
|
|
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS app_settings (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
"""
|
|
|
|
|
|
def init_db(db_path: Path) -> None:
|
|
global _db_path
|
|
_db_path = db_path
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(str(db_path))
|
|
conn.executescript(_SCHEMA)
|
|
conn.commit()
|
|
_seed_admin(conn)
|
|
conn.close()
|
|
|
|
|
|
def _seed_admin(conn: sqlite3.Connection) -> None:
|
|
username = os.environ.get("ADMIN_USERNAME", "").strip()
|
|
password = os.environ.get("ADMIN_PASSWORD", "").strip()
|
|
if not username or not password:
|
|
return
|
|
row = conn.execute("SELECT COUNT(*) FROM users").fetchone()
|
|
if row[0] > 0:
|
|
return
|
|
conn.execute(
|
|
"INSERT INTO users (id, username, password_hash, role, created_at) VALUES (?, ?, ?, 'admin', ?)",
|
|
(str(uuid.uuid4()), username, generate_password_hash(password), time.time()),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def get_db() -> sqlite3.Connection:
|
|
conn = getattr(_local, "conn", None)
|
|
if conn is None:
|
|
conn = sqlite3.connect(str(_db_path), check_same_thread=False)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
_local.conn = conn
|
|
return conn
|
|
|
|
|
|
def _row(r) -> dict | None:
|
|
return dict(r) if r else None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Users
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def create_user(username: str, password: str, role: str = "user") -> dict:
|
|
db = get_db()
|
|
user_id = str(uuid.uuid4())
|
|
try:
|
|
db.execute(
|
|
"INSERT INTO users (id, username, password_hash, role, created_at) VALUES (?, ?, ?, ?, ?)",
|
|
(user_id, username, generate_password_hash(password), role, time.time()),
|
|
)
|
|
db.commit()
|
|
except sqlite3.IntegrityError:
|
|
raise ValueError(f"Username '{username}' already exists")
|
|
return _row(db.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone())
|
|
|
|
|
|
def get_user_by_username(username: str) -> dict | None:
|
|
db = get_db()
|
|
return _row(db.execute("SELECT * FROM users WHERE username = ?", (username,)).fetchone())
|
|
|
|
|
|
def get_user_by_id(user_id: str) -> dict | None:
|
|
db = get_db()
|
|
return _row(db.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone())
|
|
|
|
|
|
def list_users() -> list[dict]:
|
|
db = get_db()
|
|
return [dict(r) for r in db.execute(
|
|
"SELECT id, username, role, created_at, last_login FROM users ORDER BY created_at"
|
|
).fetchall()]
|
|
|
|
|
|
def delete_user(user_id: str) -> None:
|
|
db = get_db()
|
|
db.execute("DELETE FROM users WHERE id = ?", (user_id,))
|
|
db.commit()
|
|
|
|
|
|
def update_user_password(user_id: str, new_password: str) -> None:
|
|
db = get_db()
|
|
db.execute(
|
|
"UPDATE users SET password_hash = ? WHERE id = ?",
|
|
(generate_password_hash(new_password), user_id),
|
|
)
|
|
db.commit()
|
|
|
|
|
|
def update_last_login(user_id: str) -> None:
|
|
db = get_db()
|
|
db.execute("UPDATE users SET last_login = ? WHERE id = ?", (time.time(), user_id))
|
|
db.commit()
|
|
|
|
|
|
def verify_password(user: dict, password: str) -> bool:
|
|
return check_password_hash(user["password_hash"], password)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Jobs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _job_from_row(row) -> dict | None:
|
|
if not row:
|
|
return None
|
|
d = dict(row)
|
|
d["urls"] = json.loads(d["urls"])
|
|
d["options"] = json.loads(d["options"])
|
|
d["output"] = json.loads(d["output"])
|
|
return d
|
|
|
|
|
|
def upsert_job(job: dict) -> None:
|
|
db = get_db()
|
|
db.execute(
|
|
"""INSERT OR REPLACE INTO jobs
|
|
(id, user_id, urls, options, status, output, command, return_code, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(
|
|
job["id"],
|
|
job.get("user_id", ""),
|
|
json.dumps(job.get("urls", [])),
|
|
json.dumps(job.get("options", {})),
|
|
job.get("status", "unknown"),
|
|
json.dumps(job.get("output", [])),
|
|
job.get("command"),
|
|
job.get("return_code"),
|
|
job.get("created_at", time.time()),
|
|
time.time(),
|
|
),
|
|
)
|
|
db.commit()
|
|
|
|
|
|
def get_job(job_id: str) -> dict | None:
|
|
db = get_db()
|
|
return _job_from_row(db.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone())
|
|
|
|
|
|
def list_jobs_for_user(user_id: str) -> list[dict]:
|
|
db = get_db()
|
|
rows = db.execute(
|
|
"SELECT * FROM jobs WHERE user_id = ? ORDER BY created_at DESC", (user_id,)
|
|
).fetchall()
|
|
return [_job_from_row(r) for r in rows]
|
|
|
|
|
|
def list_all_jobs() -> list[dict]:
|
|
db = get_db()
|
|
rows = db.execute("SELECT * FROM jobs ORDER BY created_at DESC").fetchall()
|
|
return [_job_from_row(r) for r in rows]
|
|
|
|
|
|
def delete_job(job_id: str) -> None:
|
|
db = get_db()
|
|
db.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
|
|
db.commit()
|
|
|
|
|
|
def delete_jobs_older_than(cutoff: float) -> None:
|
|
db = get_db()
|
|
db.execute("DELETE FROM jobs WHERE created_at < ?", (cutoff,))
|
|
db.commit()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Settings
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def get_setting(key: str, default: str | None = None) -> str | None:
|
|
db = get_db()
|
|
row = db.execute("SELECT value FROM app_settings WHERE key = ?", (key,)).fetchone()
|
|
return row["value"] if row else default
|
|
|
|
|
|
def set_setting(key: str, value: str) -> None:
|
|
db = get_db()
|
|
db.execute(
|
|
"INSERT OR REPLACE INTO app_settings (key, value) VALUES (?, ?)", (key, value)
|
|
)
|
|
db.commit()
|
|
|
|
|
|
def get_all_settings() -> dict:
|
|
db = get_db()
|
|
return {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM app_settings").fetchall()}
|