From f79b3c98766b5e330fb0ef1c29e27eca6e5c9f23 Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Fri, 10 May 2019 10:22:45 -0400 Subject: [PATCH] Add functionality to _valid_refs to accept actual cyber observable objects instead of just strings with their types --- stix2/base.py | 5 ++++- stix2/test/v21/test_observed_data.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/stix2/base.py b/stix2/base.py index 9fcdf56..a9a801e 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -308,7 +308,10 @@ class _Observable(_STIXBase): allowed_types = prop.valid_types try: - ref_type = self._STIXBase__valid_refs[ref] + try: + ref_type = self._STIXBase__valid_refs[ref].type + except AttributeError: + ref_type = self._STIXBase__valid_refs[ref] except TypeError: raise ValueError("'%s' must be created with _valid_refs as a dict, not a list." % self.__class__.__name__) diff --git a/stix2/test/v21/test_observed_data.py b/stix2/test/v21/test_observed_data.py index 5d0f9b1..864dd7a 100644 --- a/stix2/test/v21/test_observed_data.py +++ b/stix2/test/v21/test_observed_data.py @@ -940,6 +940,24 @@ def test_ip4_address_example(): assert ip4.resolves_to_refs == ["4", "5"] +def test_ip4_address_valid_refs(): + mac1 = stix2.v21.MACAddress( + value="a1:b2:c3:d4:e5:f6", + ) + mac2 = stix2.v21.MACAddress( + value="a7:b8:c9:d0:e1:f2", + ) + + ip4 = stix2.v21.IPv4Address( + _valid_refs={"1": mac1, "2": mac2}, + value="177.60.40.7", + resolves_to_refs=["1", "2"], + ) + + assert ip4.value == "177.60.40.7" + assert ip4.resolves_to_refs == ["1", "2"] + + def test_ip4_address_example_cidr(): ip4 = stix2.v21.IPv4Address(value="198.51.100.0/24")