Godstrument/auth.py

"""
auth.py — invite-only accounts + per-user presets for the Godstrument.

The planet is shared; the *instrument* is personal. This stores each user's
saved patches (routes, tweaks, godtime, synth) so a logged-in browser can load
its own instrument off the shared feed. Pure stdlib: SQLite + scrypt hashing +
signed stateless session cookies. No frameworks, no new services.

  python auth.py --init                 # create the db
  python auth.py --mint-invite [N]      # print N new invite codes (default 1)
  python auth.py --users                # list accounts
  python auth.py --selftest             # assert the whole flow works

The hub mounts handle_api() at /api/* (see hub.py). Secret comes from
GODSTRUMENT_SECRET or an auto-generated auth_secret file next to the db.
"""

from __future__ import annotations
import argparse
import base64
import hashlib
import hmac
import json
import os
import re
import secrets
import sqlite3
import threading
import time

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_DB = os.path.join(HERE, "godstrument_users.db")
EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
USERNAME_RE = re.compile(r"^[A-Za-z0-9_.\-]{2,24}$")
MIN_PW = 8
SESSION_DAYS = 30
# whoever signs up with this email is an admin (set in the systemd unit)
ADMIN_EMAIL = (os.environ.get("GODSTRUMENT_ADMIN") or "").strip().lower()

SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
  id INTEGER PRIMARY KEY, username TEXT UNIQUE NOT NULL, email TEXT UNIQUE NOT NULL,
  pw_hash TEXT NOT NULL, salt TEXT NOT NULL, created REAL);
CREATE TABLE IF NOT EXISTS invites (
  code TEXT PRIMARY KEY, created REAL, used_by INTEGER);
CREATE TABLE IF NOT EXISTS presets (
  user_id INTEGER, name TEXT, data TEXT, updated REAL,
  PRIMARY KEY (user_id, name));
CREATE TABLE IF NOT EXISTS feedback (
  id INTEGER PRIMARY KEY, user_id INTEGER, username TEXT, text TEXT, created REAL);
"""

# ---- basic in-memory login throttle (friends-scale; resets on restart) ------
_ATTEMPTS: dict[str, list[float]] = {}
_ATT_LOCK = threading.Lock()
_MAX_TRIES, _WINDOW = 8, 300.0     # 8 failures / 5 min per email


def _throttled(email: str) -> bool:
    now = time.time()
    with _ATT_LOCK:
        q = [t for t in _ATTEMPTS.get(email, []) if now - t < _WINDOW]
        _ATTEMPTS[email] = q
        return len(q) >= _MAX_TRIES


def _note_fail(email: str):
    with _ATT_LOCK:
        _ATTEMPTS.setdefault(email, []).append(time.time())


# ---- db ---------------------------------------------------------------------
def _con(db: str) -> sqlite3.Connection:
    con = sqlite3.connect(db)
    con.execute("PRAGMA journal_mode=WAL")
    con.executescript(SCHEMA)
    try:                                    # migrate: add is_admin to older dbs
        con.execute("ALTER TABLE users ADD COLUMN is_admin INTEGER DEFAULT 0")
        con.commit()
    except sqlite3.OperationalError:
        pass                                # column already exists
    return con


def init_db(db: str = DEFAULT_DB):
    _con(db).close()


# ---- secret / password hashing ----------------------------------------------
def _secret(db: str = DEFAULT_DB) -> bytes:
    env = os.environ.get("GODSTRUMENT_SECRET")
    if env:
        return env.encode()
    path = os.path.join(os.path.dirname(os.path.abspath(db)), "auth_secret")
    if not os.path.exists(path):
        with open(os.open(path, os.O_CREAT | os.O_WRONLY, 0o600), "w") as f:
            f.write(secrets.token_hex(32))
    with open(path) as f:
        return f.read().strip().encode()


def hash_pw(pw: str, salt: str | None = None) -> tuple[str, str]:
    salt = salt or secrets.token_hex(16)
    h = hashlib.scrypt(pw.encode(), salt=bytes.fromhex(salt),
                       n=16384, r=8, p=1, dklen=32, maxmem=64 * 1024 * 1024)
    return salt, h.hex()


def verify_pw(pw: str, salt: str, expected: str) -> bool:
    _, h = hash_pw(pw, salt)
    return hmac.compare_digest(h, expected)


# ---- stateless signed session token -----------------------------------------
def make_token(user_id: int, db: str = DEFAULT_DB, days: int = SESSION_DAYS) -> str:
    exp = int(time.time()) + days * 86400
    payload = f"{user_id}:{exp}"
    sig = hmac.new(_secret(db), payload.encode(), hashlib.sha256).hexdigest()
    return base64.urlsafe_b64encode(f"{payload}:{sig}".encode()).decode()


def read_token(token: str, db: str = DEFAULT_DB) -> int | None:
    try:
        raw = base64.urlsafe_b64decode(token.encode()).decode()
        uid, exp, sig = raw.split(":")
        good = hmac.new(_secret(db), f"{uid}:{exp}".encode(), hashlib.sha256).hexdigest()
        if not hmac.compare_digest(sig, good):
            return None
        if int(exp) < time.time():
            return None
        return int(uid)
    except (ValueError, TypeError, AttributeError):
        return None


# ---- account operations -----------------------------------------------------
def mint_invite(db: str = DEFAULT_DB) -> str:
    code = secrets.token_urlsafe(9)
    con = _con(db)
    con.execute("INSERT INTO invites VALUES (?,?,?)", (code, time.time(), None))
    con.commit(); con.close()
    return code


def signup(username: str, email: str, pw: str, invite: str,
           db: str = DEFAULT_DB) -> tuple[bool, str, int | None]:
    username = (username or "").strip()
    email = (email or "").strip().lower()
    if not USERNAME_RE.match(username):
        return False, "username must be 2-24 letters, numbers, . _ or -", None
    if not EMAIL_RE.match(email):
        return False, "invalid email", None
    if len(pw or "") < MIN_PW:
        return False, f"password must be at least {MIN_PW} characters", None
    con = _con(db)
    try:
        row = con.execute("SELECT used_by FROM invites WHERE code=?", (invite or "",)).fetchone()
        if row is None:
            return False, "unknown invite code", None
        if row[0] is not None:
            return False, "invite code already used", None
        if con.execute("SELECT 1 FROM users WHERE email=?", (email,)).fetchone():
            return False, "email already registered", None
        if con.execute("SELECT 1 FROM users WHERE username=? COLLATE NOCASE", (username,)).fetchone():
            return False, "username already taken", None
        salt, ph = hash_pw(pw)
        cur = con.execute("INSERT INTO users (username,email,pw_hash,salt,created) VALUES (?,?,?,?,?)",
                          (username, email, ph, salt, time.time()))
        uid = cur.lastrowid
        con.execute("UPDATE invites SET used_by=? WHERE code=?", (uid, invite))
        con.commit()
        return True, "ok", uid
    finally:
        con.close()


def login(identifier: str, pw: str, db: str = DEFAULT_DB) -> int | None:
    identifier = (identifier or "").strip()
    if _throttled(identifier.lower()):
        return None
    con = _con(db)
    try:                                    # log in with email OR username
        row = con.execute("SELECT id,salt,pw_hash FROM users WHERE email=? OR username=? COLLATE NOCASE",
                          (identifier.lower(), identifier)).fetchone()
    finally:
        con.close()
    if row and verify_pw(pw or "", row[1], row[2]):
        return row[0]
    _note_fail(identifier.lower())
    return None


def user_info(user_id: int, db: str = DEFAULT_DB) -> dict | None:
    con = _con(db)
    try:
        row = con.execute("SELECT username,email,is_admin FROM users WHERE id=?", (user_id,)).fetchone()
    finally:
        con.close()
    if not row:
        return None
    admin = bool(row[2] or (ADMIN_EMAIL and row[1] == ADMIN_EMAIL))
    return {"username": row[0], "email": row[1], "admin": admin}


def _is_admin(user_id: int, db: str) -> bool:
    info = user_info(user_id, db)
    return bool(info and info.get("admin"))


def make_admin(email: str, db: str = DEFAULT_DB) -> bool:
    con = _con(db)
    cur = con.execute("UPDATE users SET is_admin=1 WHERE email=?", ((email or "").strip().lower(),))
    con.commit(); con.close()
    return cur.rowcount > 0


# ---- feedback (from friends testing) + admin stats --------------------------
def submit_feedback(user_id: int, username: str, text: str, db: str = DEFAULT_DB) -> bool:
    text = (text or "").strip()[:4000]
    if not text:
        return False
    con = _con(db)
    con.execute("INSERT INTO feedback (user_id,username,text,created) VALUES (?,?,?,?)",
                (user_id, username, text, time.time()))
    con.commit(); con.close()
    return True


def admin_stats(db: str = DEFAULT_DB) -> dict:
    con = _con(db)
    try:
        users = con.execute("SELECT COUNT(*) FROM users").fetchone()[0]
        unused = con.execute("SELECT COUNT(*) FROM invites WHERE used_by IS NULL").fetchone()[0]
        presets = con.execute("SELECT COUNT(*) FROM presets").fetchone()[0]
        accounts = [{"username": u, "email": e,
                     "joined": time.strftime("%Y-%m-%d", time.localtime(c))}
                    for u, e, c in con.execute("SELECT username,email,created FROM users ORDER BY created DESC")]
        fb = [{"username": u, "text": t,
               "when": time.strftime("%Y-%m-%d %H:%M", time.localtime(c))}
              for u, t, c in con.execute("SELECT username,text,created FROM feedback ORDER BY created DESC LIMIT 100")]
    finally:
        con.close()
    return {"users": users, "unused_codes": unused, "presets": presets,
            "accounts": accounts, "feedback": fb}


# ---- presets (a full personal patch, stored as one JSON blob) ---------------
def save_preset(user_id: int, name: str, data, db: str = DEFAULT_DB) -> bool:
    name = (name or "").strip()[:60]
    if not name:
        return False
    blob = json.dumps(data)
    if len(blob) > 256 * 1024:          # sanity cap; a patch is small
        return False
    con = _con(db)
    con.execute("INSERT INTO presets VALUES (?,?,?,?) "
                "ON CONFLICT(user_id,name) DO UPDATE SET data=excluded.data, updated=excluded.updated",
                (user_id, name, blob, time.time()))
    con.commit(); con.close()
    return True


def list_presets(user_id: int, db: str = DEFAULT_DB) -> list[dict]:
    con = _con(db)
    try:
        rows = con.execute("SELECT name,updated FROM presets WHERE user_id=? ORDER BY updated DESC",
                           (user_id,)).fetchall()
    finally:
        con.close()
    return [{"name": n, "updated": u} for n, u in rows]


def load_preset(user_id: int, name: str, db: str = DEFAULT_DB):
    con = _con(db)
    try:
        row = con.execute("SELECT data FROM presets WHERE user_id=? AND name=?",
                          (user_id, name)).fetchone()
    finally:
        con.close()
    return json.loads(row[0]) if row else None


def delete_preset(user_id: int, name: str, db: str = DEFAULT_DB) -> bool:
    con = _con(db)
    cur = con.execute("DELETE FROM presets WHERE user_id=? AND name=?", (user_id, name))
    con.commit(); con.close()
    return cur.rowcount > 0


# ---- HTTP dispatch (called by hub.py's request handler) ---------------------
COOKIE = "gs_session"


def _cookie(token: str, clear: bool = False) -> str:
    if clear:
        return f"{COOKIE}=; Path=/; Max-Age=0; HttpOnly; Secure; SameSite=Lax"
    age = SESSION_DAYS * 86400
    return f"{COOKIE}={token}; Path=/; Max-Age={age}; HttpOnly; Secure; SameSite=Lax"


def _uid_from_cookies(cookie_header: str, db: str) -> int | None:
    for part in (cookie_header or "").split(";"):
        if "=" in part:
            k, v = part.strip().split("=", 1)
            if k == COOKIE:
                return read_token(v, db)
    return None


def handle_api(method: str, path: str, body: bytes, cookie_header: str,
               db: str = DEFAULT_DB):
    """Returns (status, response_dict, set_cookie_or_None). Pure enough to test."""
    try:
        payload = json.loads(body or b"{}") if body else {}
    except (ValueError, TypeError):
        payload = {}
    uid = _uid_from_cookies(cookie_header, db)

    if path == "/api/me":
        info = user_info(uid, db) if uid else None
        return (200, info, None) if info else (401, {"error": "not logged in"}, None)

    if path == "/api/signup" and method == "POST":
        ok, msg, new_uid = signup(payload.get("username"), payload.get("email"),
                                  payload.get("password"), payload.get("invite"), db)
        if ok:
            return 200, user_info(new_uid, db), _cookie(make_token(new_uid, db))
        return 400, {"error": msg}, None

    if path == "/api/login" and method == "POST":
        u = login(payload.get("identifier") or payload.get("email"), payload.get("password"), db)
        if u:
            return 200, user_info(u, db), _cookie(make_token(u, db))
        return 401, {"error": "wrong login or password"}, None

    if path == "/api/logout" and method == "POST":
        return 200, {"ok": True}, _cookie("", clear=True)

    # everything below requires a session
    if not uid:
        return 401, {"error": "not logged in"}, None

    if path == "/api/feedback" and method == "POST":
        info = user_info(uid, db) or {}
        ok = submit_feedback(uid, info.get("username"), payload.get("text"), db)
        return (200, {"ok": True}, None) if ok else (400, {"error": "empty feedback"}, None)

    if path.startswith("/api/admin/"):
        if not _is_admin(uid, db):
            return 403, {"error": "admins only"}, None
        if path == "/api/admin/stats" and method == "GET":
            return 200, admin_stats(db), None
        if path == "/api/admin/invite" and method == "POST":
            return 200, {"code": mint_invite(db)}, None
        return 404, {"error": "not found"}, None

    if path == "/api/presets" and method == "GET":
        return 200, {"presets": list_presets(uid, db)}, None

    if path.startswith("/api/presets/"):
        name = _unquote(path[len("/api/presets/"):])
        if method == "GET":
            data = load_preset(uid, name, db)
            return (200, {"name": name, "data": data}, None) if data is not None else (404, {"error": "no such preset"}, None)
        if method in ("PUT", "POST"):
            return (200, {"ok": True}, None) if save_preset(uid, name, payload.get("data"), db) else (400, {"error": "bad preset"}, None)
        if method == "DELETE":
            return (200, {"ok": True}, None) if delete_preset(uid, name, db) else (404, {"error": "no such preset"}, None)

    return 404, {"error": "not found"}, None


def _unquote(s: str) -> str:
    from urllib.parse import unquote
    return unquote(s)


# ---- CLI + self-test --------------------------------------------------------
def _selftest():
    import tempfile
    db = os.path.join(tempfile.mkdtemp(), "t.db")
    os.environ["GODSTRUMENT_SECRET"] = "test-secret-not-real"
    init_db(db)
    code = mint_invite(db)
    # signup gates
    assert signup("bo", "bad", "password123", code, db)[0] is False, "bad email rejected"
    assert signup("bo", "a@b.co", "short", code, db)[0] is False, "short password rejected"
    assert signup("no spaces!", "a@b.co", "password123", code, db)[0] is False, "bad username rejected"
    ok, _, uid = signup("Nova", "Friend@Example.com", "hunter2pass", code, db)
    assert ok and uid, "signup should succeed with a valid invite"
    assert signup("Other", "other@x.co", "hunter2pass", code, db)[0] is False, "invite is single-use"
    assert signup("Nova2", "friend@example.com", "hunter2pass", mint_invite(db), db)[0] is False, "email unique"
    assert signup("nova", "new@x.co", "hunter2pass", mint_invite(db), db)[0] is False, "username unique (case-insensitive)"
    # login by email OR username
    assert login("friend@example.com", "hunter2pass", db) == uid
    assert login("Nova", "hunter2pass", db) == uid, "login by username works"
    assert login("friend@example.com", "wrong", db) is None
    # session round-trip
    tok = make_token(uid, db)
    assert read_token(tok, db) == uid
    assert read_token(tok + "x", db) is None, "tampered token rejected"
    assert read_token("garbage", db) is None
    # presets
    assert save_preset(uid, "my patch", {"routes": [1, 2, 3], "bpm": 120}, db)
    assert load_preset(uid, "my patch", db)["bpm"] == 120
    assert save_preset(uid, "my patch", {"bpm": 96}, db)  # overwrite
    assert load_preset(uid, "my patch", db)["bpm"] == 96
    assert [p["name"] for p in list_presets(uid, db)] == ["my patch"]
    assert load_preset(999, "my patch", db) is None, "presets are per-user"
    assert delete_preset(uid, "my patch", db) and not list_presets(uid, db)
    # http layer: unauth blocked, signup+login cookie flow works
    st, _, _ = handle_api("GET", "/api/presets", b"", "", db)
    assert st == 401, "presets require a session"
    st, resp, ck = handle_api("POST", "/api/signup",
                              json.dumps({"username": "Mika", "email": "mika@x.co",
                                          "password": "orbit12345", "invite": mint_invite(db)}).encode(), "", db)
    assert st == 200 and resp["username"] == "Mika" and ck and "HttpOnly" in ck and "Secure" in ck
    st, resp, ck = handle_api("POST", "/api/login",
                              json.dumps({"identifier": "Nova", "password": "hunter2pass"}).encode(), "", db)
    assert st == 200 and resp["email"] == "friend@example.com"
    cookie = ck.split(";")[0]
    st, resp, _ = handle_api("GET", "/api/me", b"", cookie, db)
    assert st == 200 and resp["username"] == "Nova"
    # feedback + admin gating
    assert handle_api("POST", "/api/feedback", json.dumps({"text": "love it, one bug tho"}).encode(), cookie, db)[0] == 200
    assert handle_api("GET", "/api/admin/stats", b"", cookie, db)[0] == 403, "non-admin blocked from admin"
    assert make_admin("friend@example.com", db), "grant admin"
    st, stats, _ = handle_api("GET", "/api/admin/stats", b"", cookie, db)
    assert st == 200 and stats["users"] >= 1, "admin can read stats"
    assert any("one bug tho" in f["text"] for f in stats["feedback"]), "feedback shows in stats"
    assert handle_api("POST", "/api/admin/invite", b"", cookie, db)[1].get("code"), "admin can mint codes"
    print("auth self-test: all checks passed.")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", default=DEFAULT_DB)
    ap.add_argument("--init", action="store_true")
    ap.add_argument("--mint-invite", nargs="?", type=int, const=1, metavar="N")
    ap.add_argument("--users", action="store_true")
    ap.add_argument("--make-admin", metavar="EMAIL")
    ap.add_argument("--feedback", action="store_true", help="print recent feedback")
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()
    if args.selftest:
        _selftest(); return
    if args.make_admin:
        print("granted admin" if make_admin(args.make_admin, args.db) else "no such user"); return
    if args.feedback:
        for f in admin_stats(args.db)["feedback"]:
            print(f"  [{f['when']}] @{f['username']}: {f['text']}")
        return
    if args.init:
        init_db(args.db); print(f"initialized {args.db}"); return
    if args.mint_invite:
        for _ in range(args.mint_invite):
            print("invite:", mint_invite(args.db))
        return
    if args.users:
        con = _con(args.db)
        for row in con.execute("SELECT id,username,email,created FROM users ORDER BY id"):
            print(f"  #{row[0]:3} {row[1]:16} {row[2]:30} joined {time.strftime('%Y-%m-%d', time.localtime(row[3]))}")
        con.close(); return
    ap.print_help()


if __name__ == "__main__":
    main()