Clean up TAXII Datastore to use latest TAXII client
parent
f4833c05f6
commit
b82606ba38
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue