From 4af46db0038db7c3722b4aeaa9aeb1b23ffe3a7c Mon Sep 17 00:00:00 2001 From: Michael Chisholm Date: Wed, 6 Dec 2023 22:18:47 -0500 Subject: [PATCH] Fix an issue where reference graph generator's "halt" generator didn't inherit the original object generator's specs or providers. This can cause crashes when reference graph generator is used with object generators which have custom specs and/or semantics providers. --- stix2generator/generation/object_generator.py | 26 +++++- .../generation/reference_graph_generator.py | 2 +- .../test/test_reference_graph_generator.py | 87 +++++++++++++++++++ 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/stix2generator/generation/object_generator.py b/stix2generator/generation/object_generator.py index 1558a2b..0869dc4 100644 --- a/stix2generator/generation/object_generator.py +++ b/stix2generator/generation/object_generator.py @@ -1,3 +1,4 @@ +import copy import itertools import logging import math @@ -106,8 +107,8 @@ def __init__( :param spec_registry: A name->specification mapping used to look up references inside of specifications - :param semantic_providers: A list of semantic providers (e.g. - instances of subclasses of SemanticsProvider) + :param semantic_providers: An iterable of semantic providers (e.g. + instances of subclasses of SemanticsProvider), or None :param config: A Config instance giving user settings regarding generation. If None, defaults will be used. """ @@ -142,6 +143,27 @@ def spec_names(self): """ return self.__specs.keys() + def derive_generator(self, new_config): + """ + Create an object generator using this generator's registry and + semantics providers, but with a given config. + + :param new_config: An instance of Config giving desired configuration + settings for the new object generator + :return: A new object generator + """ + + our_providers = copy.deepcopy(set(self.__semantics.values())) + our_specs = copy.deepcopy(self.__specs) + + new_generator = ObjectGenerator( + spec_registry=our_specs, + semantic_providers=our_providers, + config=new_config + ) + + return new_generator + def generate( self, spec_name, expected_type=None, spec_name_stack=None, value_constraint=None diff --git a/stix2generator/generation/reference_graph_generator.py b/stix2generator/generation/reference_graph_generator.py index 05c4646..4184401 100644 --- a/stix2generator/generation/reference_graph_generator.py +++ b/stix2generator/generation/reference_graph_generator.py @@ -515,7 +515,7 @@ def __init__( .Config( **halt_generator_config_dict ) - self.__halt_generator = stix2generator.create_object_generator( + self.__halt_generator = object_generator.derive_generator( halt_generator_config ) diff --git a/stix2generator/test/test_reference_graph_generator.py b/stix2generator/test/test_reference_graph_generator.py index b82cdd2..3faa91f 100644 --- a/stix2generator/test/test_reference_graph_generator.py +++ b/stix2generator/test/test_reference_graph_generator.py @@ -1,12 +1,14 @@ import copy import pytest +import stix2 import stix2.base import stix2.utils import stix2generator import stix2generator.generation.object_generator import stix2generator.generation.reference_graph_generator +import stix2generator.generation.semantics import stix2generator.test.utils _TLP_MARKING_DEFINITION_IDS = { @@ -959,3 +961,88 @@ def test_not_parsing(num_trials): else: assert isinstance(obj, dict) + + +def test_ref_gen_with_custom(): + """ + Set up a reference graph generator with an object generator with stuff + you would not get by default via + stix2generator.create_object_generator(...). Reference graph generator + creates a derivative "halt" generator, and we need to ensure that the + latter inherits the same config as the original object generator, including + registry and semantics providers. + """ + class TestSemantics(stix2generator.generation.semantics.SemanticsProvider): + def get_semantics(self): + return ["testxyz"] + + def testxyz(self, spec, generator, constraint): + return "test" + + custom_registry = { + # A self-referential type, to cause simple reference chains + "test-ref-type": { + "type": "object", + "optional": ["obj_ref"], + "properties": { + "id": { + "type": "string", + "semantics": "stix-id", + "stix-type": "test-ref-type" + }, + "type": "test-ref-type", + "test_prop": { + "type": "string", + "semantics": "testxyz" + }, + "obj_ref": { + "type": "string", + "semantics": "stix-id", + "stix-type": "test-ref-type" + } + } + } + } + + @stix2.CustomObject("test-ref-type", []) + class TestRefType: + # This class won't be used since we'll turn the parse setting off; but + # we need a registration so test-ref-type is seen as a generatable + # type. So we can leave it empty. + pass + + semantics_providers = [ + TestSemantics(), + stix2generator.generation.semantics.STIXSemantics() + ] + + obj_gen_config = stix2generator.generation.object_generator.Config( + optional_property_probability=1, + minimize_ref_properties=False + ) + + obj_gen = stix2generator.generation.object_generator.ObjectGenerator( + custom_registry, semantics_providers, obj_gen_config + ) + + ref_gen_config = stix2generator.generation.reference_graph_generator \ + .Config( + graph_type="TREE", + parse=False + ) + + ref_gen = stix2generator.generation.reference_graph_generator \ + .ReferenceGraphGenerator( + obj_gen, ref_gen_config + ) + + _, graph = ref_gen.generate() + + for obj in graph.values(): + # Ensure our semantics provider was invoked properly for all objects + assert obj["test_prop"] == "test" + + # Ensure no dangling references + obj_ref = obj.get("obj_ref") + if obj_ref is not None: + assert obj_ref in graph