from xml.etree import ElementTree

from cti.models import Indicator

from cti.rules import build_cdb_lists, build_rules_xml

def sample():

    return [

        Indicator(type="ip", value="45.137.21.9", source="threatfox", threat_type="botnet_cc",

                  confidence=100, malware="Cobalt Strike", techniques=["T1071.001"]),

        Indicator(type="domain", value="cdn-jquery-min.net", source="threatfox",

                  threat_type="botnet_cc", confidence=85, malware="AgentTesla",

                  techniques=["T1056.001"]),

        Indicator(type="sha256", value="5d41402abc4b2a76b9719d911017c592e1b2c3d4f5a6978899aabbccddeeff00",

                  source="threatfox", threat_type="payload_delivery", confidence=95,

                  malware="AgentTesla", techniques=["T1204.002"]),

        Indicator(type="url", value="http://update-flashplayer.org/payload/load.php",

                  source="urlhaus", threat_type="payload_delivery", confidence=80,

                  techniques=["T1105"]),

    ]

def test_cdb_lists_have_expected_buckets():

    lists = build_cdb_lists(sample())

    assert "45.137.21.9:Cobalt Strike" in lists["cti-malicious-ip"]

    assert lists["cti-malicious-domain"].startswith("cdn-jquery-min.net")

    assert "cti-malware-hash" in lists

def test_cdb_label_is_sanitized():

    indicators = [Indicator(type="ip", value="1.2.3.4", source="x",

                            threat_type="weird:type", confidence=90)]

    lists = build_cdb_lists(indicators)

    assert ":weird-type" in lists["cti-malicious-ip"]

def test_rules_xml_is_wellformed_and_tagged():

    xml = build_rules_xml(sample(), base_id=100300)

    root = ElementTree.fromstring(xml)

    assert root.tag == "group"

    rule_ids = [r.get("id") for r in root.findall("rule")]

    assert "100300" in rule_ids

    techniques = [t.text for t in root.iter("id")]

    assert "T1071.001" in techniques

def test_rules_reference_generated_lists():

    xml = build_rules_xml(sample(), base_id=100300)

    assert "etc/lists/cti-malicious-ip" in xml

    assert "etc/lists/cti-malware-hash" in xml