# SPDX-License-Identifier: GPL-3.0-or-later
import json
import tempfile
import unittest
from pathlib import Path

from enodia_sentinel import correlation, incident
from enodia_sentinel.alert import Alert, Severity
from enodia_sentinel.config import Config


def _alert(sig: str, sid: int) -> Alert:
    return Alert(Severity.HIGH, sig, sig, sig, (4242,), sid=sid,
                 classtype=sig.rsplit(".", 1)[-1])


class TestCorrelation(unittest.TestCase):
    def test_web_rce_plus_suspicious_egress_correlates(self):
        inc = {
            "first_ts": 1000.0,
            "last_ts": 1050.0,
            "signatures": ["exec_rule.web-rce", "host_rule.suspicious-egress"],
        }
        out = correlation.correlate(inc)
        self.assertEqual(out[0]["sid"], correlation.SID_MULTI_STAGE_INTRUSION)

    def test_incident_record_persists_correlation(self):
        with tempfile.TemporaryDirectory() as d:
            cfg = Config()
            cfg.log_dir = Path(d)
            iid = incident.record(
                cfg, "alert-1.log", [_alert("exec_rule.web-rce", 100003)],
                {4242}, 1000.0, "h")
            incident.record(
                cfg, "alert-2.log", [_alert("host_rule.suspicious-egress", 100067)],
                {4242}, 1050.0, "h")
            data = json.loads((Path(d) / "incidents.json").read_text())[iid]
            self.assertEqual(data["severity"], "CRITICAL")
            self.assertIn(correlation.SID_MULTI_STAGE_INTRUSION, data["sids"])
            self.assertEqual(data["correlations"][0]["signature"],
                             "correlation.multi_stage_intrusion")


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