SHA256
591 lines
19 KiB
Python
591 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""sjcode daemon.
|
|
|
|
Runs on the server, talks to the local ollama API (127.0.0.1) and serves
|
|
sjcode clients on its own port (from config.toml).
|
|
|
|
On startup it drops into a minimal admin CLI; while that CLI runs, an
|
|
HTTP server thread serves the client API:
|
|
|
|
POST /login {"username", "password"} -> {"token"}
|
|
POST /chat {"token", "message", "model"?}
|
|
-> {"response"} or {"error"} with 4xx
|
|
GET /models?token=... -> {"models": [...]}
|
|
|
|
User accounts live in SQLite (sjcode.db) with:
|
|
username, password_hash, last_logged_in, role, used_today (minutes).
|
|
|
|
Roles decide the daily usage limit (minutes) from config.toml [limits].
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cmd
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import sqlite3
|
|
import sys
|
|
import threading
|
|
import time
|
|
import tomllib
|
|
import urllib.error
|
|
import urllib.request
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
CONFIG_FILE = ROOT / "config.toml"
|
|
DB_FILE = ROOT / "sjcode.db" # overwritten by config value once loaded
|
|
|
|
PEPPER = b"sjcode-v1"
|
|
SESSION_TTL = 24 * 3600 # login token valid for one day
|
|
|
|
CONF: dict = {
|
|
"server": {"host": "127.0.0.1", "port": 8420},
|
|
"ollama": {"base_url": "http://127.0.0.1:11434", "model": "llama3.2"},
|
|
"database": {"path": "sjcode.db"},
|
|
"limits": {"user": 120, "admin": -1},
|
|
}
|
|
|
|
|
|
def load_config() -> None:
|
|
"""Read config.toml on top of the defaults."""
|
|
global DB_FILE
|
|
if not CONFIG_FILE.exists():
|
|
sys.exit(
|
|
f"error: {CONFIG_FILE.name} not found.\n"
|
|
f"Copy {ROOT / 'config.toml.example'} to {CONFIG_FILE} and edit it."
|
|
)
|
|
try:
|
|
with open(CONFIG_FILE, "rb") as fh:
|
|
loaded = tomllib.load(fh)
|
|
except tomllib.TOMLDecodeError as exc:
|
|
sys.exit(f"error: invalid {CONFIG_FILE.name}: {exc}")
|
|
|
|
for section, values in loaded.items():
|
|
if isinstance(values, dict) and isinstance(CONF.get(section), dict):
|
|
CONF[section].update(values)
|
|
DB_FILE = ROOT / CONF["database"]["path"]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# database
|
|
# --------------------------------------------------------------------------
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT NOT NULL UNIQUE,
|
|
password_hash TEXT NOT NULL,
|
|
role TEXT NOT NULL DEFAULT 'user',
|
|
used_today REAL NOT NULL DEFAULT 0,
|
|
last_logged_in TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
token TEXT PRIMARY KEY,
|
|
username TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL
|
|
);
|
|
"""
|
|
|
|
|
|
def db() -> sqlite3.Connection:
|
|
conn = sqlite3.connect(DB_FILE, timeout=10)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
return conn
|
|
|
|
|
|
def init_db() -> None:
|
|
with db() as conn:
|
|
conn.executescript(SCHEMA)
|
|
|
|
|
|
def hash_password(password: str) -> str:
|
|
salt = f"{time.time_ns()}:{threading.get_ident()}".encode()
|
|
digest = hashlib.scrypt(password.encode(), salt=PEPPER + salt, n=2**14, r=8, p=1)
|
|
return f"{salt.decode()}${digest.hex()}"
|
|
|
|
|
|
def verify_password(password: str, stored: str) -> bool:
|
|
try:
|
|
salt_s, hexdigest = stored.split("$", 1)
|
|
except ValueError:
|
|
return False
|
|
calc = hashlib.scrypt(
|
|
password.encode(), salt=PEPPER + salt_s.encode(), n=2**14, r=8, p=1
|
|
)
|
|
return hmac.compare_digest(calc.hex(), hexdigest)
|
|
|
|
|
|
def new_token() -> str:
|
|
import secrets
|
|
|
|
return secrets.token_hex(32)
|
|
|
|
|
|
def limit_for(role: str) -> float:
|
|
limits = CONF["limits"]
|
|
return float(limits.get(role, limits.get("user", 0)))
|
|
|
|
|
|
def ask_password(prompt: str) -> str:
|
|
"""Hidden input on a terminal, plain stdin otherwise (scriptable)."""
|
|
if sys.stdin.isatty():
|
|
import getpass
|
|
|
|
return getpass.getpass(prompt)
|
|
print(prompt, end="", flush=True)
|
|
return input("")
|
|
|
|
|
|
def usage_reset_if_needed(conn: sqlite3.Connection, row: sqlite3.Row) -> None:
|
|
"""used_today is a counter in minutes; reset it on a new calendar day (UTC)."""
|
|
last_date = row["last_logged_in"][:10]
|
|
if last_date and last_date != time.strftime("%Y-%m-%d", time.gmtime()):
|
|
conn.execute(
|
|
"UPDATE users SET used_today = 0 WHERE username = ?", (row["username"],)
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def get_user(username: str):
|
|
with db() as conn:
|
|
row = conn.execute(
|
|
"SELECT * FROM users WHERE username = ?", (username,)
|
|
).fetchone()
|
|
if row is not None:
|
|
usage_reset_if_needed(conn, row)
|
|
conn.commit()
|
|
return row
|
|
|
|
|
|
def remaining_minutes(username: str) -> float:
|
|
row = get_user(username)
|
|
if row is None:
|
|
return 0.0
|
|
limit = limit_for(row["role"])
|
|
if limit < 0: # unlimited
|
|
return float("inf")
|
|
return max(0.0, limit - float(row["used_today"]))
|
|
|
|
|
|
def add_used_minutes(username: str, minutes: float) -> None:
|
|
with db() as conn:
|
|
conn.execute(
|
|
"UPDATE users SET used_today = used_today + ? WHERE username = ?",
|
|
(minutes, username),
|
|
)
|
|
|
|
|
|
def set_last_login(username: str) -> None:
|
|
stamp = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime())
|
|
with db() as conn:
|
|
conn.execute(
|
|
"UPDATE users SET last_logged_in = ? WHERE username = ?", (stamp, username)
|
|
)
|
|
|
|
|
|
def create_session(username: str) -> str:
|
|
token = new_token()
|
|
with db() as conn:
|
|
conn.execute(
|
|
"INSERT INTO sessions (token, username, created_at) VALUES (?, ?, ?)",
|
|
(token, username, int(time.time())),
|
|
)
|
|
return token
|
|
|
|
|
|
def session_user(token: str):
|
|
"""Return the username for a valid session token, else None."""
|
|
with db() as conn:
|
|
row = conn.execute(
|
|
"SELECT username, created_at FROM sessions WHERE token = ?", (token,)
|
|
).fetchone()
|
|
if row is None or int(time.time()) - row["created_at"] > SESSION_TTL:
|
|
return None
|
|
return row["username"]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# ollama
|
|
# --------------------------------------------------------------------------
|
|
|
|
def ollama_json(path: str, payload: dict | None = None, timeout: int = 300) -> tuple[int, dict]:
|
|
url = CONF["ollama"]["base_url"].rstrip("/") + path
|
|
data = None if payload is None else json.dumps(payload).encode()
|
|
req = urllib.request.Request(
|
|
url, data=data, headers={"Content-Type": "application/json"}
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
return resp.status, json.loads(resp.read().decode())
|
|
except urllib.error.HTTPError as exc:
|
|
try:
|
|
body = json.loads(exc.read().decode())
|
|
except Exception:
|
|
body = {"error": str(exc)}
|
|
return exc.code, body
|
|
except (urllib.error.URLError, TimeoutError, ConnectionError) as exc:
|
|
return 502, {"error": f"ollama not reachable: {exc}"}
|
|
|
|
|
|
def ollama_generate(prompt: str, model: str | None) -> tuple[int, str]:
|
|
status, body = ollama_json(
|
|
"/api/generate",
|
|
{
|
|
"model": model or CONF["ollama"]["model"],
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
},
|
|
)
|
|
if status != 200:
|
|
return status, str(body.get("error", "ollama error"))
|
|
return 200, str(body.get("response", ""))
|
|
|
|
|
|
def list_models() -> tuple[int, list[str]]:
|
|
status, body = ollama_json("/api/tags", timeout=10)
|
|
if status != 200:
|
|
return status, []
|
|
names = [m.get("name", "") for m in body.get("models", [])]
|
|
return 200, names
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# HTTP API
|
|
# --------------------------------------------------------------------------
|
|
|
|
class ApiHandler(BaseHTTPRequestHandler):
|
|
server_version = "sjcode-daemon/0.1"
|
|
|
|
# -- helpers ----------------------------------------------------------
|
|
def _json(self, status: int, payload: dict) -> None:
|
|
body = json.dumps(payload).encode()
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _read_body(self) -> dict:
|
|
try:
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
return json.loads(self.rfile.read(length).decode())
|
|
except (ValueError, json.JSONDecodeError):
|
|
return {}
|
|
|
|
def log_message(self, fmt: str, *args) -> None: # quiet by default
|
|
return
|
|
|
|
# -- routes -----------------------------------------------------------
|
|
def do_GET(self) -> None: # noqa: N802 (http.server API)
|
|
if self.path.startswith("/models"):
|
|
status, models = list_models()
|
|
self._json(200 if status == 200 else 502, {"models": models})
|
|
return
|
|
if self.path == "/status":
|
|
self._json(200, {"status": "ok", "daemon": "sjcode"})
|
|
return
|
|
self._json(404, {"error": "not found"})
|
|
|
|
def do_POST(self) -> None: # noqa: N802 (http.server API)
|
|
data = self._read_body()
|
|
|
|
if self.path == "/login":
|
|
self._handle_login(data)
|
|
return
|
|
if self.path == "/chat":
|
|
self._handle_chat(data)
|
|
return
|
|
self._json(404, {"error": "not found"})
|
|
|
|
# -- handlers ---------------------------------------------------------
|
|
def _handle_login(self, data: dict) -> None:
|
|
username = str(data.get("username", "")).strip()
|
|
password = str(data.get("password", ""))
|
|
if not username or not password:
|
|
self._json(400, {"error": "username and password required"})
|
|
return
|
|
|
|
row = get_user(username)
|
|
if row is None or not verify_password(password, row["password_hash"]):
|
|
self._json(401, {"error": "invalid credentials"})
|
|
return
|
|
|
|
limit = limit_for(row["role"])
|
|
left = remaining_minutes(username)
|
|
if left <= 0:
|
|
self._json(
|
|
403,
|
|
{"error": "daily limit reached", "limit_minutes": limit},
|
|
)
|
|
return
|
|
|
|
set_last_login(username)
|
|
token = create_session(username)
|
|
self._json(
|
|
200,
|
|
{
|
|
"token": token,
|
|
"role": row["role"],
|
|
"remaining_minutes": None if left == float("inf") else round(left, 1),
|
|
},
|
|
)
|
|
|
|
def _handle_chat(self, data: dict) -> None:
|
|
token = str(data.get("token", ""))
|
|
username = session_user(token)
|
|
if username is None:
|
|
self._json(401, {"error": "not logged in"})
|
|
return
|
|
|
|
message = str(data.get("message", "")).strip()
|
|
if not message:
|
|
self._json(400, {"error": "empty message"})
|
|
return
|
|
|
|
left = remaining_minutes(username)
|
|
if left <= 0:
|
|
self._json(403, {"error": "daily limit reached"})
|
|
return
|
|
|
|
model = str(data.get("model") or "") or CONF["ollama"]["model"]
|
|
|
|
started = time.monotonic()
|
|
status, answer = ollama_generate(message, model)
|
|
elapsed = time.monotonic() - started
|
|
|
|
if status != 200:
|
|
self._json(status, {"error": answer})
|
|
return
|
|
|
|
used_min = elapsed / 60.0
|
|
add_used_minutes(username, used_min)
|
|
left_now = remaining_minutes(username)
|
|
self._json(
|
|
200,
|
|
{
|
|
"response": answer,
|
|
"model": model,
|
|
"used_minutes": round(used_min, 4),
|
|
"remaining_minutes": None if left_now == float("inf") else round(left_now, 2),
|
|
},
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# admin CLI
|
|
# --------------------------------------------------------------------------
|
|
|
|
class DaemonCLI(cmd.Cmd):
|
|
"""Minimal admin shell shown after daemon startup."""
|
|
|
|
intro = (
|
|
"sjcode daemon CLI. Type 'help' for commands.\n"
|
|
"Accounts are stored in SQLite; the HTTP API runs in the background."
|
|
)
|
|
prompt = "sjcode-admin> "
|
|
|
|
def __init__(self, server_ref: dict) -> None:
|
|
super().__init__()
|
|
self.server_ref = server_ref
|
|
|
|
# -- account management ----------------------------------------------
|
|
def do_create(self, arg: str) -> None:
|
|
"""create <username> [role] - create an account (asks for password)"""
|
|
parts = arg.split()
|
|
if not parts:
|
|
print("usage: create <username> [role]")
|
|
return
|
|
username = parts[0]
|
|
role = parts[1] if len(parts) > 1 else "user"
|
|
if get_user(username) is not None:
|
|
print(f"user '{username}' already exists")
|
|
return
|
|
pw1 = ask_password("password: ")
|
|
pw2 = ask_password("password (again): ")
|
|
if len(pw1) < 4:
|
|
print("password too short (min 4 chars)")
|
|
return
|
|
if pw1 != pw2:
|
|
print("passwords do not match")
|
|
return
|
|
with db() as conn:
|
|
conn.execute(
|
|
"INSERT INTO users (username, password_hash, role, created_at) "
|
|
"VALUES (?, ?, ?, ?)",
|
|
(
|
|
username,
|
|
hash_password(pw1),
|
|
role,
|
|
time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
|
|
),
|
|
)
|
|
print(f"created '{username}' with role '{role}' "
|
|
f"(limit: {limit_for(role)} min/day)")
|
|
|
|
def do_list(self, arg: str) -> None:
|
|
"""list - show all accounts"""
|
|
with db() as conn:
|
|
rows = conn.execute(
|
|
"SELECT username, role, used_today, last_logged_in "
|
|
"FROM users ORDER BY username"
|
|
).fetchall()
|
|
if not rows:
|
|
print("no accounts yet. use 'create <username>'")
|
|
return
|
|
print(f"{'username':<20} {'role':<10} {'used today':>10} last login")
|
|
for r in rows:
|
|
print(
|
|
f"{r['username']:<20} {r['role']:<10} "
|
|
f"{r['used_today']:>10.1f} {r['last_logged_in'] or '-'}"
|
|
)
|
|
|
|
def do_delete(self, arg: str) -> None:
|
|
"""delete <username> - remove an account"""
|
|
username = arg.strip().split()[0] if arg.strip() else ""
|
|
if not username:
|
|
print("usage: delete <username>")
|
|
return
|
|
with db() as conn:
|
|
cur = conn.execute("DELETE FROM users WHERE username = ?", (username,))
|
|
if cur.rowcount:
|
|
print(f"deleted '{username}'")
|
|
else:
|
|
print(f"no user '{username}'")
|
|
|
|
def do_role(self, arg: str) -> None:
|
|
"""role <username> <role> - change a user's role"""
|
|
parts = arg.split()
|
|
if len(parts) != 2:
|
|
print("usage: role <username> <role>")
|
|
return
|
|
username, role = parts
|
|
with db() as conn:
|
|
cur = conn.execute(
|
|
"UPDATE users SET role = ? WHERE username = ?", (role, username)
|
|
)
|
|
if cur.rowcount:
|
|
print(f"'{username}' now has role '{role}' "
|
|
f"(limit: {limit_for(role)} min/day)")
|
|
else:
|
|
print(f"no user '{username}'")
|
|
|
|
def do_passwd(self, arg: str) -> None:
|
|
"""passwd <username> - set a new password for a user"""
|
|
username = arg.strip().split()[0] if arg.strip() else ""
|
|
if not username:
|
|
print("usage: passwd <username>")
|
|
return
|
|
row = get_user(username)
|
|
if row is None:
|
|
print(f"no user '{username}'")
|
|
return
|
|
pw = ask_password("new password: ")
|
|
if len(pw) < 4:
|
|
print("password too short (min 4 chars)")
|
|
return
|
|
with db() as conn:
|
|
conn.execute(
|
|
"UPDATE users SET password_hash = ? WHERE username = ?",
|
|
(hash_password(pw), username),
|
|
)
|
|
print(f"password updated for '{username}'")
|
|
|
|
def do_usage(self, arg: str) -> None:
|
|
"""usage [username] - show used minutes / limit for today"""
|
|
with db() as conn:
|
|
if arg.strip():
|
|
rows = conn.execute(
|
|
"SELECT username, role, used_today FROM users "
|
|
"WHERE username = ?",
|
|
(arg.strip(),),
|
|
).fetchall()
|
|
else:
|
|
rows = conn.execute(
|
|
"SELECT username, role, used_today FROM users "
|
|
"ORDER BY username"
|
|
).fetchall()
|
|
if not rows:
|
|
print("no such user" if arg.strip() else "no accounts yet")
|
|
return
|
|
for r in rows:
|
|
limit = limit_for(r["role"])
|
|
lim = "unlimited" if limit < 0 else f"{limit:g} min"
|
|
print(f"{r['username']:<20} {r['used_today']:>7.1f} min / {lim}")
|
|
|
|
def do_reset(self, arg: str) -> None:
|
|
"""reset [username] - reset the used_today counter"""
|
|
with db() as conn:
|
|
if arg.strip():
|
|
conn.execute(
|
|
"UPDATE users SET used_today = 0 WHERE username = ?",
|
|
(arg.strip(),),
|
|
)
|
|
else:
|
|
conn.execute("UPDATE users SET used_today = 0")
|
|
print("usage counter reset")
|
|
|
|
# -- server control ---------------------------------------------------
|
|
def do_status(self, arg: str) -> None:
|
|
"""status - show daemon settings"""
|
|
host, port = CONF["server"]["host"], CONF["server"]["port"]
|
|
print(f"listening on http://{host}:{port}")
|
|
print(f"ollama at {CONF['ollama']['base_url']} (model: {CONF['ollama']['model']})")
|
|
print(f"database at {DB_FILE}")
|
|
|
|
def do_quit(self, arg: str) -> bool:
|
|
"""quit - stop the daemon"""
|
|
print("shutting down...")
|
|
self.server_ref["stop"] = True
|
|
srv: ThreadingHTTPServer = self.server_ref["httpd"]
|
|
srv.shutdown()
|
|
return True
|
|
|
|
def do_exit(self, arg: str) -> bool:
|
|
"""exit - stop the daemon (same as quit)"""
|
|
return self.do_quit(arg)
|
|
|
|
do_EOF = do_quit
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# main
|
|
# --------------------------------------------------------------------------
|
|
|
|
def main() -> None:
|
|
load_config()
|
|
init_db()
|
|
|
|
host = CONF["server"]["host"]
|
|
port = int(CONF["server"]["port"])
|
|
|
|
try:
|
|
httpd = ThreadingHTTPServer((host, port), ApiHandler)
|
|
except OSError as exc:
|
|
sys.exit(f"error: cannot listen on {host}:{port}: {exc}")
|
|
|
|
ref = {"httpd": httpd, "stop": False}
|
|
thread = threading.Thread(target=httpd.serve_forever, daemon=True)
|
|
thread.start()
|
|
|
|
print(f"sjcode daemon listening on http://{host}:{port}")
|
|
print(f"ollama: {CONF['ollama']['base_url']} model: {CONF['ollama']['model']}")
|
|
print(f"database: {DB_FILE}")
|
|
|
|
try:
|
|
DaemonCLI(ref).cmdloop()
|
|
except KeyboardInterrupt:
|
|
print()
|
|
finally:
|
|
httpd.shutdown()
|
|
httpd.server_close()
|
|
print("bye")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|