<commit_before>"""
Python STIX 2.0 Memory Source/Sink
"""

import itertools
import json
import os

from stix2.base import _STIXBase
from stix2.core import Bundle, parse
from stix2.datastore import DataSink, DataSource, DataStoreMixin
from stix2.datastore.filters import FilterSet, apply_common_filters


def _add(store, stix_data=None, allow_custom=True, version=None):
    """Add STIX objects to MemoryStore/Sink.

    Adds STIX objects to an in-memory dictionary for fast lookup.
    Recursive function, breaks down STIX Bundles and lists.

    Args:
        stix_data (list OR dict OR STIX object): STIX objects to be added
        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
            None, use latest version.

    """
    if isinstance(stix_data, list):
        # STIX objects are in a list- recurse on each object
        for stix_obj in stix_data:
            _add(store, stix_obj, allow_custom, version)

    elif stix_data["type"] == "bundle":
        # adding a json bundle - so just grab STIX objects
        for stix_obj in stix_data.get("objects", []):
            _add(store, stix_obj, allow_custom, version)

    else:
        # Adding a single non-bundle object
        if isinstance(stix_data, _STIXBase):
            stix_obj = stix_data
        else:
            stix_obj = parse(stix_data, allow_custom, version)

        # Map ID directly to the object, if it is a marking.  Otherwise,
        # map to a family, so we can track multiple versions.
        if _is_marking(stix_obj):
            store._data[stix_obj["id"]] = stix_obj

        else:
            if stix_obj.id in store._data:
                obj_family = store._data[stix_obj["id"]]
            else:
                obj_family = _ObjectFamily()
                store._data[stix_obj["id"]] = obj_family

            obj_family.add(stix_obj)


def _is_marking(obj_or_id):
    """Determines whether the given object or object ID is/is for a marking
    definition.

    :param obj_or_id: A STIX object or object ID as a string.
    :return: True if a marking definition, False otherwise.
    """

    if isinstance(obj_or_id, (_STIXBase, dict)):
        id_ = obj_or_id["id"]
    else:
        id_ = obj_or_id

    return id_.startswith("marking-definition--")


class _ObjectFamily(object):
    """
    An internal implementation detail of memory sources/sinks/stores.
    Represents a "family" of STIX objects: all objects with a particular
    ID.  (I.e. all versions.)  The latest version is also tracked so that it
    can be obtained quickly.
    """
    def __init__(self):
        self.all_versions = {}
        self.latest_version = None

    def add(self, obj):
        self.all_versions[obj.modified] = obj
        if self.latest_version is None or \
                obj.modified > self.latest_version.modified:
            self.latest_version = obj

    def __str__(self):
        return "<<{}; latest={}>>".format(self.all_versions,
                                          self.latest_version.modified)

    def __repr__(self):
        return str(self)


class MemoryStore(DataStoreMixin):
    """Interface to an in-memory dictionary of STIX objects.

    MemoryStore is a wrapper around a paired MemorySink and MemorySource.

    Note: It doesn't make sense to create a MemoryStore by passing
    in existing MemorySource and MemorySink because there could
    be data concurrency issues. As well, just as easy to create new MemoryStore.

    Args:
        stix_data (list OR dict OR STIX object): STIX content to be added
        allow_custom (bool): whether to allow custom STIX content.
            Only applied when export/input functions called, i.e.
            load_from_file() and save_to_file(). Defaults to True.
        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
            None, use latest version.

    Attributes:
        _data (dict): the in-memory dict that holds STIX objects
        source (MemorySource): MemorySource
        sink (MemorySink): MemorySink

    """
    def __init__(self, stix_data=None, allow_custom=True, version=None):
        self._data = {}

        if stix_data:
            _add(self, stix_data, allow_custom, version=version)

        super(MemoryStore, self).__init__(
            source=MemorySource(stix_data=self._data, allow_custom=allow_custom, version=version, _store=True),
            sink=MemorySink(stix_data=self._data, allow_custom=allow_custom, version=version, _store=True)
        )

    def save_to_file(self, *args, **kwargs):
        """Write SITX objects from in-memory dictionary to JSON file, as a STIX
        Bundle.

        Args:
            file_path (str): file path to write STIX data to

        """
        return self.sink.save_to_file(*args, **kwargs)

    def load_from_file(self, *args, **kwargs):
        """Load STIX data from JSON file.

        File format is expected to be a single JSON
        STIX object or JSON STIX bundle.

        Args:
            file_path (str): file path to load STIX data from
            version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
                None, use latest version.

        """
        return self.source.load_from_file(*args, **kwargs)


class MemorySink(DataSink):
    """Interface for adding/pushing STIX objects to an in-memory dictionary.

    Designed to be paired with a MemorySource, together as the two
    components of a MemoryStore.

    Args:
        stix_data (dict OR list): valid STIX 2.0 content in
            bundle or a list.
        _store (bool): whether the MemorySink is a part of a MemoryStore,
            in which case "stix_data" is a direct reference to
            shared memory with DataSource. Not user supplied
        allow_custom (bool): whether to allow custom objects/properties
            when exporting STIX content to file.
            Default: True.

    Attributes:
        _data (dict): the in-memory dict that holds STIX objects.
            If part of a MemoryStore, the dict is shared with a MemorySource

    """
    def __init__(self, stix_data=None, allow_custom=True, version=None, _store=False):
        super(MemorySink, self).__init__()
        self.allow_custom = allow_custom

        if _store:
            self._data = stix_data
        else:
            self._data = {}
            if stix_data:
                _add(self, stix_data, allow_custom, version=version)

    def add(self, stix_data, version=None):
        _add(self, stix_data, self.allow_custom, version)
    add.__doc__ = _add.__doc__

    def save_to_file(self, file_path):
        file_path = os.path.abspath(file_path)

        all_objs = itertools.chain.from_iterable(
            value.all_versions.values() if isinstance(value, _ObjectFamily)
            else [value]
            for value in self._data.values()
        )

        if not os.path.exists(os.path.dirname(file_path)):
            os.makedirs(os.path.dirname(file_path))
        with open(file_path, "w") as f:<commit_msg>Missing changes to key/value approach<commit_after>"""
Python STIX 2.0 Memory Source/Sink
"""

import itertools
import json
import os

from stix2.base import _STIXBase
from stix2.core import Bundle, parse
from stix2.datastore import DataSink, DataSource, DataStoreMixin
from stix2.datastore.filters import FilterSet, apply_common_filters


def _add(store, stix_data=None, allow_custom=True, version=None):
    """Add STIX objects to MemoryStore/Sink.

    Adds STIX objects to an in-memory dictionary for fast lookup.
    Recursive function, breaks down STIX Bundles and lists.

    Args:
        stix_data (list OR dict OR STIX object): STIX objects to be added
        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
            None, use latest version.

    """
    if isinstance(stix_data, list):
        # STIX objects are in a list- recurse on each object
        for stix_obj in stix_data:
            _add(store, stix_obj, allow_custom, version)

    elif stix_data["type"] == "bundle":
        # adding a json bundle - so just grab STIX objects
        for stix_obj in stix_data.get("objects", []):
            _add(store, stix_obj, allow_custom, version)

    else:
        # Adding a single non-bundle object
        if isinstance(stix_data, _STIXBase):
            stix_obj = stix_data
        else:
            stix_obj = parse(stix_data, allow_custom, version)

        # Map ID directly to the object, if it is a marking.  Otherwise,
        # map to a family, so we can track multiple versions.
        if _is_marking(stix_obj):
            store._data[stix_obj["id"]] = stix_obj

        else:
            if stix_obj["id"] in store._data:
                obj_family = store._data[stix_obj["id"]]
            else:
                obj_family = _ObjectFamily()
                store._data[stix_obj["id"]] = obj_family

            obj_family.add(stix_obj)


def _is_marking(obj_or_id):
    """Determines whether the given object or object ID is/is for a marking
    definition.

    :param obj_or_id: A STIX object or object ID as a string.
    :return: True if a marking definition, False otherwise.
    """

    if isinstance(obj_or_id, (_STIXBase, dict)):
        id_ = obj_or_id["id"]
    else:
        id_ = obj_or_id

    return id_.startswith("marking-definition--")


class _ObjectFamily(object):
    """
    An internal implementation detail of memory sources/sinks/stores.
    Represents a "family" of STIX objects: all objects with a particular
    ID.  (I.e. all versions.)  The latest version is also tracked so that it
    can be obtained quickly.
    """
    def __init__(self):
        self.all_versions = {}
        self.latest_version = None

    def add(self, obj):
        self.all_versions[obj["modified"]] = obj
        if self.latest_version is None or \
                obj["modified"] > self.latest_version["modified"]:
            self.latest_version = obj

    def __str__(self):
        return "<<{}; latest={}>>".format(self.all_versions,
                                          self.latest_version["modified"])

    def __repr__(self):
        return str(self)


class MemoryStore(DataStoreMixin):
    """Interface to an in-memory dictionary of STIX objects.

    MemoryStore is a wrapper around a paired MemorySink and MemorySource.

    Note: It doesn't make sense to create a MemoryStore by passing
    in existing MemorySource and MemorySink because there could
    be data concurrency issues. As well, just as easy to create new MemoryStore.

    Args:
        stix_data (list OR dict OR STIX object): STIX content to be added
        allow_custom (bool): whether to allow custom STIX content.
            Only applied when export/input functions called, i.e.
            load_from_file() and save_to_file(). Defaults to True.
        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
            None, use latest version.

    Attributes:
        _data (dict): the in-memory dict that holds STIX objects
        source (MemorySource): MemorySource
        sink (MemorySink): MemorySink

    """
    def __init__(self, stix_data=None, allow_custom=True, version=None):
        self._data = {}

        if stix_data:
            _add(self, stix_data, allow_custom, version=version)

        super(MemoryStore, self).__init__(
            source=MemorySource(stix_data=self._data, allow_custom=allow_custom, version=version, _store=True),
            sink=MemorySink(stix_data=self._data, allow_custom=allow_custom, version=version, _store=True)
        )

    def save_to_file(self, *args, **kwargs):
        """Write SITX objects from in-memory dictionary to JSON file, as a STIX
        Bundle.

        Args:
            file_path (str): file path to write STIX data to

        """
        return self.sink.save_to_file(*args, **kwargs)

    def load_from_file(self, *args, **kwargs):
        """Load STIX data from JSON file.

        File format is expected to be a single JSON
        STIX object or JSON STIX bundle.

        Args:
            file_path (str): file path to load STIX data from
            version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
                None, use latest version.

        """
        return self.source.load_from_file(*args, **kwargs)


class MemorySink(DataSink):
    """Interface for adding/pushing STIX objects to an in-memory dictionary.

    Designed to be paired with a MemorySource, together as the two
    components of a MemoryStore.

    Args:
        stix_data (dict OR list): valid STIX 2.0 content in
            bundle or a list.
        _store (bool): whether the MemorySink is a part of a MemoryStore,
            in which case "stix_data" is a direct reference to
            shared memory with DataSource. Not user supplied
        allow_custom (bool): whether to allow custom objects/properties
            when exporting STIX content to file.
            Default: True.

    Attributes:
        _data (dict): the in-memory dict that holds STIX objects.
            If part of a MemoryStore, the dict is shared with a MemorySource

    """
    def __init__(self, stix_data=None, allow_custom=True, version=None, _store=False):
        super(MemorySink, self).__init__()
        self.allow_custom = allow_custom

        if _store:
            self._data = stix_data
        else:
            self._data = {}
            if stix_data:
                _add(self, stix_data, allow_custom, version=version)

    def add(self, stix_data, version=None):
        _add(self, stix_data, self.allow_custom, version)
    add.__doc__ = _add.__doc__

    def save_to_file(self, file_path, encoding="utf-8"):
        file_path = os.path.abspath(file_path)

        all_objs = itertools.chain.from_iterable(
            value.all_versions.values() if isinstance(value, _ObjectFamily)
            else [value]
            for value in self._data.values()
        )

        if not os.path.exists(os.path.dirname(file_path)):
            os.makedirs(os.path.dirname(file_path))
        with open(file_path, "w") as f:
{'fraction_non_alphanumeric': 0.08700902052702607, 'fraction_numerical': 0.0033383052773634492, 'mean_word_length': 2.408375695957395, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6507881', 'n_tokens_mistral': 4245, 'n_tokens_neox': 4091, 'n_words': 1360}