Allow attribute and key-based access. Make immutable.

stix2.1
Greg Back 2017-01-17 15:52:03 -08:00
parent eeec5a4ce3
commit ef0b80ad44
2 changed files with 58 additions and 30 deletions

View File

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

View File

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