Clean up TAXII Datastore to use latest TAXII client

stix2.1
Greg Back 2017-07-20 19:36:54 +00:00
parent f4833c05f6
commit b82606ba38
3 changed files with 48 additions and 165 deletions

View File

@ -1,7 +1,7 @@
[settings]
check=1
diff=1
known_third_party=dateutil,pytest,pytz,six,requests
known_third_party=dateutil,pytest,pytz,six,requests,taxii2_client
known_first_party=stix2
not_skip=__init__.py
force_sort_within_sections=1

View File

@ -21,67 +21,28 @@ TAXII_FILTERS = ['added_after', 'id', 'type', 'version']
class TAXIICollectionStore(DataStore):
"""
"""
def __init__(self,
taxii_client=None,
api_root_name=None,
collection_id=None,
user=None,
password=None,
name="TAXIICollectionStore"):
def __init__(self, collection, name="TAXIICollectionStore"):
"""
Create a new TAXII Collection Data store
Args:
collection (taxii2.Collection): Collection instance
"""
self.name = name
self.id = make_id()
self.source = TAXIICollectionSource(taxii_client, api_root_name, collection_id, user, password)
self.sink = TAXIICollectionSink(taxii_client, api_root_name, collection_id, user, password)
# file system sink API calls
def add(self, stix_objs):
return self.sink.add(stix_objs=stix_objs)
# file sytem source API calls
def get(self, stix_id):
return self.source.get(stix_id=stix_id)
def all_versions(self, stix_id):
return self.source.all_versions(stix_id=stix_id)
def query(self, query):
return self.source.query(query=query)
self.source = TAXIICollectionSource(collection)
self.sink = TAXIICollectionSink(collection)
class TAXIICollectionSink(DataSink):
"""
"""
def __init__(self, taxii_client=None, api_root_name=None, collection_id=None, user=None, password=None, name="TAXIICollectionSink"):
def __init__(self, collection, name="TAXIICollectionSink"):
super(TAXIICollectionSink, self).__init__(name=name)
self.taxii_client = taxii_client
self.taxii_client.populate_available_information()
if not api_root_name:
raise ValueError("No api_root specified.")
else:
self.api_root = None
for a_r in self.taxii_client.api_roots:
if api_root_name == a_r.name:
self.api_root = a_r
break
if not self.api_root:
raise ValueError("The api_root %s is not found on this taxii server" % api_root_name)
if not collection_id:
raise ValueError("No collection specified.")
else:
self.collection = None
for c in self.api_root.collections:
if c.id_ == collection_id:
self.collection = c
break
if not self.collection:
raise ValueError("The collection %s is not found on the api_root %s of this taxii server" %
(collection_id, api_root_name))
self.collection = collection
def add(self, stix_obj):
"""
@ -95,53 +56,14 @@ class TAXIICollectionSink(DataSink):
spec_version="2.0",
type="bundle")
# utility functions for the current set collection and api root
def get_api_root_info(self):
"""
"""
return self.api_root.get_information()
def get_api_root_collections(self):
"""
"""
return self.api_root.get_collections()
def get_collection_manifest(self):
"""
"""
return self.collection.get_collection_manifest()
class TAXIICollectionSource(DataSource):
"""
"""
def __init__(self, taxii_client=None, api_root_name=None, collection_id=None, user=None, password=None, name="TAXIICollectionSourc"):
def __init__(self, collection, name="TAXIICollectionSource"):
super(TAXIICollectionSource, self).__init__(name=name)
self.taxii_client = taxii_client
self.taxii_client.populate_available_information()
if not api_root_name:
raise ValueError("No api_root specified.")
else:
self.api_root = None
for a_r in self.taxii_client.api_roots:
if api_root_name == a_r.name:
self.api_root = a_r
break
if not self.api_root:
raise ValueError("The api_root %s is not found on this taxii server" % api_root_name)
if not collection_id:
raise ValueError("No collection specified.")
else:
self.collection = None
for c in self.api_root.collections:
if c.id_ == collection_id:
self.collection = c
break
if not self.collection:
raise ValueError("The collection %s is not found on the api_root %s of this taxii server" %
(collection_id, api_root_name))
self.collection = collection
def get(self, stix_id, _composite_filters=None):
"""
@ -244,54 +166,3 @@ class TAXIICollectionSource(DataSource):
taxii_field = "match[" + filter_["field"] + ']'
params[taxii_field] = filter_["value"]
return params
# utility functions for the current attached collection and api root
def get_api_root_info(self):
"""
"""
return self.api_root.get_information()
def get_api_root_collections(self):
"""
"""
return self.api_root.get_collections()
def get_collection_manifest(self):
"""
"""
return self.collection.get_collection_manifest()
def get_server_api_roots(taxii_client):
"""
"""
api_root_info = []
taxii_client.populate_available_information()
for api_root in taxii_client.api_roots:
api_root_info.append(api_root.information())
return api_root_info
def get_server_collections(taxii_client):
"""
"""
server_collections = []
taxii_client.populate_available_information()
for api_root in taxii_client.api_roots:
server_collections.extend(api_root.get_collections())
return server_collections
def get_api_root_collections(taxii_client, api_root_name):
"""
"""
taxii_client.populate_available_information()
for api_root in taxii_client.api_roots:
if api_root == api_root_name:
return api_root.get_collections()

View File

@ -1,25 +1,37 @@
import pytest
from taxii2_client import Collection
from stix2.sources import taxii
def test_ds_taxii():
ds = taxii.TAXIICollectionSource()
assert ds.name == 'TAXII'
COLLECTION_URL = 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'
def test_ds_taxii_name():
ds = taxii.TAXIICollectionSource(name='My Data Source Name')
class MockTAXIIClient(object):
"""Mock for taxii2_client.TAXIIClient"""
def get(self):
return {}
def post(self):
return {}
@pytest.fixture
def collection():
return Collection(COLLECTION_URL, MockTAXIIClient())
def test_ds_taxii(collection):
ds = taxii.TAXIICollectionSource(collection)
assert ds.name == 'TAXIICollectionSource'
def test_ds_taxii_name(collection):
ds = taxii.TAXIICollectionSource(collection, name='My Data Source Name')
assert ds.name == "My Data Source Name"
def test_ds_params():
url = "http://taxii_url.com:5000"
creds = {"username": "Wade", "password": "Wilson"}
ds = taxii.TAXIICollectionSource(api_root=url, auth=creds)
assert ds.taxii_info['api_root']['url'] == url
assert ds.taxii_info['auth'] == creds
def test_parse_taxii_filters():
def test_parse_taxii_filters(collection):
query = [
{
"field": "added_after",
@ -55,14 +67,14 @@ def test_parse_taxii_filters():
"match[version]": "first"
}
ds = taxii.TAXIICollectionSource()
ds = taxii.TAXIICollectionSource(collection)
taxii_filters = ds._parse_taxii_filters(query)
assert taxii_filters == expected_params
def test_add_get_remove_filter():
def test_add_get_remove_filter(collection):
class dummy(object):
x = 4
@ -114,7 +126,7 @@ def test_add_get_remove_filter():
"Filter 'value' type is not supported. The type(value) must be python immutable type or dictionary"
]
ds = taxii.TAXIICollectionSource()
ds = taxii.TAXIICollectionSource(collection)
# add
ids, statuses = ds.add_filter(filters)
@ -156,7 +168,7 @@ def test_add_get_remove_filter():
assert id_ in ids[:3]
def test_apply_common_filters():
def test_apply_common_filters(collection):
stix_objs = [
{
"created": "2017-01-27T13:49:53.997Z",
@ -210,7 +222,7 @@ def test_apply_common_filters():
}
]
ds = taxii.TAXIICollectionSource()
ds = taxii.TAXIICollectionSource(collection)
resp = ds.apply_common_filters(stix_objs, [filters[0]])
ids = [r['id'] for r in resp]
@ -224,7 +236,7 @@ def test_apply_common_filters():
assert resp[0]['id'] == stix_objs[0]['id']
def test_deduplicate():
def test_deduplicate(collection):
stix_objs = [
{
"created": "2017-01-27T13:49:53.935Z",
@ -288,7 +300,7 @@ def test_deduplicate():
}
]
ds = taxii.TAXIICollectionSource()
ds = taxii.TAXIICollectionSource(collection)
unique = ds.deduplicate(stix_objs)
# Only 3 objects are unique