// Relay server for ArDrvLoader agents.
//
// The Cloudflare Worker is only a proxy in front of this process; all state,
// auth and the panel live here. Node >= 22 for node:sqlite.
//
//   node server.js
//
// Env:
//   PORT             listen port            (default 8790)
//   DATA_DIR         sqlite + payloads      (default ./data)
//   PANEL_PATH       panel.html             (default ../cloudflare/panel.html)
//   AGENT_SECRET     master key for agent key derivation
//   OPERATOR_KEY     panel / operator API key
//   ENROLL_SECRET    enrollment token
//   LEGACY_MASTER_AUTH  "true" accepts the shared key during rollout
//   PROXY_SECRET     when set, requests must carry X-Proxy-Secret

const http = require("node:http");
const fs = require("node:fs");
const path = require("node:path");
const crypto = require("node:crypto");
const { DatabaseSync } = require("node:sqlite");

const PORT = Number(process.env.PORT || 8790);
const DATA_DIR = path.resolve(process.env.DATA_DIR || path.join(__dirname, "data"));
const PANEL_PATH = path.resolve(process.env.PANEL_PATH || path.join(__dirname, "panel.html"));

const AGENT_SECRET = process.env.AGENT_SECRET || "";
const OPERATOR_KEY = process.env.OPERATOR_KEY || "";
const ENROLL_SECRET = process.env.ENROLL_SECRET || OPERATOR_KEY;
const PROXY_SECRET = process.env.PROXY_SECRET || "";
const LEGACY_MASTER_AUTH = process.env.LEGACY_MASTER_AUTH === "true";

fs.mkdirSync(DATA_DIR, { recursive: true });
fs.mkdirSync(path.join(DATA_DIR, "files"), { recursive: true });

const db = new DatabaseSync(path.join(DATA_DIR, "ardrv.sqlite"));

db.exec("PRAGMA journal_mode=WAL");
db.exec("PRAGMA busy_timeout=5000");
db.exec("PRAGMA synchronous=NORMAL");

db.exec(`
CREATE TABLE IF NOT EXISTS agents (
    id TEXT PRIMARY KEY,
    first_seen INTEGER NOT NULL,
    last_seen INTEGER NOT NULL,
    host TEXT, user TEXT, sid TEXT, version TEXT,
    interval INTEGER, system INTEGER, tags TEXT, revoked INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS commands (
    id TEXT PRIMARY KEY,
    agent TEXT NOT NULL,
    text TEXT NOT NULL,
    state TEXT NOT NULL,
    created INTEGER NOT NULL,
    sent INTEGER, finished INTEGER,
    operator TEXT, result TEXT, error TEXT
);
CREATE INDEX IF NOT EXISTS idx_commands_agent ON commands(agent, created DESC);
CREATE INDEX IF NOT EXISTS idx_commands_state ON commands(state);
CREATE TABLE IF NOT EXISTS events (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    ts INTEGER NOT NULL, kind TEXT NOT NULL, agent TEXT, detail TEXT
);
`);

const now = () => Math.floor(Date.now() / 1000);

function hmacHex(secret, message)
{
    return crypto.createHmac("sha256", String(secret)).update(String(message)).digest("hex");
}

function sha256Hex(buf)
{
    return crypto.createHash("sha256").update(buf || Buffer.alloc(0)).digest("hex");
}

function deriveAgentSecret(agent)
{
    return hmacHex(AGENT_SECRET, "agent:" + agent);
}

function safeEqual(a, b)
{
    const x = Buffer.from(String(a));
    const y = Buffer.from(String(b));
    return x.length === y.length && crypto.timingSafeEqual(x, y);
}

function json(res, data, status = 200)
{
    const body = JSON.stringify(data);
    res.writeHead(status, {
        "Content-Type": "application/json; charset=utf-8",
        "Content-Length": Buffer.byteLength(body),
        "Cache-Control": "no-store"
    });
    res.end(body);
}

function text(res, body, status = 200, type = "text/plain; charset=utf-8")
{
    const buf = Buffer.isBuffer(body) ? body : Buffer.from(String(body), "utf8");
    res.writeHead(status, { "Content-Type": type, "Content-Length": buf.length, "Cache-Control": "no-store" });
    res.end(buf);
}

function readBody(req)
{
    return new Promise((resolve, reject) => {
        const chunks = [];
        req.on("data", c => chunks.push(c));
        req.on("end", () => resolve(Buffer.concat(chunks)));
        req.on("error", reject);
    });
}

function authenticateAgent(req, url, body)
{
    const agent = (req.headers["x-agent-id"] || "").trim();
    const ts = (req.headers["x-agent-ts"] || "").trim();
    const sign = (req.headers["x-agent-sign"] || "").trim().toLowerCase();

    if (!AGENT_SECRET || !agent || !ts || !sign)
        return null;

    if (!/^[A-Za-z0-9._-]{1,64}$/.test(agent))
        return null;

    const tsNum = Number(ts);

    if (!Number.isFinite(tsNum) || Math.abs(now() - tsNum) > 300)
        return null;

    const canonicalPath = url.pathname.replace(/^\//, "") + url.search;
    const canonical = `${agent}\n${ts}\n${req.method}\n${canonicalPath}\n${sha256Hex(body)}`;

    if (safeEqual(hmacHex(deriveAgentSecret(agent), canonical), sign))
        return agent;

    if (LEGACY_MASTER_AUTH && safeEqual(hmacHex(AGENT_SECRET, canonical), sign))
        return agent;

    return null;
}

function operatorOk(req, url)
{
    const key = req.headers["x-operator-key"] || url.searchParams.get("key") || "";
    return OPERATOR_KEY && safeEqual(key, OPERATOR_KEY);
}

function agentRow(id)
{
    return db.prepare("SELECT * FROM agents WHERE id = ?").get(id);
}

function poll(agent, req, res)
{
    const row = agentRow(agent);

    if (row && row.revoked)
        return text(res, "revoked", 403);

    const t = now();

    db.prepare(`
        INSERT INTO agents (id, first_seen, last_seen, host, user, sid, version, interval, system)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(id) DO UPDATE SET
            last_seen = excluded.last_seen,
            host = COALESCE(excluded.host, agents.host),
            user = COALESCE(excluded.user, agents.user),
            sid = COALESCE(excluded.sid, agents.sid),
            version = COALESCE(excluded.version, agents.version),
            interval = COALESCE(excluded.interval, agents.interval),
            system = excluded.system
    `).run(
        agent, t, t,
        req.headers["x-agent-host"] || null,
        req.headers["x-agent-user"] || null,
        req.headers["x-agent-sid"] || null,
        req.headers["x-agent-version"] || null,
        Number(req.headers["x-agent-interval"]) || null,
        req.headers["x-agent-system"] === "1" ? 1 : 0
    );

    const cmd = db.prepare(`
        SELECT id, text FROM commands
        WHERE agent = ? AND state = 'queued'
        ORDER BY created ASC LIMIT 1
    `).get(agent);

    if (!cmd)
        return text(res, "", 204);

    db.prepare("UPDATE commands SET state = 'sent', sent = ? WHERE id = ?").run(t, cmd.id);

    const b64 = Buffer.from(String(cmd.text), "utf8").toString("base64");
    text(res, cmd.id + "\n" + b64);
}

function ack(agent, url, body, res)
{
    const id = url.searchParams.get("id") || "";
    const raw = body.toString("utf8");

    let result = raw;

    try { result = Buffer.from(raw, "base64").toString("utf8"); } catch { }

    if (result.length > 32768)
        result = result.slice(0, 32768) + "\n[truncated]";

    db.prepare("UPDATE commands SET state='done', finished=?, result=? WHERE id=? AND agent=?")
        .run(now(), result, id, agent);

    db.prepare("UPDATE agents SET last_seen=? WHERE id=?").run(now(), agent);

    text(res, "ok");
}

let lastPrune = 0;

// Keeps the database from growing forever. Commands are pruned by age and by
// count, events by age. Runs at most once a minute, from the agents endpoint.
function prune()
{
    if (now() - lastPrune < 60)
        return;

    lastPrune = now();

    const cmdDays = Number(process.env.RETAIN_COMMAND_DAYS || 7);
    const evtDays = Number(process.env.RETAIN_EVENT_DAYS || 30);
    const maxCommands = Number(process.env.RETAIN_MAX_COMMANDS || 50000);

    db.prepare("DELETE FROM commands WHERE created < ?").run(now() - cmdDays * 86400);
    db.prepare("DELETE FROM events WHERE ts < ?").run(now() - evtDays * 86400);

    db.prepare(`DELETE FROM commands WHERE id IN (
        SELECT id FROM commands ORDER BY created DESC LIMIT -1 OFFSET ?
    )`).run(maxCommands);
}

function reapStale()
{
    db.prepare("UPDATE commands SET state='queued', sent=NULL WHERE state='sent' AND sent IS NOT NULL AND sent < ?")
        .run(now() - 180);
}

function agentsPayload()
{
    reapStale();
    prune();

    const t = now();
    const rows = db.prepare(`
        SELECT a.*, (SELECT COUNT(*) FROM commands c WHERE c.agent = a.id AND c.state='queued') AS queued
        FROM agents a ORDER BY a.last_seen DESC
    `).all();

    return rows.map(a => {
        const interval = a.interval || 15000;
        const intervalSec = Math.max(1, Math.round(interval / 1000));
        const age = t - a.last_seen;

        return {
            id: a.id, host: a.host, user: a.user, sid: a.sid, version: a.version,
            interval, system: !!a.system, tags: a.tags, queued: a.queued,
            first_seen: a.first_seen, last_seen: a.last_seen,
            age,
            // interval is stored in ms, timestamps are seconds.
            online: age < Math.max(intervalSec * 3, 45),
            revoked: !!a.revoked
        };
    });
}

function enqueue(agent, cmdText, operator)
{
    const id = crypto.randomUUID();

    db.prepare("INSERT INTO commands (id, agent, text, state, created, operator) VALUES (?,?,?,?,?,?)")
        .run(id, agent, cmdText, "queued", now(), operator);

    db.prepare("INSERT INTO events (ts, kind, agent, detail) VALUES (?,?,?,?)")
        .run(now(), "enqueue", agent, String(cmdText).slice(0, 200));

    return id;
}

async function handle(req, res)
{
    const url = new URL(req.url, "http://local");
    const p = url.pathname;

    if (PROXY_SECRET && req.headers["x-proxy-secret"] !== PROXY_SECRET)
        return text(res, "denied", 403);

    const body = ["GET", "HEAD"].includes(req.method) ? Buffer.alloc(0) : await readBody(req);

    if (p === "/" || p === "/panel")
    {
        try
        {
            return text(res, fs.readFileSync(PANEL_PATH), 200, "text/html; charset=utf-8");
        }
        catch
        {
            return text(res, "panel missing at " + PANEL_PATH, 500);
        }
    }

    if (p.startsWith("/file/"))
    {
        const name = path.basename(decodeURIComponent(p.slice(6)));
        const file = path.join(DATA_DIR, "files", name);

        if (!fs.existsSync(file))
            return text(res, "not found", 404);

        return text(res, fs.readFileSync(file), 200, "application/octet-stream");
    }

    if (p.startsWith("/api/"))
    {
        if (p !== "/api/enroll" && !operatorOk(req, url))
            return json(res, { ok: false, error: "denied" }, 403);

        if (p === "/api/enroll" && req.method === "POST")
        {
            const provided = req.headers["x-enroll-secret"] || "";

            if (!ENROLL_SECRET || !safeEqual(provided, ENROLL_SECRET))
                return json(res, { ok: false, error: "enroll denied" }, 403);

            let payload;

            try { payload = JSON.parse(body.toString("utf8")); } catch { payload = null; }

            if (!payload || !/^[A-Za-z0-9._-]{1,64}$/.test(payload.agent || ""))
                return json(res, { ok: false, error: "valid agent id required" }, 400);

            const t = now();

            db.prepare(`INSERT INTO agents (id, first_seen, last_seen, revoked) VALUES (?,?,?,0)
                        ON CONFLICT(id) DO UPDATE SET revoked = 0`).run(payload.agent, t, t);

            db.prepare("INSERT INTO events (ts, kind, agent, detail) VALUES (?,?,?,?)")
                .run(t, "enroll", payload.agent, "issued");

            return json(res, { ok: true, agent: payload.agent, secret: deriveAgentSecret(payload.agent) });
        }

        if (p === "/api/revoke" && req.method === "POST")
        {
            let payload;

            try { payload = JSON.parse(body.toString("utf8")); } catch { payload = null; }

            if (!payload || !payload.agent)
                return json(res, { ok: false, error: "agent required" }, 400);

            const revoked = payload.revoked === false ? 0 : 1;

            db.prepare("UPDATE agents SET revoked=? WHERE id=?").run(revoked, payload.agent);

            db.prepare("INSERT INTO events (ts, kind, agent, detail) VALUES (?,?,?,?)")
                .run(now(), "revoke", payload.agent, revoked ? "revoked" : "restored");

            return json(res, { ok: true, agent: payload.agent, revoked: !!revoked });
        }

        if (p === "/api/agents")
            return json(res, { ok: true, now: now(), agents: agentsPayload() });

        if (p === "/api/stats")
        {
            const counts = db.prepare("SELECT state, COUNT(*) AS n FROM commands GROUP BY state").all();
            const agents = db.prepare(`
                SELECT COUNT(*) AS n, SUM(CASE WHEN last_seen > ? THEN 1 ELSE 0 END) AS online FROM agents
            `).get(now() - 300);

            return json(res, { ok: true, agents: agents.n || 0, online: agents.online || 0, commands: counts });
        }

        if (p === "/api/commands")
        {
            const agent = url.searchParams.get("agent") || "";
            const limit = Math.min(Number(url.searchParams.get("limit") || 50), 200);

            const rows = agent
                ? db.prepare(`SELECT id, agent, text, state, created, sent, finished, operator,
                              substr(COALESCE(result,''),1,2000) AS result
                              FROM commands WHERE agent=? ORDER BY created DESC LIMIT ?`).all(agent, limit)
                : db.prepare(`SELECT id, agent, text, state, created, sent, finished, operator,
                              substr(COALESCE(result,''),1,2000) AS result
                              FROM commands ORDER BY created DESC LIMIT ?`).all(limit);

            return json(res, { ok: true, commands: rows });
        }

        if (p === "/api/enqueue" && req.method === "POST")
        {
            let payload;

            try { payload = JSON.parse(body.toString("utf8")); } catch { payload = null; }

            if (!payload || !payload.agent || !payload.text)
                return json(res, { ok: false, error: "agent and text required" }, 400);

            return json(res, { ok: true, id: enqueue(payload.agent, payload.text, "panel") });
        }

        if (p === "/api/broadcast" && req.method === "POST")
        {
            let payload;

            try { payload = JSON.parse(body.toString("utf8")); } catch { payload = null; }

            if (!payload || !payload.text)
                return json(res, { ok: false, error: "text required" }, 400);

            const t = now();
            const rows = db.prepare("SELECT id FROM agents WHERE last_seen > ?").all(t - 300);

            for (const a of rows)
                enqueue(a.id, payload.text, "broadcast");

            return json(res, { ok: true, sent: rows.length, agents: rows.map(a => a.id) });
        }

        if (p.startsWith("/api/command/"))
        {
            const row = db.prepare("SELECT * FROM commands WHERE id=?").get(p.slice(13));

            return row ? json(res, { ok: true, command: row }) : json(res, { ok: false, error: "not found" }, 404);
        }

        return json(res, { ok: false, error: "unknown endpoint" }, 404);
    }

    if (p === "/poll" || p === "/ack" || p === "/status" || p === "/enqueue")
    {
        const agent = authenticateAgent(req, url, body);

        if (!agent)
            return text(res, "denied", 403);

        if (p === "/poll")
            return poll(agent, req, res);

        if (p === "/ack" && req.method === "POST")
            return ack(agent, url, body, res);

        if (p === "/status")
        {
            const row = db.prepare("SELECT result FROM commands WHERE agent=? AND state='done' ORDER BY finished DESC LIMIT 1").get(agent);
            const last = row && row.result ? Buffer.from(row.result, "utf8").toString("base64") : null;
            return json(res, { ok: true, agent, last });
        }

        if (p === "/enqueue" && req.method === "POST")
        {
            const cmdText = body.toString("utf8").trim();

            if (!cmdText)
                return text(res, "empty", 400);

            return json(res, { ok: true, id: enqueue(agent, cmdText, "agent") });
        }
    }

    text(res, "ok");
}

const server = http.createServer((req, res) => {
    handle(req, res).catch(err => {
        try { json(res, { ok: false, error: String(err && err.message ? err.message : err) }, 500); } catch { }
    });
});

server.listen(PORT, () => {
    console.log(`relay server on http://127.0.0.1:${PORT}`);
    console.log(`data: ${DATA_DIR}`);
    console.log(`panel: ${PANEL_PATH}`);
    console.log(`legacy master auth: ${LEGACY_MASTER_AUTH}`);
console.log(`journal mode: ${db.prepare("PRAGMA journal_mode").get().journal_mode}`);
});