Add Malware object with required fields.
parent
a1dfa09fb9
commit
3e7adef792
|
@ -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.
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue