feat: implemented account system
This commit is contained in:
232
db.py
Normal file
232
db.py
Normal file
@@ -0,0 +1,232 @@
|
||||
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()}
|
||||
Reference in New Issue
Block a user