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()