Add Malware object with required fields.

stix2.1
Greg Back 2017-01-17 16:53:27 -08:00
parent a1dfa09fb9
commit 3e7adef792
2 changed files with 156 additions and 9 deletions

View File

@ -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.

View File

@ -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)