def testHasAttributeContainers(self):
"""Tests the HasAttributeContainers method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertFalse(result)
test_store.AddAttributeContainer(event_data_stream)
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertTrue(result)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainerByIdentifier(self):
"""Tests the _GetAttributeContainerByIdentifier method."""
test_events = []
for test_event, _, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_events.append(test_event)
test_event_tags = self._CreateTestEventTags(test_events)
test_event_tag = test_event_tags[0]
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
store.AddEventTag(test_event_tag)
identifier = test_event_tag.GetIdentifier()
retrieved_event_tag = store._GetAttributeContainerByIdentifier(
test_event_tag.CONTAINER_TYPE, identifier)
test_event_tag_dict = test_event_tag.CopyToDict()
retrieved_event_tag_dict = retrieved_event_tag.CopyToDict()
self.assertEqual(test_event_tag_dict, retrieved_event_tag_dict)
store.Close()
def testGetRedisHashName(self):
"""Tests the _GetRedisHashName function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
redis_hash_name = test_store._GetRedisHashName(
event_data_stream.CONTAINER_TYPE)
expected_redis_hash_name = '{0:s}-{1:s}-{2:s}'.format(
task.session_identifier, task.identifier,
event_data_stream.CONTAINER_TYPE)
self.assertEqual(redis_hash_name, expected_redis_hash_name)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetSerializedAttributeContainers(self):
"""Tests the GetSerializedAttributeContainers method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event_data)
cursor, serialized_containers = (
test_store.GetSerializedAttributeContainers(
'event_data', 0, 0))
self.assertEqual(len(serialized_containers), 4)
for serialized_container in serialized_containers:
self.assertIsInstance(serialized_container, bytes)
self.assertIsInstance(cursor, int)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetSortedEvents(self):
"""Tests the GetSortedEvents method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
for event, _, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event)
retrieved_events = list(test_store.GetSortedEvents())
self.assertEqual(len(retrieved_events), 4)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def Open(
self, redis_client=None, session_identifier=None, task_identifier=None,
**unused_kwargs):
"""Opens the storage writer.
Args:
redis_client (Optional[Redis]): Redis client to query. If specified, no
new client will be created. If no client is specified a new client
will be opened connected to the Redis instance specified by 'url'.
session_identifier (Optional[str]): session identifier.
task_identifier (Optional[str]): task identifier.
Raises:
IOError: if the storage writer is already opened.
OSError: if the storage writer is already opened.
"""
if self._store:
raise IOError('Storage writer already opened.')
self._store = redis_store.RedisStore(storage_type=self._storage_type)
if self._serializers_profiler:
self._store.SetSerializersProfiler(self._serializers_profiler)
if self._storage_profiler:
self._store.SetStorageProfiler(self._storage_profiler)
self._store.Open(
redis_client=redis_client, session_identifier=session_identifier,
task_identifier=task_identifier)
def testGetNumberOfAttributeContainers(self):
"""Tests the GetNumberOfAttributeContainers function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
test_store.AddAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainerByIndex(self):
"""Tests the GetAttributeContainerByIndex function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNone(container)
test_store.AddAttributeContainer(event_data_stream)
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNotNone(container)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainerByIdentifier(self):
"""Tests the GetAttributeContainerByIdentifier method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore()
test_store.Open(
redis_client=redis_client, session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
test_store.AddAttributeContainer(event_data_stream)
identifier = event_data_stream.GetIdentifier()
container = test_store.GetAttributeContainerByIdentifier(
event_data_stream.CONTAINER_TYPE, identifier)
self.assertIsNotNone(container)
identifier.sequence_number = 99
container = test_store.GetAttributeContainerByIdentifier(
event_data_stream.CONTAINER_TYPE, identifier)
self.assertIsNone(container)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testAddAttributeContainer(self):
"""Tests the AddAttributeContainer method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore()
test_store.Open(
redis_client=redis_client, session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
test_store.AddAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
with self.assertRaises(IOError):
test_store.AddAttributeContainer(event_data_stream)
def testRemoveAttributeContainer(self):
"""Tests the RemoveAttributeContainer method."""
redis_client = self._CreateRedisClient()
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
test_store.AddAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
identifier = event_data_stream.GetIdentifier()
test_store.RemoveAttributeContainer(event_data_stream.CONTAINER_TYPE,
identifier)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
test_store.Close()
def testWriteExistingAttributeContainer(self):
"""Tests the _WriteExistingAttributeContainer function."""
redis_client = self._CreateRedisClient()
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
with self.assertRaises(IOError):
test_store._WriteExistingAttributeContainer(event_data_stream)
test_store._WriteNewAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
test_store._WriteExistingAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
test_store.Close()
def testGetEventTagByEventIdentifier(self):
"""Tests the GetEventTagByEventIdentifier function."""
redis_client = self._CreateRedisClient()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
index = 0
for event, event_data, event_data_stream in (
containers_test_lib.CreateEventsFromValues(self._TEST_EVENTS)):
test_store.AddAttributeContainer(event_data_stream)
event_data.SetEventDataStreamIdentifier(
event_data_stream.GetIdentifier())
test_store.AddAttributeContainer(event_data)
event.SetEventDataIdentifier(event_data.GetIdentifier())
test_store.AddAttributeContainer(event)
if index == 1:
event_tag = events.EventTag()
event_tag.AddLabels(['Malware', 'Benign'])
event_identifier = event.GetIdentifier()
event_tag.SetEventIdentifier(event_identifier)
test_store.AddAttributeContainer(event_tag)
index += 1
test_store.Close()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
test_event = test_store.GetAttributeContainerByIndex(
events.EventObject.CONTAINER_TYPE, 1)
# Note that this method is not implemented.
self.assertIsNone(test_event)
test_store.Close()
def __init__(self, task):
"""Initializes a Redis Storage Reader.
Args:
task (Task): the task whose results the store contains.
"""
super(RedisStorageReader, self).__init__()
self._store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
def testFinalization(self):
"""Tests that a store is correctly finalized."""
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
self.assertFalse(store.IsFinalized())
store.Close()
store.Open(redis_client=redis_client)
self.assertFalse(store.IsFinalized())
store.Finalize()
self.assertTrue(store.IsFinalized())
store.Close()
def testGetSortedEvents(self):
"""Tests the GetSortedEvents method."""
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
for event, _, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
store.AddEvent(event)
retrieved_events = list(store.GetSortedEvents())
self.assertEqual(len(retrieved_events), 4)
store.Close()
def __init__(self, session_identifier, task_identifier, redis_client=None):
"""Initializes a Redis storage reader.
Args:
session_identifier (str): session identifier.
task_identifier (str): task identifier.
redis_client (Optional[Redis]): Redis client to query. If specified, no
new client will be created. If no client is specified a new client
will be opened connected to the Redis instance specified by 'url'.
"""
super(RedisStorageReader, self).__init__()
self._store = redis_store.RedisStore()
self._store.Open(redis_client=redis_client,
session_identifier=session_identifier,
task_identifier=task_identifier)
def testGetAttributeContainers(self):
"""Tests the _GetAttributeContainers method."""
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
store.AddEventData(event_data)
retrieved_event_datas = list(
store._GetAttributeContainers(store._CONTAINER_TYPE_EVENT_DATA))
self.assertEqual(len(retrieved_event_datas), 4)
store.Close()
def testGetSortedEvents(self):
"""Tests the GetSortedEvents method."""
redis_client = self._CreateRedisClient()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
for event, _, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event)
retrieved_events = list(test_store.GetSortedEvents())
self.assertEqual(len(retrieved_events), 4)
test_store.Close()
def Open(self, redis_client=None, **unused_kwargs):
"""Opens the storage writer.
Raises:
IOError: if the storage writer is already opened.
OSError: if the storage writer is already opened.
"""
if self._store:
raise IOError('Storage writer already opened.')
self._store = redis_store.RedisStore(
storage_type=self._storage_type,
session_identifier=self._task.session_identifier,
task_identifier=self._task.identifier)
self._store.Open(redis_client=redis_client)
def testGetAttributeContainers(self):
"""Tests the GetAttributeContainers method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
event_data_stream.md5_hash = '8f0bf95a7959baad9666b21a7feed79d'
containers = list(
test_store.GetAttributeContainers(
event_data_stream.CONTAINER_TYPE))
self.assertEqual(len(containers), 0)
test_store.AddAttributeContainer(event_data_stream)
containers = list(
test_store.GetAttributeContainers(
event_data_stream.CONTAINER_TYPE))
self.assertEqual(len(containers), 1)
filter_expression = 'md5_hash == "8f0bf95a7959baad9666b21a7feed79d"'
containers = list(
test_store.GetAttributeContainers(
event_data_stream.CONTAINER_TYPE,
filter_expression=filter_expression))
self.assertEqual(len(containers), 1)
filter_expression = 'md5_hash != "8f0bf95a7959baad9666b21a7feed79d"'
containers = list(
test_store.GetAttributeContainers(
event_data_stream.CONTAINER_TYPE,
filter_expression=filter_expression))
self.assertEqual(len(containers), 0)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainers(self):
"""Tests the GetAttributeContainers method."""
redis_client = self._CreateRedisClient()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event_data)
containers = list(
test_store.GetAttributeContainers(
test_store._CONTAINER_TYPE_EVENT_DATA))
self.assertEqual(len(containers), 4)
test_store.Close()
def testAddEvent(self):
"""Tests the _AddEvent method."""
event, _, _ = containers_test_lib.CreateEventFromValues(
self._TEST_EVENTS[0])
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
self.assertEqual(
store._GetNumberOfAttributeContainers(event.CONTAINER_TYPE), 0)
store.AddEvent(event)
self.assertEqual(
store._GetNumberOfAttributeContainers(event.CONTAINER_TYPE), 1)
store.Close()
def testGetSerializedAttributeContainers(self):
"""Tests the GetSerializedAttributeContainers method."""
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
store.AddEventData(event_data)
cursor, serialized_containers = store.GetSerializedAttributeContainers(
'event_data', 0, 0)
self.assertEqual(len(serialized_containers), 4)
for serialized_container in serialized_containers:
self.assertIsInstance(serialized_container, bytes)
self.assertIsInstance(cursor, int)
store.Close()
def testAddAttributeContainers(self):
"""Tests the _AddAttributeContainer method."""
event_data = events.EventData()
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
self.assertEqual(
store._GetNumberOfAttributeContainers(event_data.CONTAINER_TYPE), 0)
store._AddAttributeContainer(store._CONTAINER_TYPE_EVENT_DATA, event_data)
self.assertEqual(
store._GetNumberOfAttributeContainers(event_data.CONTAINER_TYPE), 1)
has_containers = store._HasAttributeContainers(event_data.CONTAINER_TYPE)
self.assertTrue(has_containers)
store.Close()
def testGetSerializedAttributeContainers(self):
"""Tests the GetSerializedAttributeContainers method."""
redis_client = self._CreateRedisClient()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event_data)
cursor, serialized_containers = test_store.GetSerializedAttributeContainers(
'event_data', 0, 0)
self.assertEqual(len(serialized_containers), 4)
for serialized_container in serialized_containers:
self.assertIsInstance(serialized_container, bytes)
self.assertIsInstance(cursor, int)
test_store.Close()
def testRemoveAttributeContainer(self):
"""Tests the RemoveAttributeContainer method."""
event_data = events.EventData()
store = redis_store.RedisStore()
redis_client = self._GetRedisClient()
store.Open(redis_client=redis_client)
store._AddAttributeContainer(store._CONTAINER_TYPE_EVENT_DATA, event_data)
self.assertEqual(
store._GetNumberOfAttributeContainers(event_data.CONTAINER_TYPE), 1)
identifier = event_data.GetIdentifier()
store.RemoveAttributeContainer(
store._CONTAINER_TYPE_EVENT_DATA, identifier)
self.assertEqual(
store._GetNumberOfAttributeContainers(event_data.CONTAINER_TYPE), 0)
store.Close()
def testHasAttributeContainers(self):
"""Tests the HasAttributeContainers method."""
redis_client = self._CreateRedisClient()
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertFalse(result)
test_store.AddAttributeContainer(event_data_stream)
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertTrue(result)
test_store.Close()
def __init__(self, storage_writer, task, redis_client=None):
"""Initializes a Redis storage merge reader.
Args:
storage_writer (StorageWriter): storage writer.
task (Task): the task whose store is being merged.
redis_client (Optional[Redis]): Redis client to query. If specified, no
new client will be created.
Raises:
RuntimeError: if an add container method is missing.
"""
super(RedisMergeReader, self).__init__(storage_writer)
self._active_container_type = None
self._container_types = []
self._active_cursor = 0
self._add_active_container_method = None
self._store = redis_store.RedisStore(
definitions.STORAGE_TYPE_TASK,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
self._store.Open(redis_client=redis_client)
self._event_data_identifier_mappings = {}
self._event_data_stream_identifier_mappings = {}
self._add_container_type_methods = {}
self._active_extra_containers = []
# Create a runtime lookup table for the add container type method. This
# prevents having to create a series of if-else checks for container types.
# The table is generated at runtime as there are no forward function
# declarations in Python.
for container_type, method_name in self._ADD_CONTAINER_TYPE_METHODS.items(
):
method = getattr(self, method_name, None)
if not method:
raise RuntimeError(
'Add method missing for container type: {0:s}'.format(
container_type))
self._add_container_type_methods[container_type] = method
def testGetAttributeContainerByIndex(self):
"""Tests the GetAttributeContainerByIndex function."""
redis_client = self._CreateRedisClient()
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client)
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNone(container)
test_store.AddAttributeContainer(event_data_stream)
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNotNone(container)
with self.assertRaises(IOError):
test_store.GetAttributeContainerByIndex('bogus', 0)
test_store.Close()
