From 5285934034d575d57dd08e7b77f1de17f52a1941 Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Wed, 29 Nov 2017 19:25:52 -0500 Subject: [PATCH] Make Workbench use implicit ObjectFactory This is needed to implement functions like `set_default_creator`. The changes to Tox are so that the wrapping we do in workbench doesn't affect the rest of our tests. If we test them all in one go, pytest will import all the tests before running any of them. This will cause the workbench versions of the SDO classes to be used in all tests. --- stix2/test/test_tool.py | 6 +++++ stix2/test/test_workbench.py | 31 +++++++++++++------------ stix2/workbench.py | 44 ++++++++++++++++++++++++++++-------- tox.ini | 3 ++- 4 files changed, 60 insertions(+), 24 deletions(-) diff --git a/stix2/test/test_tool.py b/stix2/test/test_tool.py index 21ece24..ce99fb8 100644 --- a/stix2/test/test_tool.py +++ b/stix2/test/test_tool.py @@ -58,4 +58,10 @@ def test_parse_tool(data): assert tool.labels == ["remote-access"] assert tool.name == "VNC" + +def test_tool_no_workbench_wrappers(): + tool = stix2.Tool(name='VNC', labels=['remote-access']) + with pytest.raises(AttributeError): + tool.created_by() + # TODO: Add other examples diff --git a/stix2/test/test_workbench.py b/stix2/test/test_workbench.py index 7f3c9fc..6a33f11 100644 --- a/stix2/test/test_workbench.py +++ b/stix2/test/test_workbench.py @@ -1,5 +1,8 @@ import stix2 -from stix2.workbench import (add, all_versions, attack_patterns, campaigns, +from stix2.workbench import (AttackPattern, Campaign, CourseOfAction, Identity, + Indicator, IntrusionSet, Malware, ObservedData, + Report, ThreatActor, Tool, Vulnerability, add, + all_versions, attack_patterns, campaigns, courses_of_action, create, get, identities, indicators, intrusion_sets, malware, observed_data, query, reports, threat_actors, @@ -19,7 +22,7 @@ from .constants import (ATTACK_PATTERN_ID, ATTACK_PATTERN_KWARGS, CAMPAIGN_ID, def test_workbench_environment(): # Create a STIX object - ind = create(stix2.Indicator, id=INDICATOR_ID, **INDICATOR_KWARGS) + ind = create(Indicator, id=INDICATOR_ID, **INDICATOR_KWARGS) add(ind) resp = get(INDICATOR_ID) @@ -35,7 +38,7 @@ def test_workbench_environment(): def test_workbench_get_all_attack_patterns(): - mal = stix2.AttackPattern(id=ATTACK_PATTERN_ID, **ATTACK_PATTERN_KWARGS) + mal = AttackPattern(id=ATTACK_PATTERN_ID, **ATTACK_PATTERN_KWARGS) add(mal) resp = attack_patterns() @@ -44,7 +47,7 @@ def test_workbench_get_all_attack_patterns(): def test_workbench_get_all_campaigns(): - cam = stix2.Campaign(id=CAMPAIGN_ID, **CAMPAIGN_KWARGS) + cam = Campaign(id=CAMPAIGN_ID, **CAMPAIGN_KWARGS) add(cam) resp = campaigns() @@ -53,7 +56,7 @@ def test_workbench_get_all_campaigns(): def test_workbench_get_all_courses_of_action(): - coa = stix2.CourseOfAction(id=COURSE_OF_ACTION_ID, **COURSE_OF_ACTION_KWARGS) + coa = CourseOfAction(id=COURSE_OF_ACTION_ID, **COURSE_OF_ACTION_KWARGS) add(coa) resp = courses_of_action() @@ -62,7 +65,7 @@ def test_workbench_get_all_courses_of_action(): def test_workbench_get_all_identities(): - idty = stix2.Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) + idty = Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) add(idty) resp = identities() @@ -77,7 +80,7 @@ def test_workbench_get_all_indicators(): def test_workbench_get_all_intrusion_sets(): - ins = stix2.IntrusionSet(id=INTRUSION_SET_ID, **INTRUSION_SET_KWARGS) + ins = IntrusionSet(id=INTRUSION_SET_ID, **INTRUSION_SET_KWARGS) add(ins) resp = intrusion_sets() @@ -86,7 +89,7 @@ def test_workbench_get_all_intrusion_sets(): def test_workbench_get_all_malware(): - mal = stix2.Malware(id=MALWARE_ID, **MALWARE_KWARGS) + mal = Malware(id=MALWARE_ID, **MALWARE_KWARGS) add(mal) resp = malware() @@ -95,7 +98,7 @@ def test_workbench_get_all_malware(): def test_workbench_get_all_observed_data(): - od = stix2.ObservedData(id=OBSERVED_DATA_ID, **OBSERVED_DATA_KWARGS) + od = ObservedData(id=OBSERVED_DATA_ID, **OBSERVED_DATA_KWARGS) add(od) resp = observed_data() @@ -104,7 +107,7 @@ def test_workbench_get_all_observed_data(): def test_workbench_get_all_reports(): - rep = stix2.Report(id=REPORT_ID, **REPORT_KWARGS) + rep = Report(id=REPORT_ID, **REPORT_KWARGS) add(rep) resp = reports() @@ -113,7 +116,7 @@ def test_workbench_get_all_reports(): def test_workbench_get_all_threat_actors(): - thr = stix2.ThreatActor(id=THREAT_ACTOR_ID, **THREAT_ACTOR_KWARGS) + thr = ThreatActor(id=THREAT_ACTOR_ID, **THREAT_ACTOR_KWARGS) add(thr) resp = threat_actors() @@ -122,7 +125,7 @@ def test_workbench_get_all_threat_actors(): def test_workbench_get_all_tools(): - tool = stix2.Tool(id=TOOL_ID, **TOOL_KWARGS) + tool = Tool(id=TOOL_ID, **TOOL_KWARGS) add(tool) resp = tools() @@ -131,7 +134,7 @@ def test_workbench_get_all_tools(): def test_workbench_get_all_vulnerabilities(): - vuln = stix2.Vulnerability(id=VULNERABILITY_ID, **VULNERABILITY_KWARGS) + vuln = Vulnerability(id=VULNERABILITY_ID, **VULNERABILITY_KWARGS) add(vuln) resp = vulnerabilities() @@ -152,7 +155,7 @@ def test_workbench_relationships(): def test_workbench_created_by(): - intset = stix2.IntrusionSet(name="Breach 123", created_by_ref=IDENTITY_ID) + intset = IntrusionSet(name="Breach 123", created_by_ref=IDENTITY_ID) add(intset) creator = intset.created_by() assert creator.id == IDENTITY_ID diff --git a/stix2/workbench.py b/stix2/workbench.py index 4069309..f5f2d32 100644 --- a/stix2/workbench.py +++ b/stix2/workbench.py @@ -1,9 +1,18 @@ """Functions and class wrappers for interacting with STIX data at a high level. """ -from . import (AttackPattern, Campaign, CourseOfAction, CustomObject, Identity, - Indicator, IntrusionSet, Malware, ObservedData, Report, - ThreatActor, Tool, Vulnerability) +from . import AttackPattern as _AttackPattern +from . import Campaign as _Campaign +from . import CourseOfAction as _CourseOfAction +from . import Identity as _Identity +from . import Indicator as _Indicator +from . import IntrusionSet as _IntrusionSet +from . import Malware as _Malware +from . import ObservedData as _ObservedData +from . import Report as _Report +from . import ThreatActor as _ThreatActor +from . import Tool as _Tool +from . import Vulnerability as _Vulnerability from .environment import Environment from .sources.filters import Filter from .sources.memory import MemoryStore @@ -27,6 +36,11 @@ add_data_source = _environ.source.add_data_source # Wrap SDOs with helper functions +STIX_OBJS = [_AttackPattern, _Campaign, _CourseOfAction, _Identity, + _Indicator, _IntrusionSet, _Malware, _ObservedData, _Report, + _ThreatActor, _Tool, _Vulnerability] + + def created_by_wrapper(self, *args, **kwargs): return _environ.creator_of(self, *args, **kwargs) @@ -39,14 +53,26 @@ def related_wrapper(self, *args, **kwargs): return _environ.related_to(self, *args, **kwargs) -STIX_OBJS = [AttackPattern, Campaign, CourseOfAction, CustomObject, Identity, - Indicator, IntrusionSet, Malware, ObservedData, Report, - ThreatActor, Tool, Vulnerability] +def constructor_wrapper(obj_type): + # Use an intermediate wrapper class so the implicit environment will create objects that have our wrapper functions + wrapped_type = type(obj_type.__name__, obj_type.__bases__, dict( + created_by=created_by_wrapper, + relationships=relationships_wrapper, + related=related_wrapper, + **obj_type.__dict__ + )) + @staticmethod + def new_constructor(cls, *args, **kwargs): + return _environ.create(wrapped_type, *args, **kwargs) + return new_constructor + + +# Create wrapper classes whose constructors call the implicit environment's create() for obj_type in STIX_OBJS: - obj_type.created_by = created_by_wrapper - obj_type.relationships = relationships_wrapper - obj_type.related = related_wrapper + new_class = type(obj_type.__name__, (), {}) + new_class.__new__ = constructor_wrapper(obj_type) + globals()[obj_type.__name__] = new_class # Functions to get all objects of a specific type diff --git a/tox.ini b/tox.ini index bfc8c1b..97d9519 100644 --- a/tox.ini +++ b/tox.ini @@ -9,7 +9,8 @@ deps = pytest-cov coverage commands = - py.test --cov=stix2 stix2/test/ --cov-report term-missing + py.test --ignore=stix2/test/test_workbench.py --cov=stix2 stix2/test/ --cov-report term-missing --cov-append + py.test stix2/test/test_workbench.py --cov=stix2 --cov-report term-missing --cov-append passenv = CI TRAVIS TRAVIS_*