cti-python-stix2/stix2/base.py

"""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 __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

        if self.granular_markings:
            for m in self.granular_markings:
                # TODO: check selectors
                pass

    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)

#  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