# sb-builder - self-checks for the remote helper
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

"""Run: python3 -m unittest discover -s tests"""

import contextlib
import fcntl
import importlib.machinery
import importlib.util
import io
import json
import os
import shutil
import sqlite3
import tempfile
import threading
import time
import unittest
from unittest import mock

SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "remote", "sbb-remote")
_loader = importlib.machinery.SourceFileLoader("sbb_remote", SCRIPT)
sbb = importlib.util.module_from_spec(importlib.util.spec_from_loader("sbb_remote", _loader))
_loader.exec_module(sbb)


class Base(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, self.tmp)
        conf = os.path.join(self.tmp, "config.toml")
        with open(conf, "w") as f:
            f.write(f'state_dir = "{self.tmp}"\n'
                    f'repo_dir = "{self.tmp}/repo"\n'
                    f'slackrepo_db = "{self.tmp}/slackrepo.sqlite3"\n'
                    f'hintdir = "{self.tmp}/hints"\n'
                    f'genreprc = "{self.tmp}/genreprc"\n')
        patcher = mock.patch.dict(os.environ, {"SBB_CONFIG": conf})
        patcher.start()
        self.addCleanup(patcher.stop)
        self.cfg = sbb.load_config()
        os.makedirs(os.path.join(self.tmp, "jobs"))

    def make_job(self, argv, **meta):
        jid = "20260101-000000-test"
        d = os.path.join(self.tmp, "jobs", jid)
        os.mkdir(d)
        sbb.write_meta(d, {"id": jid, "kind": "test", "args": [], "argv": argv,
                           "status": "starting", "created": time.time(), **meta})
        open(os.path.join(d, "log"), "wb").close()
        return jid, d

    def make_pkg(self, path, info="", hint=None, desc=None):
        d = os.path.join(self.cfg["repo_dir"], path)
        pkg = os.path.basename(path)
        os.makedirs(d, exist_ok=True)
        open(os.path.join(d, pkg + ".SlackBuild"), "w").close()
        with open(os.path.join(d, pkg + ".info"), "w") as f:
            f.write(info)
        if desc:
            with open(os.path.join(d, "slack-desc"), "w") as f:
                f.write(f"{pkg}: {pkg} ({desc})\n{pkg}:\n")
        if hint is not None:
            os.makedirs(self.cfg["hintdir"], exist_ok=True)
            with open(os.path.join(self.cfg["hintdir"], pkg + ".hint"), "w") as f:
                f.write(hint)

    def fake_bin(self, name, script):
        bindir = os.path.join(self.tmp, "bin")
        os.makedirs(bindir, exist_ok=True)
        path = os.path.join(bindir, name)
        with open(path, "w") as f:
            f.write("#!/bin/sh\n" + script)
        os.chmod(path, 0o755)
        if bindir not in os.environ["PATH"]:
            patcher = mock.patch.dict(os.environ, {"PATH": bindir + os.pathsep + os.environ["PATH"]})
            patcher.start()
            self.addCleanup(patcher.stop)
        return path

    def hold_lock(self):
        f = open(os.path.join(self.tmp, "lock"), "a")
        fcntl.flock(f, fcntl.LOCK_EX)
        self.addCleanup(f.close)

    def run_main(self, *argv):
        out = io.StringIO()
        with contextlib.redirect_stdout(out):
            rc = sbb.main(list(argv))
        return rc, json.loads(out.getvalue())

    def wait_status(self, d, statuses):
        for _ in range(100):
            if sbb.read_meta(d)["status"] in statuses:
                break
            time.sleep(0.05)
        return sbb.read_meta(d)


class ConfigTest(Base):
    def write(self, text):
        p = os.path.join(self.tmp, "c.toml")
        with open(p, "w") as f:
            f.write(text)
        return p

    def test_missing_file_gives_defaults(self):
        self.assertEqual(sbb.load_config(os.path.join(self.tmp, "nope.toml")), sbb.DEFAULTS)

    def test_override(self):
        self.assertEqual(sbb.load_config(self.write('overlays = ["a"]\n'))["overlays"], ["a"])

    def test_example_matches_defaults(self):
        example = os.path.join(os.path.dirname(SCRIPT), "config.example.toml")
        self.assertEqual(sbb.load_config(example), sbb.DEFAULTS)

    def test_unknown_key(self):
        with self.assertRaisesRegex(sbb.Error, "bogus: unknown key"):
            sbb.load_config(self.write("bogus = 1\n"))

    def test_wrong_type(self):
        with self.assertRaisesRegex(sbb.Error, "state_dir: expected str"):
            sbb.load_config(self.write("state_dir = 1\n"))


class KindsTest(Base):
    def test_sbpull_argv(self):
        self.assertEqual(sbb.KINDS["sbpull"](self.cfg, ["personal", "--no"]),
                         (["bash", "-c", '. "$1" && sbpull "$2" "$3"', "_",
                           self.cfg["sbpull_script"], "personal", "--no"], {}))

    def test_sbshadow_argv(self):
        self.assertEqual(sbb.KINDS["sbshadow"](self.cfg, ["--yes"]), (["sbshadow", "--yes"], {}))

    def test_slackrepo_argv(self):
        self.make_pkg("misc/nuclei")
        self.make_pkg("personal/foo")
        kind = sbb.KINDS["slackrepo"]
        self.assertEqual(kind(self.cfg, ["update", "--preview", "nuclei", "foo", "nuclei"]),
                         (["slackrepo", "--repo=SBo-danix", "update", "--preview", "nuclei", "foo"], {}))
        self.assertEqual(kind(self.cfg, ["build", "--no-publish", "foo"]),
                         (["slackrepo", "--repo=SBo-danix", "build", "foo"], {"SBB_PUBLISH": "n"}))
        self.assertEqual(kind(self.cfg, ["update"]), (["slackrepo", "--repo=SBo-danix", "update"], {}))

    def test_slackrepo_rejects(self):
        self.make_pkg("misc/nuclei")
        for args in (["install", "nuclei"], ["build"], ["remove"], ["build", "nope"],
                     ["build", "../etc"], ["build", "--force", "nuclei"]):
            with self.subTest(args=args), self.assertRaises(sbb.Error):
                sbb.KINDS["slackrepo"](self.cfg, args)

    def test_publish(self):
        self.assertEqual(sbb.KINDS["publish"](self.cfg, []), (["sb-publish"], {}))
        with self.assertRaises(sbb.Error):
            sbb.KINDS["publish"](self.cfg, ["x"])

    def test_rejects_bad_args(self):
        for kind, args in (("sbpull", ["evil", "--no"]), ("sbpull", ["personal"]),
                           ("sbpull", ["personal", "--force"]), ("sbshadow", []),
                           ("sbshadow", ["--yes", "--no"])):
            with self.subTest(kind=kind, args=args), self.assertRaises(sbb.Error):
                sbb.KINDS[kind](self.cfg, args)

    def test_unknown_kind(self):
        with self.assertRaisesRegex(sbb.Error, "unknown kind"):
            sbb.cmd_job_start(self.cfg, "rm", "-rf", "/")


class MainTest(Base):
    def test_unknown_command(self):
        rc, data = self.run_main("nope")
        self.assertEqual(rc, 1)
        self.assertIn("usage", data["error"])

    def test_version(self):
        self.assertEqual(self.run_main("version"),
                         (0, {"version": sbb.VERSION, "overlays": ["personal", "pentesting"]}))


class InventoryTest(Base):
    def test_rows(self):
        self.make_pkg("misc/nuclei", 'VERSION="3.11.0"\nMAINTAINER="Some One"\n'
                      'REQUIRES="go \\\n  git"\n', desc="Fast scanner")
        self.make_pkg("personal/foo", 'VERSION="1.0"\n',
                      hint='VERSION="1.1"  # bumped\nDELREQUIRES="a b"\nOPTIONS="X=$(uname -m)"\n')
        os.makedirs(os.path.join(self.cfg["repo_dir"], "misc", "notapkg"))
        db = sqlite3.connect(self.cfg["slackrepo_db"])
        db.executescript("""
            create table buildresults (itemid text primary key, time text, result text);
            create table revisions (itemid text, dep text, deplist text, version text,
                built text, rev text, os text, hintcksum text, primary key (itemid, dep));
            insert into buildresults values ('misc/nuclei', '5', 'ok'), ('nuclei', '1', 'bad');
            insert into revisions values ('misc/nuclei', '/', '/', '3.11.0', '1788978222', 'r', 'o', '/');
        """)
        db.close()
        inv = sbb.cmd_inventory(self.cfg)
        rows = [dict(zip(inv["columns"], r)) for r in inv["rows"]]
        self.assertEqual([r["path"] for r in rows], ["misc/nuclei", "personal/foo"])
        nuclei, foo = rows
        self.assertEqual((nuclei["version"], nuclei["built_version"], nuclei["built_time"],
                          nuclei["result"], nuclei["maintainer"], nuclei["requires"],
                          nuclei["summary"], nuclei["hint_path"]),
                         ("3.11.0", "3.11.0", 1788978222, "ok", "Some One", ["go", "git"],
                          "Fast scanner", None))
        self.assertEqual((foo["version"], foo["info_version"], foo["hint_version"], foo["result"]),
                         ("1.1", "1.0", "1.1", None))
        self.assertEqual(foo["hints"]["DELREQUIRES"], "a b")
        self.assertEqual(foo["hints"]["OPTIONS"], {"value": '"X=$(uname -m)"', "shell": True})
        json.dumps(inv)

    def test_hint_search_order(self):
        self.make_pkg("misc/foo")
        sub = os.path.join(self.cfg["hintdir"], "misc", "foo")
        os.makedirs(sub)
        with open(os.path.join(sub, "foo.hint"), "w") as f:
            f.write('VERSION="2"\n')
        self.assertEqual(sbb.find_hint(self.cfg, "misc/foo"), os.path.join(sub, "foo.hint"))
        with open(os.path.join(self.cfg["hintdir"], "foo.hint"), "w") as f:
            f.write('VERSION="3"\n')
        self.assertEqual(sbb.find_hint(self.cfg, "misc/foo"),
                         os.path.join(self.cfg["hintdir"], "foo.hint"))
        in_repo = os.path.join(self.cfg["repo_dir"], "misc", "foo", "foo.hint")
        open(in_repo, "w").close()
        self.assertEqual(sbb.find_hint(self.cfg, "misc/foo"), in_repo)

    def test_missing_db(self):
        self.make_pkg("misc/foo", 'VERSION="1"\n')
        row = sbb.cmd_inventory(self.cfg)["rows"][0]
        self.assertEqual(row[4:7], [None, None, None])


class GpgTest(Base):
    def setup_key(self, cached):
        with open(self.cfg["genreprc"], "w") as f:
            f.write('REPOSOWNERGPG="ABCD1234"\nUSE_GPGAGENT=1\n')
        self.fake_bin("gpg", 'case "$*" in\n'
                      '  *--list-secret-keys*) printf "sec:u:255:22:K:1::::::scESC:\\n'
                      'grp:::::::::GRIPSIGN:\\nssb:u:255:18:S:1::::::e:\\ngrp:::::::::GRIPENC:\\n" ;;\n'
                      f'  *-bas*) [ -f "{self.tmp}/good" ] ;;\n'
                      'esac\n')
        flag = "1" if cached else "-"
        self.fake_bin("gpg-connect-agent", f'echo "S KEYINFO GRIPSIGN D - - {flag} P - - -"; echo OK\n')

    def test_status(self):
        self.assertEqual(sbb.cmd_gpg_status(self.cfg), {"unlocked": None, "key": None})
        self.setup_key(cached=False)
        self.assertEqual(sbb.cmd_gpg_status(self.cfg), {"unlocked": False, "key": "ABCD1234"})
        self.setup_key(cached=True)
        self.assertEqual(sbb.cmd_gpg_status(self.cfg)["unlocked"], True)

    def test_locked_refuses_real_run_only(self):
        self.setup_key(cached=False)
        self.make_pkg("misc/foo")
        with self.assertRaisesRegex(sbb.Error, "signing key locked"):
            sbb.cmd_job_start(self.cfg, "slackrepo", "build", "foo")
        self.fake_bin("slackrepo", "exit 0\n")
        self.assertIn("id", sbb.cmd_job_start(self.cfg, "slackrepo", "build", "--preview", "foo"))

    def test_unlock(self):
        self.setup_key(cached=False)
        log = os.path.join(self.tmp, "preset.log")
        preset = self.fake_bin("preset", f'echo "$* $(cat)" >> {log}\n')
        with mock.patch.object(sbb, "PRESET", preset):
            with self.assertRaisesRegex(sbb.Error, "bad passphrase"):
                sbb.cmd_gpg_unlock(self.cfg, stdin=io.BytesIO(b"wrong\n"))
            open(os.path.join(self.tmp, "good"), "w").close()
            self.assertEqual(sbb.cmd_gpg_unlock(self.cfg, stdin=io.BytesIO(b"right\n")),
                             {"unlocked": True})
        with open(log) as f:
            self.assertEqual(f.read().splitlines(),
                             ["--preset GRIPSIGN wrong", "--forget GRIPSIGN ", "--preset GRIPSIGN right"])


class JobsTest(Base):
    def log(self, d):
        with open(os.path.join(d, "log"), "rb") as f:
            return f.read()

    def test_run_ok(self):
        jid, d = self.make_job(["sh", "-c", "echo hello"])
        sbb.cmd_run(self.cfg, jid)
        m = sbb.read_meta(d)
        self.assertEqual((m["status"], m["exit"]), ("ok", 0))
        self.assertEqual(self.log(d), b"hello\n")

    def test_run_failed(self):
        jid, d = self.make_job(["sh", "-c", "exit 3"])
        sbb.cmd_run(self.cfg, jid)
        m = sbb.read_meta(d)
        self.assertEqual((m["status"], m["exit"]), ("failed", 3))

    def test_run_missing_program(self):
        jid, d = self.make_job(["/nonexistent/prog"])
        sbb.cmd_run(self.cfg, jid)
        self.assertEqual(sbb.read_meta(d)["status"], "failed")
        self.assertIn(b"sbb-remote:", self.log(d))

    def test_kill(self):
        jid, d = self.make_job(["sleep", "30"])
        t = threading.Thread(target=sbb.cmd_run, args=(self.cfg, jid))
        t.start()
        self.wait_status(d, ("running",))
        self.assertEqual(sbb.cmd_job_kill(self.cfg, jid), {"ok": True})
        t.join(5)
        self.assertEqual(sbb.read_meta(d)["status"], "killed")

    def test_kill_not_running(self):
        jid, _ = self.make_job(["true"], status="ok")
        with self.assertRaisesRegex(sbb.Error, "not running"):
            sbb.cmd_job_kill(self.cfg, jid)

    def test_busy(self):
        self.hold_lock()
        jid, d = self.make_job(["true"])
        sbb.cmd_run(self.cfg, jid)
        self.assertEqual(sbb.read_meta(d)["status"], "busy")

    def test_lost_when_lock_free(self):
        self.make_job(["true"], status="running", pid=1)
        self.assertEqual(sbb.cmd_jobs(self.cfg)["jobs"][0]["status"], "lost")

    def test_running_kept_while_lock_held(self):
        self.hold_lock()
        self.make_job(["true"], status="running", pid=1)
        self.assertEqual(sbb.cmd_jobs(self.cfg)["jobs"][0]["status"], "running")

    def test_starting_grace(self):
        _, d = self.make_job(["true"])
        self.assertEqual(sbb.cmd_jobs(self.cfg)["jobs"][0]["status"], "starting")
        sbb.write_meta(d, dict(sbb.read_meta(d), created=time.time() - 60))
        self.assertEqual(sbb.cmd_jobs(self.cfg)["jobs"][0]["status"], "lost")

    def test_follow_from_offset(self):
        jid, d = self.make_job(["true"], status="ok")
        with open(os.path.join(d, "log"), "wb") as f:
            f.write(b"hello world")
        out = io.BytesIO()
        sbb.cmd_job_follow(self.cfg, jid, "6", out=out)
        self.assertEqual(out.getvalue(), b"world")

    def test_bad_job_id(self):
        with self.assertRaisesRegex(sbb.Error, "bad job id"):
            sbb.cmd_job_follow(self.cfg, "../../etc", "0")

    def fake_sbshadow(self):
        bindir = os.path.join(self.tmp, "bin")
        os.mkdir(bindir)
        fake = os.path.join(bindir, "sbshadow")
        with open(fake, "w") as f:
            f.write('#!/bin/sh\necho "shadow:   cat/$1"\n')
        os.chmod(fake, 0o755)
        patcher = mock.patch.dict(os.environ, {"PATH": bindir + os.pathsep + os.environ["PATH"]})
        patcher.start()
        self.addCleanup(patcher.stop)

    def test_env_reaches_child(self):
        jid, d = self.make_job(["sh", "-c", 'echo "publish=$SBB_PUBLISH"'], env={"SBB_PUBLISH": "n"})
        sbb.cmd_run(self.cfg, jid)
        self.assertEqual(self.log(d), b"publish=n\n")

    def test_start_end_to_end(self):
        self.fake_sbshadow()
        rc, data = self.run_main("job-start", "sbshadow", "--no")
        self.assertEqual(rc, 0)
        d = os.path.join(self.tmp, "jobs", data["id"])
        self.assertEqual(self.wait_status(d, ("ok", "failed", "lost"))["status"], "ok")
        self.assertEqual(self.log(d), b"shadow:   cat/--no\n")

    def test_start_busy(self):
        self.fake_sbshadow()
        self.hold_lock()
        jid, _ = self.make_job(["true"], status="running", pid=1)
        self.assertEqual(self.run_main("job-start", "sbshadow", "--no"), (0, {"busy": jid}))
        self.assertEqual(os.listdir(os.path.join(self.tmp, "jobs")), [jid])


class StatsTest(Base):
    def test_missing_sources_are_null(self):
        s = sbb.cmd_stats(self.cfg)
        self.assertIsNotNone(s["load"])
        self.assertEqual(set(s["disk"]), {"total", "free"})
        self.assertIn("slackpkg", s)
        self.assertIsNone(s["packages"])
        self.assertIsNone(s["slackrepo"])
        self.assertEqual(s["sync"], {"personal": None, "pentesting": None})
        json.dumps(s)

    def test_vm_figures(self):
        meminfo = os.path.join(self.tmp, "meminfo")
        with open(meminfo, "w") as f:
            f.write("MemTotal:       16000 kB\nMemFree:  9000 kB\nMemAvailable:   12000 kB\n")
        self.assertEqual(sbb._meminfo(meminfo), {"total": 16000 * 1024, "available": 12000 * 1024})
        with mock.patch.object(sbb, "_cpu_times", side_effect=[(100, 1000), (150, 1100)]), \
                mock.patch.object(sbb.time, "sleep"):
            self.assertEqual(sbb._cpu(), 50.0)
        stat = os.path.join(self.tmp, "stat")
        with open(stat, "w") as f:
            f.write("cpu  10 0 10 70 10 0 0 0 0 0\ncpu0 1 2 3 4\n")
        self.assertEqual(sbb._cpu_times(stat), (20, 100))  # idle + iowait count as not busy

    def test_slackpkg_times(self):
        last = os.path.join(self.tmp, "LASTUPDATE")
        with open(last, "w") as f:
            f.write("1791760212\n")
        pkgs = os.path.join(self.tmp, "packages")
        os.mkdir(pkgs)
        for name, t in (("a-1-x86_64-1", 1000), ("b-2-x86_64-1", 3000)):
            open(os.path.join(pkgs, name), "w").close()
            os.utime(os.path.join(pkgs, name), (t, t))
        self.assertEqual(sbb._slackpkg(last, pkgs), {"update": 1791760212, "upgrade": 3000})
        self.assertEqual(sbb._slackpkg("/nonexistent", "/nonexistent"), {"update": None, "upgrade": None})

    def test_counts(self):
        r = self.cfg["repo_dir"]
        for p in ("cat/foo", "cat/bar", "personal/foo", "personal/baz", "pentesting/nmap"):
            os.makedirs(os.path.join(r, p))
            open(os.path.join(r, p, os.path.basename(p) + ".SlackBuild"), "w").close()
        os.makedirs(os.path.join(r, "cat", "notapkg"))

        def fake_git(cfg, *args):
            if args[0] == "ls-tree":
                return "cat/foo/foo.SlackBuild\ncat/foo/README\ncat/old/old.SlackBuild\n"
            return "1792351223\n"

        with mock.patch.object(sbb, "_git", fake_git):
            s = sbb.cmd_stats(self.cfg)
        self.assertEqual(s["packages"], {"total": 5, "personal": 2, "pentesting": 1, "shadowed": 1})
        self.assertEqual(s["sync"]["personal"], 1792351223)

    def test_slackrepo(self):
        db = sqlite3.connect(self.cfg["slackrepo_db"])
        db.executescript("""
            create table buildresults (itemid text primary key, time text, result text);
            create table revisions (itemid text, dep text, deplist text, version text,
                built text, rev text, os text, hintcksum text, primary key (itemid, dep));
            insert into buildresults values ('cat/a', '1', 'ok'), ('cat/b', '1', 'ok'),
                ('cat/c', '1', 'failed'), ('ghost', '1', 'bad'), ('c-1.0-x86_64-1_x', '1', 'bad');
            insert into revisions values ('cat/a', '/', '/', '1', '1', 'r', 'o', '/'),
                                         ('cat/a', 'dep', '/', '1', '1', 'r', 'o', '/');
        """)
        db.close()
        # legacy keys without a category/ match no package and are not counted
        self.assertEqual(sbb.cmd_stats(self.cfg)["slackrepo"],
                         {"results": {"ok": 2, "failed": 1}, "built": 1})


class ResultsTest(Base):
    def setUp(self):
        super().setUp()
        logdir = os.path.join(self.tmp, "logs")
        self.cfg["slackrepo_logdir"] = logdir
        db = sqlite3.connect(self.cfg["slackrepo_db"])
        db.executescript("""
            create table buildresults (itemid text primary key, time text, result text);
            create table revisions (itemid text, dep text, deplist text, version text,
                built text, rev text, os text, hintcksum text, primary key (itemid, dep));
            insert into buildresults values ('cat/a', '100', 'ok'), ('cat/c', '300', 'failed'),
                                            ('ghost', '200', 'bad');
            insert into revisions values ('cat/a', '/', '/', '1.0', '100', 'r', 'o', '/');
        """)
        db.close()
        os.makedirs(os.path.join(logdir, "cat", "c"))
        with open(os.path.join(logdir, "cat", "c", "build.log"), "w") as f:
            f.write("x" * 100 + "\nERROR: make failed\n")
        with open(os.path.join(logdir, "cat", "c", "resource.log"), "w") as f:
            f.write("peak memory 3308M\n")

    def test_rows(self):
        self.assertEqual(sbb.cmd_results(self.cfg), {"rows": [
            ["cat/c", "failed", 300, None, True],
            ["cat/a", "ok", 100, "1.0", False]]})

    def test_build_log_tail(self):
        with mock.patch.object(sbb, "LOG_TAIL", 20):
            log = sbb.cmd_build_log(self.cfg, "cat/c")
        self.assertEqual(log["build"], "\nERROR: make failed\n"[-20:])
        self.assertTrue(log["truncated"])
        self.assertEqual(log["resource"], "peak memory 3308M\n")
        self.assertIsInstance(log["mtime"], int)

    def test_build_log_rejects(self):
        for bad in ("../etc/passwd", "cat/../../x", "cat", "cat/a", "/etc/x"):
            with self.subTest(bad=bad), self.assertRaises(sbb.Error):
                sbb.cmd_build_log(self.cfg, bad)


if __name__ == "__main__":
    unittest.main()