import os
import re
import shlex
import stat
import subprocess
import tempfile
import unittest
from pathlib import Path


class LinpeasBuilderTests(unittest.TestCase):
    def setUp(self):
        self.repo_root = Path(__file__).resolve().parents[2]
        self.linpeas_dir = self.repo_root / "linPEAS"

    def _run_builder(self, args, output_path):
        cmd = ["python3", "-m", "builder.linpeas_builder"] + args + ["--output", str(output_path)]
        result = subprocess.run(cmd, cwd=str(self.linpeas_dir), capture_output=True, text=True)
        if result.returncode != 0:
            raise AssertionError(
                f"linpeas_builder failed:\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
            )

    def test_small_build_creates_executable(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas_small.sh"
            self._run_builder(["--small"], output_path)
            self.assertTrue(output_path.exists(), "linpeas_small.sh was not created.")
            mode = output_path.stat().st_mode
            self.assertTrue(mode & stat.S_IXUSR, "linpeas_small.sh is not executable.")

    def test_include_exclude_modules(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas_include.sh"
            self._run_builder(["--include", "system_information,container", "--exclude", "container"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            self.assertIn("Operative system", content)
            self.assertNotIn("Am I Containered?", content)

    def test_exclude_matches_module_ids_case_insensitively(self):
        """Regression: --exclude must match module IDs such as SY_Copy_Fail."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas_exclude_copyfail.sh"
            self._run_builder(
                [
                    "--include",
                    "SY_Copy_Fail",
                    "--exclude",
                    "SY_Copy_Fail,checkCopyFail",
                ],
                output_path,
            )
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            self.assertNotIn("Checking for Copy Fail", content)
            self.assertNotIn("checkCopyFail", content)

    def test_copyfail_python_probe_removes_temp_script(self):
        """Regression: cf31_run_python_probe must remove its generated probe file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = Path(tmpdir)
            marker_path = tmp_path / "probe_path"
            fake_python = tmp_path / "fake-python"
            fake_python.write_text(
                "#!/bin/sh\n"
                "printf '%s\\n' \"$1\" > \"$CF31_PROBE_PATH_FILE\"\n"
                "printf '%s\\n' probe-output\n"
                "exit 2\n",
                encoding="utf-8",
            )
            fake_python.chmod(0o755)

            function_file = self.linpeas_dir / "builder" / "linpeas_parts" / "functions" / "checkCopyFail.sh"
            script = "\n".join(
                [
                    f". {shlex.quote(str(function_file))}",
                    "CF31_PY_TIMEOUT=5",
                    f"export CF31_PROBE_PATH_FILE={shlex.quote(str(marker_path))}",
                    f"CF31_MSG=$(cf31_run_python_probe {shlex.quote(str(fake_python))})",
                    "CF31_RC=$?",
                    f"CF31_PROBE_FILE=$(cat {shlex.quote(str(marker_path))})",
                    'printf "rc=%s\\n" "$CF31_RC"',
                    'printf "msg=%s\\n" "$CF31_MSG"',
                    'printf "path=%s\\n" "$CF31_PROBE_FILE"',
                    'if [ -e "$CF31_PROBE_FILE" ]; then echo "exists=yes"; else echo "exists=no"; fi',
                ]
            )

            result = subprocess.run(
                ["sh", "-c", script],
                cwd=str(self.repo_root),
                capture_output=True,
                text=True,
            )

            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertIn("rc=2", result.stdout)
            self.assertIn("msg=probe-output", result.stdout)
            self.assertIn("exists=no", result.stdout)

    def test_threads_flag_present_in_getopts(self):
        """Regression: -z must appear in the getopts string so it is actually parsed."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas.sh"
            self._run_builder(["--all-no-fat"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            # Match the actual option-parsing line: 'while getopts' followed by
            # either a single or double quoted option string, to avoid matching
            # comments or help text that happen to contain 'getopts'.
            getopts_line = next(
                (l for l in content.splitlines()
                 if re.match(r'\s*while\s+getopts\s+[\'"]', l)),
                None
            )
            self.assertIsNotNone(getopts_line,
                                 "'while getopts' line not found in built script.")
            self.assertIn("z:", getopts_line,
                          "-z: option is missing from the getopts string in the built script.")

    def test_threads_flag_present_in_help_text(self):
        """Regression: -z must be documented in the help text of the built script."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas.sh"
            self._run_builder(["--all-no-fat"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            self.assertIn("-z <N>", content,
                          "-z <N> help entry is missing from the built script.")

    def test_mitre_flag_present_in_getopts(self):
        """The -T flag must appear in the getopts string so it is actually parsed."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas.sh"
            self._run_builder(["--all-no-fat"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            getopts_line = next(
                (l for l in content.splitlines()
                 if re.match(r'\s*while\s+getopts\s+[\'"]', l)),
                None
            )
            self.assertIsNotNone(getopts_line,
                                 "'while getopts' line not found in built script.")
            self.assertIn("T:", getopts_line,
                          "-T: option is missing from the getopts string in the built script.")

    def test_mitre_flag_present_in_help_text(self):
        """The -T flag must be documented in the help text of the built script."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas.sh"
            self._run_builder(["--all-no-fat"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            self.assertIn("-T", content,
                          "-T help entry is missing from the built script.")

    def test_mitre_filter_function_present(self):
        """check_mitre_filter() must be emitted into the built script."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "linpeas.sh"
            self._run_builder(["--all-no-fat"], output_path)
            content = output_path.read_text(encoding="utf-8", errors="ignore")
            self.assertIn("check_mitre_filter", content,
                          "check_mitre_filter function is missing from the built script.")

    def _run_base_mitre_filter(self, mitre_filter, check_ids):
        base_file = self.linpeas_dir / "builder" / "linpeas_parts" / "linpeas_base" / "0_variables_base.sh"
        result = subprocess.run(
            [
                "bash",
                "-lc",
                (
                    f'source "{base_file}" >/dev/null 2>&1 || true; '
                    f'MITRE_FILTER="{mitre_filter}"; '
                    f'check_mitre_filter "{check_ids}"; '
                    'echo $?'
                ),
            ],
            capture_output=True,
            text=True,
            cwd=str(self.repo_root),
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        return result.stdout.strip().splitlines()[-1]

    def test_mitre_parent_filter_matches_subtechnique(self):
        """Regression: filtering by a base technique must include child sub-techniques."""
        exit_code = self._run_base_mitre_filter("T1552", "T1552.001")
        self.assertEqual("0", exit_code,
                         "Parent MITRE filter T1552 should match sub-technique T1552.001.")

    def test_mitre_subtechnique_filter_does_not_match_parent(self):
        """Regression: filtering by a sub-technique must not include a parent-only tag."""
        exit_code = self._run_base_mitre_filter("T1552.001", "T1552")
        self.assertEqual("1", exit_code,
                         "Sub-technique MITRE filter T1552.001 should not match parent tag T1552.")


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