# SPDX-License-Identifier: GPL-3.0-or-later
"""Tests for post-alert snapshot enrichment."""
import tempfile
import unittest
from dataclasses import dataclass
from pathlib import Path

from enodia_sentinel import enrich
from enodia_sentinel.config import Config
from enodia_sentinel.system import SystemState


@dataclass
class FakeProc:
    pid: int
    comm: str
    exe: str = ""
    ppid: int = 0
    uid: int = 0


class TestEnrichment(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.TemporaryDirectory()
        self.tmp = Path(self.dir.name)
        self.cfg = Config()
        self.cfg.log_dir = self.tmp

    def tearDown(self):
        self.dir.cleanup()

    def test_process_remote_file_and_integrity_context(self):
        exe = self.tmp / "payload"
        exe.write_text("stage-one")
        watched = self.tmp / "persistence.service"
        watched.write_text("[Service]\nExecStart=/tmp/payload\n")
        self.cfg.fim_baseline.write_text("{}")
        (self.tmp / "pkgdb-anchor.json").write_text("{}")

        state = SystemState(processes=[
            FakeProc(1, "systemd"),
            FakeProc(10, "nginx", ppid=1),
            FakeProc(42, "payload", exe=str(exe), ppid=10),
        ])
        report = {
            "alerts": [
                {"signature": "reverse_shell",
                 "detail": "pid=42 peer=[8.8.8.8:4444]",
                 "key": "rsh:42"},
                {"signature": "persistence",
                 "detail": f"persistence file modified: {watched}",
                 "key": "k:persistence"},
            ],
            "processes": [
                {"pid": 42, "comm": "payload", "exe": str(exe), "ppid": 10},
            ],
        }

        out = enrich.build(
            report, state, self.cfg, lineage={42, 10, 1},
            recent_files=[str(watched)],
            owner=lambda p: "testpkg 1.0-1" if p == str(exe) else None,
        )

        self.assertEqual(out["processes"][0]["package_owner"], "testpkg 1.0-1")
        self.assertRegex(out["processes"][0]["exe_sha256"], r"^[0-9a-f]{64}$")
        self.assertEqual(out["processes"][0]["parent_chain"][0]["pid"], 10)
        self.assertEqual(out["remotes"][0],
                         {"ip": "8.8.8.8", "classification": "public"})
        files = {f["path"]: f for f in out["files"]}
        self.assertTrue(files[str(watched)]["exists"])
        self.assertIn(str(watched), out["recent_writes"])
        self.assertEqual(out["integrity"]["fim_baseline"]["status"], "present")
        self.assertEqual(out["integrity"]["pkgdb_anchor"]["status"], "present")

    def test_private_remote_and_missing_file_context(self):
        state = SystemState(processes=[])
        report = {
            "alerts": [
                {"signature": "egress",
                 "detail": "python connected to 10.1.2.3:443 and /gone",
                 "key": "k"},
            ],
            "processes": [],
        }
        out = enrich.build(report, state, self.cfg, owner=lambda _p: None)
        self.assertEqual(out["remotes"][0]["classification"], "private")
        self.assertFalse(out["files"][0]["exists"])

    def test_text_summary_is_compact(self):
        text = enrich.format_text({
            "processes": [{"pid": 42, "comm": "payload",
                           "package_owner": "pkg", "exe_sha256": "a" * 64,
                           "parent_chain": [{"pid": 10, "comm": "nginx"}]}],
            "remotes": [{"ip": "8.8.8.8", "classification": "public"}],
            "files": [],
            "integrity": {"fim_baseline": {"status": "present"},
                          "pkgdb_anchor": {"status": "missing"},
                          "rootcheck": {"enabled": True}},
        })
        self.assertIn("process pid=42", text)
        self.assertIn("remote 8.8.8.8: public", text)
        self.assertIn("integrity:", text)


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