From 3e7adef792d586ed027011ba2d6d0906ee03ef8e Mon Sep 17 00:00:00 2001 From: Greg Back Date: Tue, 17 Jan 2017 16:53:27 -0800 Subject: [PATCH] Add Malware object with required fields. --- stix2/__init__.py | 66 +++++++++++++++++++++++++-- stix2/test/test_stix2.py | 99 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 156 insertions(+), 9 deletions(-) diff --git a/stix2/__init__.py b/stix2/__init__.py index c1f3feb..dd24dc9 100644 --- a/stix2/__init__.py +++ b/stix2/__init__.py @@ -7,14 +7,19 @@ import pytz def format_datetime(dt): + # TODO: how to handle naive datetime + # 1. Convert to UTC # 2. Format in isoformat # 3. Strip off "+00:00" # 4. Add "Z" + + # TODO: how to handle timestamps with subsecond 0's return dt.astimezone(pytz.utc).isoformat()[:-6] + "Z" -class Indicator(collections.Mapping): +class _STIXBase(collections.Mapping): + """Base class for STIX object types""" def __getitem__(self, key): return self._inner[key] @@ -32,7 +37,10 @@ class Indicator(collections.Mapping): def __setattr__(self, name, value): if name != '_inner': raise ValueError("Cannot modify properties after creation.") - super(Indicator, self).__setattr__(name, value) + super(_STIXBase, self).__setattr__(name, value) + + +class Indicator(_STIXBase): def __init__(self, type='indicator', id=None, created=None, modified=None, labels=None, pattern=None, valid_from=None): @@ -67,7 +75,7 @@ class Indicator(collections.Mapping): self._inner = { 'type': type, - 'id': id or now, + 'id': id, 'created': created or now, 'modified': modified or now, 'labels': labels, @@ -86,3 +94,55 @@ class Indicator(collections.Mapping): 'pattern': self['pattern'], 'valid_from': format_datetime(self['valid_from']), }, indent=4, sort_keys=True, separators=(",", ": ")) # Don't include spaces after commas. + + +class Malware(_STIXBase): + + def __init__(self, type='malware', id=None, created=None, modified=None, + labels=None, name=None): + # TODO: + # - created_by_ref + # - revoked + # - external_references + # - object_marking_refs + # - granular_markings + + # - description + # - kill_chain_phases + + if not created or not modified: + now = datetime.now(tz=pytz.UTC) + + if type != 'malware': + raise ValueError("Malware must have type='malware'.") + + if not id: + id = 'malware--' + str(uuid.uuid4()) + if not id.startswith('malware--'): + raise ValueError("Malware id values must begin with 'malware--'.") + + if not labels: + raise ValueError("Missing required field for Malware: 'labels'.") + + if not name: + raise ValueError("Missing required field for Malware: 'name'.") + + self._inner = { + 'type': type, + 'id': id, + 'created': created or now, + 'modified': modified or now, + 'labels': labels, + 'name': name, + } + + def __str__(self): + # TODO: put keys in specific order. Probably need custom JSON encoder. + return json.dumps({ + 'type': self['type'], + 'id': self['id'], + 'created': format_datetime(self['created']), + 'modified': format_datetime(self['modified']), + 'labels': self['labels'], + 'name': self['name'], + }, indent=4, sort_keys=True, separators=(",", ": ")) # Don't include spaces after commas. diff --git a/stix2/test/test_stix2.py b/stix2/test/test_stix2.py index ab45602..0bb841f 100644 --- a/stix2/test/test_stix2.py +++ b/stix2/test/test_stix2.py @@ -21,7 +21,7 @@ def test_timestamp_formatting(dt, timestamp): assert stix2.format_datetime(dt) == timestamp -EXPECTED = """{ +EXPECTED_INDICATOR = """{ "created": "2017-01-01T00:00:00Z", "id": "indicator--01234567-89ab-cdef-0123-456789abcdef", "labels": [ @@ -48,18 +48,18 @@ def test_indicator_with_all_required_fields(): valid_from=epoch, ) - assert str(indicator) == EXPECTED + assert str(indicator) == EXPECTED_INDICATOR -# Minimum required args for an indicator -KWARGS = dict( +# Minimum required args for an Indicator instance +INDICATOR_KWARGS = dict( labels=['malicious-activity'], pattern="[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']", ) def test_indicator_autogenerated_fields(): - indicator = stix2.Indicator(**KWARGS) + indicator = stix2.Indicator(**INDICATOR_KWARGS) assert indicator.type == 'indicator' assert indicator.id.startswith('indicator--') assert indicator.created is not None @@ -105,9 +105,96 @@ def test_indicator_required_field_pattern(): def test_cannot_assign_to_attributes(): - indicator = stix2.Indicator(**KWARGS) + indicator = stix2.Indicator(**INDICATOR_KWARGS) with pytest.raises(ValueError) as excinfo: indicator.valid_from = datetime.datetime.now() assert "Cannot modify properties after creation." in str(excinfo) + + +EXPECTED_MALWARE = """{ + "created": "2016-05-12T08:17:27Z", + "id": "malware--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061", + "labels": [ + "ransomware" + ], + "modified": "2016-05-12T08:17:27Z", + "name": "Cryptolocker", + "type": "malware" +}""" + + +def test_malware_with_all_required_fields(): + now = datetime.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc) + + malware = stix2.Malware( + type="malware", + id="malware--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061", + created=now, + modified=now, + labels=["ransomware"], + name="Cryptolocker", + ) + + assert str(malware) == EXPECTED_MALWARE + + +# Minimum required args for a Malware instance +MALWARE_KWARGS = dict( + labels=['ransomware'], + name="Cryptolocker", +) + + +def test_malware_autogenerated_fields(): + malware = stix2.Malware(**MALWARE_KWARGS) + assert malware.type == 'malware' + assert malware.id.startswith('malware--') + assert malware.created is not None + assert malware.modified is not None + assert malware.labels == ['ransomware'] + assert malware.name == "Cryptolocker" + + assert malware['type'] == 'malware' + assert malware['id'].startswith('malware--') + assert malware['created'] is not None + assert malware['modified'] is not None + assert malware['labels'] == ['ransomware'] + assert malware['name'] == "Cryptolocker" + + +def test_malware_type_must_be_malware(): + with pytest.raises(ValueError) as excinfo: + malware = stix2.Malware(type='xxx') + + assert "Malware must have type='malware'." in str(excinfo) + + +def test_malware_id_must_start_with_malware(): + with pytest.raises(ValueError) as excinfo: + malware = stix2.Malware(id='my-prefix--') + + assert "Malware id values must begin with 'malware--'." in str(excinfo) + + +def test_malware_required_field_labels(): + with pytest.raises(ValueError) as excinfo: + malware = stix2.Malware() + assert "Missing required field for Malware: 'labels'." in str(excinfo) + + +def test_malware_required_field_name(): + with pytest.raises(ValueError) as excinfo: + # Label is checked first, so make sure that is provided + malware = stix2.Malware(labels=['ransomware']) + assert "Missing required field for Malware: 'name'." in str(excinfo) + + +def test_cannot_assign_to_attributes(): + malware = stix2.Malware(**MALWARE_KWARGS) + + with pytest.raises(ValueError) as excinfo: + malware.name = "Cryptolocker II" + + assert "Cannot modify properties after creation." in str(excinfo)