From 53c2d4fadfa0fffc63a05e89fc86ba6394271e98 Mon Sep 17 00:00:00 2001
From: Chris Lenk <clenk@mitre.org>
Date: Wed, 14 Mar 2018 14:33:45 -0400
Subject: [PATCH] Allow add'l filters in workbench query functions

---
 stix2/test/test_workbench.py | 11 +++++++
 stix2/workbench.py           | 59 +++++++++++++++++++++---------------
 2 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/stix2/test/test_workbench.py b/stix2/test/test_workbench.py
index b84b529..324011b 100644
--- a/stix2/test/test_workbench.py
+++ b/stix2/test/test_workbench.py
@@ -189,3 +189,14 @@ def test_add_data_source():
     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
+
+
+def test_additional_filter():
+    resp = tools(stix2.Filter('created_by_ref', '=', 'identity--c78cb6e5-0c4b-4611-8297-d1b8b55e40b5'))
+    assert len(resp) == 2
+
+
+def test_additional_filters_list():
+    resp = tools([stix2.Filter('created_by_ref', '=', 'identity--c78cb6e5-0c4b-4611-8297-d1b8b55e40b5'),
+                  stix2.Filter('name', '=', 'Windows Credential Editor')])
+    assert len(resp) == 1
diff --git a/stix2/workbench.py b/stix2/workbench.py
index 60d1165..49e41c9 100644
--- a/stix2/workbench.py
+++ b/stix2/workbench.py
@@ -80,49 +80,60 @@ for obj_type in STIX_OBJS:
 # Functions to get all objects of a specific type
 
 
-def attack_patterns():
-    return query(Filter('type', '=', 'attack-pattern'))
+def query_by_type(obj_type='indicator', filters=None):
+    filter_list = [Filter('type', '=', obj_type)]
+    if filters:
+        if isinstance(filters, list):
+            filter_list += filters
+        else:
+            filter_list.append(filters)
+
+    return query(filter_list)
 
 
-def campaigns():
-    return query(Filter('type', '=', 'campaign'))
+def attack_patterns(filters=None):
+    return query_by_type('attack-pattern', filters)
 
 
-def courses_of_action():
-    return query(Filter('type', '=', 'course-of-action'))
+def campaigns(filters=None):
+    return query_by_type('campaign', filters)
 
 
-def identities():
-    return query(Filter('type', '=', 'identity'))
+def courses_of_action(filters=None):
+    return query_by_type('course-of-action', filters)
 
 
-def indicators():
-    return query(Filter('type', '=', 'indicator'))
+def identities(filters=None):
+    return query_by_type('identity', filters)
 
 
-def intrusion_sets():
-    return query(Filter('type', '=', 'intrusion-set'))
+def indicators(filters=None):
+    return query_by_type('indicator', filters)
 
 
-def malware():
-    return query(Filter('type', '=', 'malware'))
+def intrusion_sets(filters=None):
+    return query_by_type('intrusion-set', filters)
 
 
-def observed_data():
-    return query(Filter('type', '=', 'observed-data'))
+def malware(filters=None):
+    return query_by_type('malware', filters)
 
 
-def reports():
-    return query(Filter('type', '=', 'report'))
+def observed_data(filters=None):
+    return query_by_type('observed-data', filters)
 
 
-def threat_actors():
-    return query(Filter('type', '=', 'threat-actor'))
+def reports(filters=None):
+    return query_by_type('report', filters)
 
 
-def tools():
-    return query(Filter('type', '=', 'tool'))
+def threat_actors(filters=None):
+    return query_by_type('threat-actor', filters)
 
 
-def vulnerabilities():
-    return query(Filter('type', '=', 'vulnerability'))
+def tools(filters=None):
+    return query_by_type('tool', filters)
+
+
+def vulnerabilities(filters=None):
+    return query_by_type('vulnerability', filters)