Repository: oasis-open/cti-python-stix2
Pull request #500 (2021-03-17 19:18:37)

fp write for STIX Objects
Adding this piece of code would help write more flexible code while also providing the same options as ``serialize()``


Patch:
diff --git a/setup.py b/setup.py
index f382412..3f82733 100644
--- a/setup.py
+++ b/setup.py
@@ -60,7 +60,7 @@ setup(
         'Bug Tracker': 'https://github.com/oasis-open/cti-python-stix2/issues/',
     },
     extras_require={
-        'taxii': ['taxii2-client>=2.3.0'],
+        'taxii': ['taxii2-client>=2.2.1'],
         'semantic': ['haversine', 'rapidfuzz'],
     },
 )
diff --git a/stix2/base.py b/stix2/base.py
index b158f06..2c48ef6 100644
--- a/stix2/base.py
+++ b/stix2/base.py
@@ -17,7 +17,8 @@ from .exceptions import (
 from .markings import _MarkingsMixin
 from .markings.utils import validate
 from .serialization import (
-    STIXJSONEncoder, STIXJSONIncludeOptionalDefaultsEncoder, serialize,
+    STIXJSONEncoder, STIXJSONIncludeOptionalDefaultsEncoder, fp_serialize,
+    serialize,
 )
 from .utils import NOW, PREFIX_21_REGEX, get_timestamp
 from .versioning import new_version as _new_version
@@ -260,6 +261,35 @@ class _STIXBase(Mapping):
         """
         return serialize(self, *args, **kwargs)
 
+    def fp_serialize(self, *args, **kwargs):
+        """
+        Serialize a STIX object to ``fp`` (a text stream file-like supporting object).
+
+        Examples:
+            >>> import stix2
+            >>> identity = stix2.Identity(name='Example Corp.', identity_class='organization')
+            >>> print(identity.serialize(sort_keys=True))
+            {"created": "2018-06-08T19:03:54.066Z", ... "name": "Example Corp.", "type": "identity"}
+            >>> print(identity.serialize(sort_keys=True, indent=4))
+            {
+                "created": "2018-06-08T19:03:54.066Z",
+                "id": "identity--d7f3e25a-ba1c-447a-ab71-6434b092b05e",
+                "identity_class": "organization",
+                "modified": "2018-06-08T19:03:54.066Z",
+                "name": "Example Corp.",
+                "type": "identity"
+            }
+            >>> with open("example.json", mode="w", encoding="utf-8") as f:
+            >>>     identity.fp_serialize(f, pretty=True)
+
+        Returns:
+            None
+
+        See Also:
+            ``stix2.serialization.fp_serialize`` for options.
+        """
+        fp_serialize(self, *args, **kwargs)
+
 
 class _DomainObject(_STIXBase, _MarkingsMixin):
     pass
diff --git a/stix2/datastore/filesystem.py b/stix2/datastore/filesystem.py
index d844115..2209116 100644
--- a/stix2/datastore/filesystem.py
+++ b/stix2/datastore/filesystem.py
@@ -13,7 +13,7 @@ from stix2.datastore import (
 )
 from stix2.datastore.filters import Filter, FilterSet, apply_common_filters
 from stix2.parsing import parse
-from stix2.serialization import serialize
+from stix2.serialization import fp_serialize
 from stix2.utils import format_datetime, get_type_from_id, parse_into_datetime
 
 
@@ -584,9 +584,8 @@ class FileSystemSink(DataSink):
         if os.path.isfile(file_path):
             raise DataSourceError("Attempted to overwrite file (!) at: {}".format(file_path))
 
-        with io.open(file_path, 'w', encoding=encoding) as f:
-            stix_obj = serialize(stix_obj, pretty=True, encoding=encoding, ensure_ascii=False)
-            f.write(stix_obj)
+        with io.open(file_path, mode='w', encoding=encoding) as f:
+            fp_serialize(stix_obj, f, pretty=True, encoding=encoding, ensure_ascii=False)
 
     def add(self, stix_data=None, version=None):
         """Add STIX objects to file directory.
diff --git a/stix2/datastore/taxii.py b/stix2/datastore/taxii.py
index 9ad6df9..41d1e54 100644
--- a/stix2/datastore/taxii.py
+++ b/stix2/datastore/taxii.py
@@ -12,8 +12,6 @@ from stix2.parsing import parse
 from stix2.utils import deduplicate
 
 try:
-    from taxii2client import v20 as tcv20
-    from taxii2client import v21 as tcv21
     from taxii2client.exceptions import ValidationError
     _taxii2_client = True
 except ImportError:
@@ -35,12 +33,9 @@ class TAXIICollectionStore(DataStoreMixin):
             side(retrieving data) and False for TAXIICollectionSink
             side(pushing data). However, when parameter is supplied, it will
             be applied to both TAXIICollectionSource/Sink.
-        items_per_page (int): How many STIX objects to request per call
-            to TAXII Server. The value can be tuned, but servers may override
-            if their internal limit is surpassed. Used by TAXIICollectionSource
 
     """
-    def __init__(self, collection, allow_custom=None, items_per_page=5000):
+    def __init__(self, collection, allow_custom=None):
         if allow_custom is None:
             allow_custom_source = True
             allow_custom_sink = False
@@ -48,7 +43,7 @@ class TAXIICollectionStore(DataStoreMixin):
             allow_custom_sink = allow_custom_source = allow_custom
 
         super(TAXIICollectionStore, self).__init__(
-            source=TAXIICollectionSource(collection, allow_custom=allow_custom_source, items_per_page=items_per_page),
+            source=TAXIICollectionSource(collection, allow_custom=allow_custom_source),
             sink=TAXIICollectionSink(collection, allow_custom=allow_custom_sink),
         )
 
@@ -149,12 +144,9 @@ class TAXIICollectionSource(DataSource):
         collection (taxii2.Collection): TAXII Collection instance
         allow_custom (bool): Whether to allow custom STIX content to be
             added to the FileSystemSink. Default: True
-        items_per_page (int): How many STIX objects to request per call
-            to TAXII Server. The value can be tuned, but servers may override
-            if their internal limit is surpassed.
 
     """
-    def __init__(self, collection, allow_custom=True, items_per_page=5000):
+    def __init__(self, collection, allow_custom=True):
         super(TAXIICollectionSource, self).__init__()
         if not _taxii2_client:
             raise ImportError("taxii2client library is required for usage of TAXIICollectionSource")
@@ -175,7 +167,6 @@ class TAXIICollectionSource(DataSource):
             )
 
         self.allow_custom = allow_custom
-        self.items_per_page = items_per_page
 
     def get(self, stix_id, version=None, _composite_filters=None):
         """Retrieve STIX object from local/remote STIX Collection
@@ -295,12 +286,8 @@ class TAXIICollectionSource(DataSource):
         taxii_filters_dict = dict((f.property, f.value) for f in taxii_filters)
 
         # query TAXII collection
-        all_data = []
         try:
-            paged_request = tcv21.as_pages if isinstance(self.collection, tcv21.Collection) else tcv20.as_pages
-
-            for resource in paged_request(self.collection.get_objects, per_request=self.items_per_page, **taxii_filters_dict):
-                all_data.extend(resource.get("objects", []))
+            all_data = self.collection.get_objects(**taxii_filters_dict).get('objects', [])
 
             # deduplicate data (before filtering as reduces wasted filtering)
             all_data = deduplicate(all_data)
diff --git a/stix2/equivalence/object/__init__.py b/stix2/equivalence/object/__init__.py
index dde52ec..da2097e 100644
--- a/stix2/equivalence/object/__init__.py
+++ b/stix2/equivalence/object/__init__.py
@@ -559,8 +559,8 @@ WEIGHTS = {
         "sectors": (20, partial_list_based),
     },
     "incident": {
-        "name": (30, partial_string_based),
-        "external_references": (70, partial_external_reference_based),
+        "name": (60, partial_string_based),
+        "external_references": (40, partial_external_reference_based),
     },
     "indicator": {
         "indicator_types": (15, partial_list_based),
diff --git a/stix2/serialization.py b/stix2/serialization.py
index 7488eb5..2784d39 100644
--- a/stix2/serialization.py
+++ b/stix2/serialization.py
@@ -2,6 +2,7 @@
 
 import copy
 import datetime as dt
+import io
 
 import simplejson as json
 
@@ -64,6 +65,37 @@ def serialize(obj, pretty=False, include_optional_defaults=False, **kwargs):
     Returns:
         str: The serialized JSON object.
 
+    Note:
+        The argument ``pretty=True`` will output the STIX object following
+        spec order. Using this argument greatly impacts object serialization
+        performance. If your use case is centered across machine-to-machine
+        operation it is recommended to set ``pretty=False``.
+
+        When ``pretty=True`` the following key-value pairs will be added or
+        overridden: indent=4, separators=(",", ": "), item_sort_key=sort_by.
+    """
+    with io.StringIO() as fp:
+        fp_serialize(obj, fp, pretty, include_optional_defaults, **kwargs)
+        return fp.getvalue()
+
+
+def fp_serialize(obj, fp, pretty=False, include_optional_defaults=False, **kwargs):
+    """
+    Serialize a STIX object to ``fp`` (a text stream file-like supporting object).
+
+    Args:
+        obj: The STIX object to be serialized.
+        fp: A text stream file-like object supporting ``.write()``.
+        pretty (bool): If True, output properties following the STIX specs
+            formatting. This includes indentation. Refer to notes for more
+            details. (Default: ``False``)
+        include_optional_defaults (bool): Determines whether to include
+            optional properties set to the default value defined in the spec.
+        **kwargs: The arguments for a json.dumps() call.
+
+    Returns:
+        None
+
     Note:
         The argument ``pretty=True`` will output the STIX object following
         spec order. Using this argument greatly impacts object serialization
@@ -80,9 +112,9 @@ def serialize(obj, pretty=False, include_optional_defaults=False, **kwargs):
         kwargs.update({'indent': 4, 'separators': (',', ': '), 'item_sort_key': sort_by})
 
     if include_optional_defaults:
-        return json.dumps(obj, cls=STIXJSONIncludeOptionalDefaultsEncoder, **kwargs)
+        json.dump(obj, fp, cls=STIXJSONIncludeOptionalDefaultsEncoder, **kwargs)
     else:
-        return json.dumps(obj, cls=STIXJSONEncoder, **kwargs)
+        json.dump(obj, fp, cls=STIXJSONEncoder, **kwargs)
 
 
 def _find(seq, val):


Test patch:
diff --git a/stix2/test/v20/test_bundle.py b/stix2/test/v20/test_bundle.py
index f53d0cb..ac5d239 100644
--- a/stix2/test/v20/test_bundle.py
+++ b/stix2/test/v20/test_bundle.py
@@ -1,3 +1,4 @@
+import io
 import json
 
 import pytest
@@ -113,6 +114,27 @@ def test_bundle_id_must_start_with_bundle():
     assert str(excinfo.value) == "Invalid value for Bundle 'id': must start with 'bundle--'."
 
 
+def test_create_bundle_fp_serialize_pretty(indicator, malware, relationship):
+    bundle = stix2.v20.Bundle(objects=[indicator, malware, relationship])
+    buffer = io.StringIO()
+
+    bundle.fp_serialize(buffer, pretty=True)
+
+    assert str(bundle) == EXPECTED_BUNDLE
+    assert bundle.serialize(pretty=True) == EXPECTED_BUNDLE
+    assert buffer.getvalue() == EXPECTED_BUNDLE
+
+
+def test_create_bundle_fp_serialize_nonpretty(indicator, malware, relationship):
+    bundle = stix2.v20.Bundle(objects=[indicator, malware, relationship])
+    buffer = io.StringIO()
+
+    bundle.fp_serialize(buffer, sort_keys=True)
+
+    assert bundle.serialize(sort_keys=True) == json.dumps(json.loads(EXPECTED_BUNDLE), sort_keys=True)
+    assert buffer.getvalue() == json.dumps(json.loads(EXPECTED_BUNDLE), sort_keys=True)
+
+
 def test_create_bundle1(indicator, malware, relationship):
     bundle = stix2.v20.Bundle(objects=[indicator, malware, relationship])
 
diff --git a/stix2/test/v20/test_datastore_taxii.py b/stix2/test/v20/test_datastore_taxii.py
index 075f0a3..34daa80 100644
--- a/stix2/test/v20/test_datastore_taxii.py
+++ b/stix2/test/v20/test_datastore_taxii.py
@@ -4,7 +4,7 @@ from medallion.filters.basic_filter import BasicFilter
 import pytest
 from requests.models import Response
 from taxii2client.common import _filter_kwargs_to_query_params
-from taxii2client.v20 import MEDIA_TYPE_STIX_V20, Collection
+from taxii2client.v20 import Collection
 
 import stix2
 from stix2.datastore import DataSourceError
@@ -34,12 +34,12 @@ class MockTAXIICollectionEndpoint(Collection):
                 {
                     "date_added": get_timestamp(),
                     "id": object["id"],
-                    "media_type": "application/stix+json;version=2.0",
+                    "media_type": "application/stix+json;version=2.1",
                     "version": object.get("modified", object.get("created", get_timestamp())),
                 },
             )
 
-    def get_objects(self, accept=MEDIA_TYPE_STIX_V20, start=0, per_request=0, **filter_kwargs):
+    def get_objects(self, **filter_kwargs):
         self._verify_can_read()
         query_params = _filter_kwargs_to_query_params(filter_kwargs)
         assert isinstance(query_params, dict)
@@ -51,12 +51,7 @@ class MockTAXIICollectionEndpoint(Collection):
             100,
         )[0]
         if objs:
-            resp = Response()
-            resp.status_code = 200
-            resp.headers["Content-Range"] = f"items 0-{len(objs)}/{len(objs)}"
-            resp.encoding = "utf-8"
-            resp._content = bytes(stix2.v20.Bundle(objects=objs).serialize(ensure_ascii=False), resp.encoding)
-            return resp
+            return stix2.v20.Bundle(objects=objs)
         else:
             resp = Response()
             resp.status_code = 404
diff --git a/stix2/test/v21/test_bundle.py b/stix2/test/v21/test_bundle.py
index 4e30c84..1cf30d0 100644
--- a/stix2/test/v21/test_bundle.py
+++ b/stix2/test/v21/test_bundle.py
@@ -1,3 +1,4 @@
+import io
 import json
 
 import pytest
@@ -123,6 +124,27 @@ def test_bundle_id_must_start_with_bundle():
     assert str(excinfo.value) == "Invalid value for Bundle 'id': must start with 'bundle--'."
 
 
+def test_create_bundle_fp_serialize_pretty(indicator, malware, relationship):
+    bundle = stix2.v21.Bundle(objects=[indicator, malware, relationship])
+    buffer = io.StringIO()
+
+    bundle.fp_serialize(buffer, pretty=True)
+
+    assert str(bundle) == EXPECTED_BUNDLE
+    assert bundle.serialize(pretty=True) == EXPECTED_BUNDLE
+    assert buffer.getvalue() == EXPECTED_BUNDLE
+
+
+def test_create_bundle_fp_serialize_nonpretty(indicator, malware, relationship):
+    bundle = stix2.v21.Bundle(objects=[indicator, malware, relationship])
+    buffer = io.StringIO()
+
+    bundle.fp_serialize(buffer, sort_keys=True)
+
+    assert bundle.serialize(sort_keys=True) == json.dumps(json.loads(EXPECTED_BUNDLE), sort_keys=True)
+    assert buffer.getvalue() == json.dumps(json.loads(EXPECTED_BUNDLE), sort_keys=True)
+
+
 def test_create_bundle1(indicator, malware, relationship):
     bundle = stix2.v21.Bundle(objects=[indicator, malware, relationship])
 
diff --git a/stix2/test/v21/test_datastore_taxii.py b/stix2/test/v21/test_datastore_taxii.py
index 62ddadc..4b7e299 100644
--- a/stix2/test/v21/test_datastore_taxii.py
+++ b/stix2/test/v21/test_datastore_taxii.py
@@ -28,14 +28,14 @@ class MockTAXIICollectionEndpoint(Collection):
         self._verify_can_write()
         if isinstance(bundle, str):
             bundle = json.loads(bundle)
-        for obj in bundle.get("objects", []):
-            self.objects.append(obj)
+        for object in bundle.get("objects", []):
+            self.objects.append(object)
             self.manifests.append(
                 {
                     "date_added": get_timestamp(),
-                    "id": obj["id"],
+                    "id": object["id"],
                     "media_type": "application/stix+json;version=2.1",
-                    "version": obj.get("modified", obj.get("created", get_timestamp())),
+                    "version": object.get("modified", object.get("created", get_timestamp())),
                 },
             )
 
@@ -51,10 +51,7 @@ class MockTAXIICollectionEndpoint(Collection):
             100,
         )[0]
         if objs:
-            return {
-                "objects": objs,
-                "more": False,
-            }
+            return stix2.v21.Bundle(objects=objs)
         else:
             resp = Response()
             resp.status_code = 404
@@ -78,10 +75,7 @@ class MockTAXIICollectionEndpoint(Collection):
         else:
             filtered_objects = []
         if filtered_objects:
-            return {
-                "objects": filtered_objects,
-                "more": False,
-            }
+            return stix2.v21.Bundle(objects=filtered_objects)
         else:
             resp = Response()
             resp.status_code = 404