cti-python-stix2/stix2/base.py

173 lines
6.3 KiB
Python

"""Base class for type definitions in the stix2 library."""
import collections
import copy
import datetime as dt
import json
from .exceptions import (ExtraFieldsError, ImmutableError, InvalidObjRefError,
InvalidValueError, MissingFieldsError, RevokeError,
UnmodifiablePropertyError)
from .utils import NOW, format_datetime, get_timestamp, parse_into_datetime
__all__ = ['STIXJSONEncoder', '_STIXBase']
DEFAULT_ERROR = "{type} must have {field}='{expected}'."
class STIXJSONEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, (dt.date, dt.datetime)):
return format_datetime(obj)
elif isinstance(obj, _STIXBase):
return dict(obj)
else:
return super(STIXJSONEncoder, self).default(obj)
def get_required_properties(properties):
return (k for k, v in properties.items() if v.required)
class _STIXBase(collections.Mapping):
"""Base class for STIX object types"""
def _check_property(self, prop_name, prop, kwargs):
if prop_name not in kwargs:
if hasattr(prop, 'default'):
value = prop.default()
if value == NOW:
value = self.__now
kwargs[prop_name] = value
if prop_name in kwargs:
try:
kwargs[prop_name] = prop.clean(kwargs[prop_name])
except ValueError as exc:
raise InvalidValueError(self.__class__, prop_name, reason=str(exc))
def _check_object_constaints(self):
if self.granular_markings:
for m in self.granular_markings:
# TODO: check selectors
pass
def __init__(self, **kwargs):
cls = self.__class__
# Use the same timestamp for any auto-generated datetimes
self.__now = get_timestamp()
# Detect any keyword arguments not allowed for a specific type
extra_kwargs = list(set(kwargs) - set(cls._properties))
if extra_kwargs:
raise ExtraFieldsError(cls, extra_kwargs)
# Remove any keyword arguments whose value is None
setting_kwargs = {}
for prop_name, prop_value in kwargs.items():
if prop_value:
setting_kwargs[prop_name] = prop_value
# Detect any missing required fields
required_fields = get_required_properties(cls._properties)
missing_kwargs = set(required_fields) - set(setting_kwargs)
if missing_kwargs:
raise MissingFieldsError(cls, missing_kwargs)
for prop_name, prop_metadata in cls._properties.items():
self._check_property(prop_name, prop_metadata, setting_kwargs)
self._inner = setting_kwargs
self._check_object_constaints()
def __getitem__(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' and not name.startswith("_STIXBase__"):
raise ImmutableError
super(_STIXBase, self).__setattr__(name, value)
def __str__(self):
# TODO: put keys in specific order. Probably need custom JSON encoder.
return json.dumps(self, indent=4, sort_keys=True, cls=STIXJSONEncoder,
separators=(",", ": ")) # Don't include spaces after commas.
def __repr__(self):
props = [(k, self[k]) for k in sorted(self._properties) if self.get(k)]
return "{0}({1})".format(self.__class__.__name__,
", ".join(["{0!s}={1!r}".format(k, v) for k, v in props]))
def __deepcopy__(self, memo):
# Assumption: we can ignore the memo argument, because no object will ever contain the same sub-object multiple times.
new_inner = copy.deepcopy(self._inner, memo)
cls = type(self)
return cls(**new_inner)
def properties_populated(self):
return list(self._inner.keys())
# Versioning API
def new_version(self, **kwargs):
unchangable_properties = []
if self.revoked:
raise RevokeError("new_version")
new_obj_inner = copy.deepcopy(self._inner)
properties_to_change = kwargs.keys()
for prop in ["created", "created_by_ref", "id", "type"]:
if prop in properties_to_change:
unchangable_properties.append(prop)
if unchangable_properties:
raise UnmodifiablePropertyError(unchangable_properties)
cls = type(self)
if 'modified' not in kwargs:
kwargs['modified'] = get_timestamp()
else:
new_modified_property = parse_into_datetime(kwargs['modified'])
if new_modified_property < self.modified:
raise InvalidValueError(cls, 'modified', "The new modified datetime cannot be before the current modified datatime.")
new_obj_inner.update(kwargs)
return cls(**new_obj_inner)
def revoke(self):
if self.revoked:
raise RevokeError("revoke")
return self.new_version(revoked=True)
class _Observable(_STIXBase):
def __init__(self, **kwargs):
# the constructor might be called independently of an observed data object
if '_valid_refs' in kwargs:
self._STIXBase__valid_refs = kwargs.pop('_valid_refs')
else:
self._STIXBase__valid_refs = []
super(_Observable, self).__init__(**kwargs)
def _check_property(self, prop_name, prop, kwargs):
super(_Observable, self)._check_property(prop_name, prop, kwargs)
if prop_name.endswith('_ref') and prop_name in kwargs:
ref = kwargs[prop_name]
if ref not in self._STIXBase__valid_refs:
raise InvalidObjRefError(self.__class__, prop_name, "'%s' is not a valid object in local scope" % ref)
elif prop_name.endswith('_refs') and prop_name in kwargs:
for ref in kwargs[prop_name]:
if ref not in self._STIXBase__valid_refs:
raise InvalidObjRefError(self.__class__, prop_name, "'%s' is not a valid object in local scope" % ref)
# TODO also check the type of the object referenced, not just that the key exists