import re

from pathlib import Path

import pytest

import yaml

RULES = sorted((Path(__file__).resolve().parents[1] / "rules").rglob("*.yml"))

TACTIC_TAGS = {

    "attack.reconnaissance", "attack.resource_development", "attack.initial_access",

    "attack.execution", "attack.persistence", "attack.privilege_escalation",

    "attack.defense_evasion", "attack.credential_access", "attack.discovery",

    "attack.lateral_movement", "attack.collection", "attack.command_and_control",

    "attack.exfiltration", "attack.impact",

}

TECHNIQUE_RE = re.compile(r"^attack\.t\d{4}(\.\d{3})?$")

def load(path):

    return yaml.safe_load(path.read_text())

def test_rules_exist():

    assert RULES, "no rules found"

@pytest.mark.parametrize("path", RULES, ids=[p.name for p in RULES])

def test_rule_schema(path):

    rule = load(path)

    for field in ("title", "id", "status", "description", "tags", "level"):

        assert rule.get(field), f"{path.name} missing {field}"

    assert "detection" in rule or "correlation" in rule, f"{path.name} has no detection/correlation"

    assert rule["level"] in {"informational", "low", "medium", "high", "critical"}

    assert rule["status"] in {"experimental", "test", "stable", "deprecated", "unsupported"}

@pytest.mark.parametrize("path", RULES, ids=[p.name for p in RULES])

def test_attack_tags(path):

    tags = load(path).get("tags", [])

    techniques = [t for t in tags if TECHNIQUE_RE.match(t)]

    assert techniques, f"{path.name} has no ATT&CK technique tag"

    for t in tags:

        if t.startswith("attack."):

            assert t in TACTIC_TAGS or TECHNIQUE_RE.match(t), f"{path.name} bad tag {t}"

def test_unique_ids():

    ids = [load(p)["id"] for p in RULES]

    dupes = {i for i in ids if ids.count(i) > 1}

    assert not dupes, f"duplicate rule ids: {dupes}"

def test_correlation_refs_resolve():

    names = {load(p).get("name") for p in RULES if load(p).get("name")}

    for path in RULES:

        rule = load(path)

        corr = rule.get("correlation")

        if corr:

            for ref in corr.get("rules", []):

                assert ref in names, f"{path.name} references unknown rule '{ref}'"