"""Tests for elway's upload ownership (2026-09-23).

elway's sudo upload used to be scp-as-user then `sudo mv`, and mv keeps the
source's owner — so every file elway installed as root landed owned by the SSH
user. These pin the fix: sudo uploads chown to root:root by default, an
explicit `owner:` overrides it, and bad input is refused before anything runs.

    python3 -m unittest scripts/test_elway.py
"""
import importlib.machinery
import importlib.util
import pathlib
import shlex
import sys
import tempfile
import textwrap
import unittest

_PATH = pathlib.Path(__file__).with_name("elway")
_loader = importlib.machinery.SourceFileLoader("elway", str(_PATH))
_spec = importlib.util.spec_from_loader("elway", _loader)
elway = importlib.util.module_from_spec(_spec)
sys.modules["elway"] = elway  # dataclasses resolve annotations via sys.modules
_loader.exec_module(elway)


class FakeCtx:
    host = "infra-ops@198.51.100.7"
    verbose = False
    sudo_password = None

    def ssh_cmd(self, remote_cmd):
        return ["ssh", self.host, remote_cmd]

    def scp_cmd(self, src, dest):
        return ["scp", src, f"{self.host}:{dest}"]


class UploadOwnershipTests(unittest.TestCase):
    def setUp(self):
        self.calls = []
        self._orig_stream = elway._stream_process
        self._orig_sudo = elway.ensure_sudo
        elway._stream_process = lambda cmd, stdin, prefix: self.calls.append(cmd) or 0
        elway.ensure_sudo = lambda ctx: None
        self.addCleanup(self._restore)
        tmp = tempfile.NamedTemporaryFile("w", suffix=".sh", delete=False)
        tmp.write("#!/bin/sh\n")
        tmp.close()
        self.src = tmp.name
        self.addCleanup(pathlib.Path(self.src).unlink)

    def _restore(self):
        elway._stream_process = self._orig_stream
        elway.ensure_sudo = self._orig_sudo

    def remote_scripts(self):
        """The shell text each ssh call runs remotely (scp calls excluded):
        for a sudo call, the script handed to `bash -c`, unwrapped."""
        out = []
        for c in self.calls:
            if c[0] != "ssh":
                continue
            words = shlex.split(c[-1])
            out.append(words[-1] if words[:2] == ["sudo", "-S"] else c[-1])
        return out

    def run_upload(self, sudo, mode=None, owner=None):
        up = elway.UploadSpec(src=self.src, dest="/etc/restic/pre-backup.sh",
                              mode=mode, owner=owner)
        step = elway.Step(name="t", upload=up, sudo=sudo)
        return elway.run_upload_step(FakeCtx(), step, "")

    def test_sudo_upload_defaults_to_root_ownership(self):
        self.assertEqual(self.run_upload(sudo=True), 0)
        (script,) = self.remote_scripts()
        self.assertIn("chown root:root /tmp/elway.", script)

    def test_sudo_upload_explicit_owner_overrides_default(self):
        self.run_upload(sudo=True, owner="booth:booth")
        (script,) = self.remote_scripts()
        self.assertIn("chown booth:booth /tmp/elway.", script)
        self.assertNotIn("root:root", script)

    def test_owner_and_mode_are_set_on_the_staged_file_before_publishing(self):
        # A chown that fails AFTER mv would leave the live path owned by the
        # SSH user — the state this exists to prevent. So chown and chmod act
        # on the staged file, and mv publishes only a fully-prepared file.
        # chown clears setuid/setgid, so chmod still comes after chown.
        self.run_upload(sudo=True, mode="0700")
        (script,) = self.remote_scripts()
        staging = self.calls[0][-1].split(":", 1)[1]   # scp dest = staging path
        chown, chmod, mv = (script.index(k) for k in (
            f"chown root:root {staging}", f"chmod 0700 {staging}", "mv -f"))
        self.assertLess(chown, chmod)
        self.assertLess(chmod, mv)
        self.assertNotIn("chown root:root /etc/restic/pre-backup.sh", script)

    def test_staged_file_is_removed_on_every_exit_path(self):
        self.run_upload(sudo=True)
        (script,) = self.remote_scripts()
        staging = self.calls[0][-1].split(":", 1)[1]
        self.assertTrue(script.startswith(f"trap 'rm -f -- {staging}' EXIT"), script)

    def test_sudo_upload_refuses_a_directory_dest(self):
        # `mv -f staged /some/dir` moves INTO the dir, and the chown/chmod that
        # follow would then hit the directory itself — e.g. dest=/home/lkraven
        # would chown the user's home to root. The remote line must refuse
        # before mv when dest is an existing directory.
        self.run_upload(sudo=True)
        (script,) = self.remote_scripts()
        guard = script.index("[ ! -d /etc/restic/pre-backup.sh ]")
        self.assertLess(guard, script.index("mv -f"))

    def test_malformed_owner_is_refused_before_anything_is_copied(self):
        # Defence in depth: even if a bad owner slipped past load-time checks,
        # nothing may be scp'd to the host first (it would orphan the staged
        # file in /tmp).
        with self.assertRaises(SystemExit):
            self.run_upload(sudo=True, owner="root; id")
        self.assertEqual(self.calls, [])

    def test_non_sudo_upload_does_not_chown(self):
        # Without sudo the file is the SSH user's by construction, and a chown
        # to anyone else would fail anyway.
        self.run_upload(sudo=False, mode="0644")
        self.assertFalse(any("chown" in s for s in self.remote_scripts()))


class PlaybookOwnerParsingTests(unittest.TestCase):
    def load(self, body):
        return elway._load_step(__import__("yaml").safe_load(textwrap.dedent(body)), "step 1")

    def test_owner_is_parsed_from_playbook(self):
        step = self.load("""
            name: x
            sudo: true
            upload: {src: a, dest: /etc/b, owner: "booth:booth"}
        """)
        self.assertEqual(step.upload.owner, "booth:booth")

    def test_owner_without_sudo_is_refused(self):
        with self.assertRaises(SystemExit):
            self.load("""
                name: x
                upload: {src: a, dest: /tmp/b, owner: "booth:booth"}
            """)

    def test_malformed_owner_is_refused(self):
        for bad in ("root; rm -rf /", "a:b:c", "", " root", "$(id)"):
            with self.subTest(owner=bad), self.assertRaises(SystemExit):
                self.load(f"""
                    name: x
                    sudo: true
                    upload: {{src: a, dest: /etc/b, owner: {bad!r}}}
                """)

    def test_user_only_and_numeric_owners_are_accepted(self):
        for ok in ("root", "1000:1000", "infra-ops:infra-ops", "_apt"):
            with self.subTest(owner=ok):
                step = self.load(f"""
                    name: x
                    sudo: true
                    upload: {{src: a, dest: /etc/b, owner: {ok!r}}}
                """)
                self.assertEqual(step.upload.owner, ok)

    def test_owner_survives_var_substitution(self):
        step = self.load("""
            name: x
            sudo: true
            upload: {src: a, dest: /etc/b, owner: "{{ svc }}:{{ svc }}"}
        """)
        new = elway.substitute_step(step, {"svc": "booth"}, "step 1")
        self.assertEqual(new.upload.owner, "booth:booth")


class ModeAndTypeTests(unittest.TestCase):
    def load(self, body):
        return elway._load_step(__import__("yaml").safe_load(textwrap.dedent(body)), "step 1")

    def test_mode_with_shell_metacharacters_is_refused(self):
        # mode is spliced into a root-run shell line; only octal digits pass.
        for bad in ("0644 && id #", "0644;id", "u+x", "", "8644", "06444"):
            with self.subTest(mode=bad), self.assertRaises(SystemExit):
                self.load(f"""
                    name: x
                    sudo: true
                    upload: {{src: a, dest: /etc/b, mode: {bad!r}}}
                """)

    def test_valid_modes_are_accepted(self):
        for ok in ("644", "0644", "0700", "4755"):
            with self.subTest(mode=ok):
                self.assertEqual(self.load(f"""
                    name: x
                    upload: {{src: a, dest: /tmp/b, mode: {ok!r}}}
                """).upload.mode, ok)

    def test_unquoted_yaml_mode_is_refused_not_reinterpreted(self):
        # YAML 1.1 reads bare 0644 as octal int 420; str() would then chmod
        # 420 (= r---w----), silently. Refuse and ask for a quoted string.
        with self.assertRaises(SystemExit):
            self.load("""
                name: x
                upload: {src: a, dest: /tmp/b, mode: 0644}
            """)

    def test_unquoted_yaml_owner_is_refused(self):
        # bare 1000:0 is sexagesimal 60000; bare yes is True — both would
        # stringify into something OWNER_RE accepts.
        for bad in ("1000:0", "yes"):
            with self.subTest(owner=bad), self.assertRaises(SystemExit):
                self.load(f"""
                    name: x
                    sudo: true
                    upload: {{src: a, dest: /etc/b, owner: {bad}}}
                """)

    def test_adhoc_bad_mode_is_refused(self):
        args = elway.parse_args(["h", "--upload", "a:/etc/b:0644;id", "--sudo"])
        with self.assertRaises(SystemExit):
            elway.build_adhoc_playbook(args)


class PreflightTests(unittest.TestCase):
    def test_templated_bad_owner_is_refused_before_any_step_runs(self):
        # Without a preflight, step 1 would already have run on the host when
        # step 2's substituted owner is found to be malformed.
        pb = elway.Playbook(vars={}, verify=[], steps=[
            elway.Step(name="stop", shell="systemctl stop x", sudo=True),
            elway.Step(name="up", sudo=True, upload=elway.UploadSpec(
                src="a", dest="/etc/b", owner="{{ o }}")),
        ])
        with self.assertRaises(SystemExit):
            elway.preflight(pb, {"o": "not a user"})

    def test_undefined_variable_is_refused_before_any_step_runs(self):
        pb = elway.Playbook(vars={}, verify=[], steps=[
            elway.Step(name="a", shell="true"),
            elway.Step(name="b", shell="echo {{ missing }}"),
        ])
        with self.assertRaises(SystemExit):
            elway.preflight(pb, {})

    def test_clean_playbook_passes_preflight(self):
        pb = elway.Playbook(vars={}, verify=[], steps=[
            elway.Step(name="up", sudo=True, upload=elway.UploadSpec(
                src="a", dest="/etc/b", owner="{{ o }}:{{ o }}"))])
        elway.preflight(pb, {"o": "booth"})

    def test_main_runs_preflight_before_opening_ssh(self):
        src = pathlib.Path(elway.__file__).read_text()
        body = src[src.index("def main("):]
        self.assertLess(body.index("preflight("), body.index("SSHContext("))


class AdhocOwnerTests(unittest.TestCase):
    def test_owner_flag_reaches_adhoc_upload(self):
        args = elway.parse_args(["h", "--upload", "a:/etc/b", "--sudo", "--owner", "booth:booth"])
        (step,) = elway.build_adhoc_playbook(args).steps
        self.assertEqual(step.upload.owner, "booth:booth")

    def test_owner_flag_outside_adhoc_upload_is_refused(self):
        # Silently ignoring --owner on --shell or --playbook would let someone
        # believe they had set an ownership that was never applied.
        for argv in (["h", "--shell", "true", "--sudo", "--owner", "x"],
                     ["h", "--playbook", "p.yaml", "--owner", "x"]):
            with self.subTest(argv=argv), self.assertRaises(SystemExit):
                elway.parse_args(argv)

    def test_owner_flag_without_sudo_is_refused(self):
        args = elway.parse_args(["h", "--upload", "a:/tmp/b", "--owner", "booth:booth"])
        with self.assertRaises(SystemExit):
            elway.build_adhoc_playbook(args)


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