Test adding a data source to the workbench

stix2/test/test_workbench.py

@@ -1,10 +1,12 @@
+import os
+
 import stix2
 from stix2.workbench import (AttackPattern, Campaign, CourseOfAction, Identity,
                              Indicator, IntrusionSet, Malware, ObservedData,
                              Report, ThreatActor, Tool, Vulnerability, add,
-                             all_versions, attack_patterns, campaigns,
+                             add_data_source, all_versions, attack_patterns,
-                             courses_of_action, create, get, identities,
+                             campaigns, courses_of_action, create, get,
-                             indicators, intrusion_sets, malware,
+                             identities, indicators, intrusion_sets, malware,
                              observed_data, query, reports, threat_actors,
                              tools, vulnerabilities)
 
@@ -171,3 +173,19 @@ def test_workbench_related():
     assert any(x['id'] == CAMPAIGN_ID for x in resp)
     assert any(x['id'] == INDICATOR_ID for x in resp)
     assert any(x['id'] == IDENTITY_ID for x in resp)
+
+    resp = get(MALWARE_ID).related(relationship_type='indicates')
+    assert len(resp) == 1
+
+
+def test_add_data_source():
+    fs_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "stix2_data")
+    fs = stix2.FileSystemSource(fs_path)
+    add_data_source(fs)
+
+    resp = tools()
+    assert len(resp) == 3
+    resp_ids = [tool.id for tool in resp]
+    assert TOOL_ID in resp_ids
+    assert 'tool--03342581-f790-4f03-ba41-e82e67392e23' in resp_ids
+    assert 'tool--242f3da3-4425-4d11-8f5c-b842886da966' in resp_ids

stix2/workbench.py

@@ -13,10 +13,11 @@ from . import Report as _Report
 from . import ThreatActor as _ThreatActor
 from . import Tool as _Tool
 from . import Vulnerability as _Vulnerability
+from .datastore.filters import Filter
+from .datastore.memory import MemoryStore
 from .environment import Environment
-from .sources.filters import Filter
-from .sources.memory import MemoryStore
 
+# Use an implicit MemoryStore
 _environ = Environment(store=MemoryStore())
 
 create = _environ.create
@@ -31,6 +32,7 @@ add_filters = _environ.add_filters
 add_filter = _environ.add_filter
 parse = _environ.parse
 add_data_source = _environ.source.add_data_source
+add_data_sources = _environ.source.add_data_sources
 
 
 # Wrap SDOs with helper functions