cti-python-stix2/stix2/test/test_malware.py

import datetime as dt
import re

import pytest
import pytz

import stix2

from .constants import FAKE_TIME, MALWARE_ID, MALWARE_KWARGS


EXPECTED_MALWARE = """{
    "created": "2016-05-12T08:17:27Z",
    "id": "malware--fedcba98-7654-3210-fedc-ba9876543210",
    "labels": [
        "ransomware"
    ],
    "modified": "2016-05-12T08:17:27Z",
    "name": "Cryptolocker",
    "type": "malware"
}"""


def test_malware_with_all_required_properties():
    now = dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)

    mal = stix2.Malware(
        type="malware",
        id=MALWARE_ID,
        created=now,
        modified=now,
        labels=["ransomware"],
        name="Cryptolocker",
    )

    assert str(mal) == EXPECTED_MALWARE


def test_malware_autogenerated_properties(malware):
    assert malware.type == 'malware'
    assert malware.id == 'malware--00000000-0000-0000-0000-000000000001'
    assert malware.created == FAKE_TIME
    assert malware.modified == FAKE_TIME
    assert malware.labels == ['ransomware']
    assert malware.name == "Cryptolocker"

    assert malware['type'] == 'malware'
    assert malware['id'] == 'malware--00000000-0000-0000-0000-000000000001'
    assert malware['created'] == FAKE_TIME
    assert malware['modified'] == FAKE_TIME
    assert malware['labels'] == ['ransomware']
    assert malware['name'] == "Cryptolocker"


def test_malware_type_must_be_malware():
    with pytest.raises(stix2.exceptions.InvalidValueError) as excinfo:
        stix2.Malware(type='xxx', **MALWARE_KWARGS)

    assert excinfo.value.cls == stix2.Malware
    assert excinfo.value.prop_name == "type"
    assert excinfo.value.reason == "must equal 'malware'."
    assert str(excinfo.value) == "Invalid value for Malware 'type': must equal 'malware'."


def test_malware_id_must_start_with_malware():
    with pytest.raises(stix2.exceptions.InvalidValueError) as excinfo:
        stix2.Malware(id='my-prefix--', **MALWARE_KWARGS)

    assert excinfo.value.cls == stix2.Malware
    assert excinfo.value.prop_name == "id"
    assert excinfo.value.reason == "must start with 'malware--'."
    assert str(excinfo.value) == "Invalid value for Malware 'id': must start with 'malware--'."


def test_malware_required_properties():
    with pytest.raises(stix2.exceptions.MissingFieldsError) as excinfo:
        stix2.Malware()

    assert excinfo.value.cls == stix2.Malware
    assert excinfo.value.properties == ["labels", "name"]


def test_malware_required_property_name():
    with pytest.raises(stix2.exceptions.MissingFieldsError) as excinfo:
        stix2.Malware(labels=['ransomware'])

    assert excinfo.value.cls == stix2.Malware
    assert excinfo.value.properties == ["name"]


def test_cannot_assign_to_malware_attributes(malware):
    with pytest.raises(stix2.exceptions.ImmutableError) as excinfo:
        malware.name = "Cryptolocker II"

    assert str(excinfo.value) == "Cannot modify properties after creation."


def test_invalid_kwarg_to_malware():
    with pytest.raises(stix2.exceptions.ExtraFieldsError) as excinfo:
        stix2.Malware(my_custom_property="foo", **MALWARE_KWARGS)

    assert excinfo.value.cls == stix2.Malware
    assert excinfo.value.properties == ['my_custom_property']
    assert str(excinfo.value) == "Unexpected properties for Malware: (my_custom_property)."


@pytest.mark.parametrize("data", [
    EXPECTED_MALWARE,
    {
        "type": "malware",
        "id": "malware--fedcba98-7654-3210-fedc-ba9876543210",
        "created": "2016-05-12T08:17:27Z",
        "modified": "2016-05-12T08:17:27Z",
        "labels": ["ransomware"],
        "name": "Cryptolocker",
    },
])
def test_parse_malware(data):
    mal = stix2.parse(data)

    assert mal.type == 'malware'
    assert mal.id == MALWARE_ID
    assert mal.created == dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)
    assert mal.modified == dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)
    assert mal.labels == ['ransomware']
    assert mal.name == "Cryptolocker"


def test_parse_malware_invalid_labels():
    data = re.compile('\[.+\]', re.DOTALL).sub('1', EXPECTED_MALWARE)
    with pytest.raises(ValueError) as excinfo:
        stix2.parse(data)
    assert "Invalid value for Malware 'labels'" in str(excinfo.value)


def test_parse_malware_kill_chain_phases():
    kill_chain = """
      "kill_chain_phases": [
        {
          "kill_chain_name": "lockheed-martin-cyber-kill-chain",
          "phase_name": "reconnaissance"
        }
      ]"""
    data = EXPECTED_MALWARE.replace('malware"', 'malware",%s' % kill_chain)
    mal = stix2.parse(data)
    assert mal.kill_chain_phases[0].kill_chain_name == "lockheed-martin-cyber-kill-chain"
    assert mal.kill_chain_phases[0].phase_name == "reconnaissance"
    assert mal['kill_chain_phases'][0]['kill_chain_name'] == "lockheed-martin-cyber-kill-chain"
    assert mal['kill_chain_phases'][0]['phase_name'] == "reconnaissance"


def test_parse_malware_clean_kill_chain_phases():
    kill_chain = """
      "kill_chain_phases": [
        {
          "kill_chain_name": "lockheed-martin-cyber-kill-chain",
          "phase_name": 1
        }
      ]"""
    data = EXPECTED_MALWARE.replace('malware"', 'malware",%s' % kill_chain)
    mal = stix2.parse(data)
    assert mal['kill_chain_phases'][0]['phase_name'] == "1"