Merge pull request #49 from chisholm/fix-find-references

Fix reference search code to search for references inside objects which are inside arrays.
oasis-open · Nov 28, 2023 · 7319f21 · 7319f21
2 parents afbe459 + b74b2b3
commit 7319f21
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/stix2generator/test/test_utils.py b/stix2generator/test/test_utils.py
@@ -18,19 +18,21 @@ def object_generator21():
     "obj, findings", [
         ({"type": "foo", "a_ref": 1, "b_ref": 2, "c": 3}, {("a_ref", 1), ("b_ref", 2)}),
         ({"type": "foo", "a_refs": [1, 2], "b_ref": 3, "c": 4}, {("a_refs", 1), ("a_refs", 2), ("b_ref", 3)}),
-        ({"type": "foo", "a": {"b": {"c_ref": 1}}}, {("c_ref", 1)})
+        ({"type": "foo", "a": {"b": {"c_ref": 1}}}, {("c_ref", 1)}),
+        ({"type": "foo", "a": [{"b_ref": 1}, {"b_ref": 2}, {"c": [{"d_ref": 3}]}]}, {("b_ref", 1), ("b_ref", 2), ("d_ref", 3)})
     ]
 )
 def test_find_references(obj, findings):
-    for ref_prop, ref_id in stix2generator.utils.find_references(obj):
-        assert (ref_prop, ref_id) in findings
+    found = set(stix2generator.utils.find_references(obj))
+    assert found == findings
 
 
 @pytest.mark.parametrize(
     "obj", [
         {"type": "foo", "a_ref": 1, "b_ref": 1, "c": 1},
         {"type": "foo", "a_refs": [1, 1], "b_ref": 1, "c": 1},
-        {"type": "foo", "a": {"b": {"c_ref": 1}}}
+        {"type": "foo", "a": {"b": {"c_ref": 1}}},
+        {"type": "foo", "a": [{"b_ref": 1}, {"b_ref": 1}, {"c": [{"d_ref": 1}]}]}
     ]
 )
 def test_find_references_assignable(obj):

diff --git a/stix2generator/utils.py b/stix2generator/utils.py
@@ -68,7 +68,7 @@ def recurse_references(obj):
     observed-data/objects.
 
     :param obj: An object.  Can be any type, but values will only be produced
-        from mappings.
+        from mappings (or iterables containing mappings).
     """
     if isinstance(obj, collections.abc.Mapping):
         for prop, value in obj.items():
@@ -82,6 +82,11 @@ def recurse_references(obj):
             else:
                 yield from recurse_references(value)
 
+    elif isinstance(obj, collections.abc.Iterable) \
+            and not isinstance(obj, str):
+        for elt in obj:
+            yield from recurse_references(elt)
+
 
 def find_references(obj):
     """
@@ -116,7 +121,7 @@ def recurse_references_assignable(obj):
     special casing for observed-data/objects.
 
     :param obj: An object.  Can be any type, but values will only be produced
-        from mappings.
+        from mappings (or iterables containing mappings).
     """
     if isinstance(obj, collections.abc.Mapping):
         for prop, value in obj.items():
@@ -130,6 +135,11 @@ def recurse_references_assignable(obj):
             else:
                 yield from recurse_references_assignable(value)
 
+    elif isinstance(obj, collections.abc.Iterable) \
+            and not isinstance(obj, str):
+        for elt in obj:
+            yield from recurse_references_assignable(elt)
+
 
 def find_references_assignable(obj):
     """