Allow attribute and key-based access. Make immutable.
parent
eeec5a4ce3
commit
ef0b80ad44
|
@ -1,3 +1,4 @@
|
|||
import collections
|
||||
from datetime import datetime
|
||||
import json
|
||||
import uuid
|
||||
|
@ -12,58 +13,76 @@ def format_datetime(dt):
|
|||
# 4. Add "Z"
|
||||
return dt.astimezone(pytz.utc).isoformat()[:-6] + "Z"
|
||||
|
||||
# REQUIRED (all):
|
||||
# - type
|
||||
# - id
|
||||
# - created
|
||||
# - modified
|
||||
|
||||
|
||||
class Indicator(object):
|
||||
# REQUIRED (Indicator):
|
||||
# - type
|
||||
# - labels
|
||||
# - pattern
|
||||
# - valid_from
|
||||
required = ['']
|
||||
class Indicator(collections.Mapping):
|
||||
|
||||
def __getitem__(self, key):
|
||||
return getattr(self, key)
|
||||
return self._inner[key]
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._inner)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._inner)
|
||||
|
||||
# Handle attribute access just like key access
|
||||
def __getattr__(self, name):
|
||||
return self.get(name)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
if name != '_inner':
|
||||
raise ValueError("Cannot modify properties after creation.")
|
||||
super(Indicator, self).__setattr__(name, value)
|
||||
|
||||
def __init__(self, type='indicator', id=None, created=None, modified=None,
|
||||
labels=None, pattern=None, valid_from=None):
|
||||
now = datetime.now(tz=pytz.UTC)
|
||||
# TODO:
|
||||
# - created_by_ref
|
||||
# - revoked
|
||||
# - external_references
|
||||
# - object_marking_refs
|
||||
# - granular_markings
|
||||
|
||||
# - name
|
||||
# - description
|
||||
# - valid_until
|
||||
# - kill_chain_phases
|
||||
|
||||
if not created or not modified or not valid_from:
|
||||
now = datetime.now(tz=pytz.UTC)
|
||||
|
||||
if type != 'indicator':
|
||||
raise ValueError("Indicators must have type='indicator'.")
|
||||
self.type = type
|
||||
|
||||
if not id:
|
||||
id = 'indicator--' + str(uuid.uuid4())
|
||||
if not id.startswith('indicator--'):
|
||||
raise ValueError("Indicator id values must begin with 'indicator--'.")
|
||||
self.id = id
|
||||
|
||||
self.created = created or now
|
||||
self.modified = modified or now
|
||||
|
||||
if not labels:
|
||||
raise ValueError("Missing required field for Indicator: 'labels'.")
|
||||
self.labels = labels
|
||||
|
||||
if not pattern:
|
||||
raise ValueError("Missing required field for Indicator: 'pattern'.")
|
||||
self.pattern = pattern
|
||||
self.valid_from = valid_from or now
|
||||
|
||||
self._inner = {
|
||||
'type': type,
|
||||
'id': id or now,
|
||||
'created': created or now,
|
||||
'modified': modified or now,
|
||||
'labels': labels,
|
||||
'pattern': pattern,
|
||||
'valid_from': valid_from or now,
|
||||
}
|
||||
|
||||
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,
|
||||
'pattern': self.pattern,
|
||||
'valid_from': format_datetime(self.valid_from),
|
||||
'type': self['type'],
|
||||
'id': self['id'],
|
||||
'created': format_datetime(self['created']),
|
||||
'modified': format_datetime(self['modified']),
|
||||
'labels': self['labels'],
|
||||
'pattern': self['pattern'],
|
||||
'valid_from': format_datetime(self['valid_from']),
|
||||
}, indent=4, sort_keys=True, separators=(",", ": ")) # Don't include spaces after commas.
|
||||
|
|
|
@ -102,3 +102,12 @@ def test_indicator_required_field_pattern():
|
|||
# Label is checked first, so make sure that is provided
|
||||
indicator = stix2.Indicator(labels=['malicious-activity'])
|
||||
assert "Missing required field for Indicator: 'pattern'." in str(excinfo)
|
||||
|
||||
|
||||
def test_cannot_assign_to_attributes():
|
||||
indicator = stix2.Indicator(**KWARGS)
|
||||
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
indicator.valid_from = datetime.datetime.now()
|
||||
|
||||
assert "Cannot modify properties after creation." in str(excinfo)
|
||||
|
|
Loading…
Reference in New Issue