From 42b2745d5f109acc6b0b94b56ff075595196fca8 Mon Sep 17 00:00:00 2001 From: Joachim Metz Date: Sat, 6 Apr 2024 11:03:54 +0200 Subject: [PATCH] Refactored Redis store --- plaso/storage/redis/redis_store.py | 53 +++++++++++++++++++----------- 1 file changed, 34 insertions(+), 19 deletions(-) diff --git a/plaso/storage/redis/redis_store.py b/plaso/storage/redis/redis_store.py index 60bd88e235..6b56527ad2 100644 --- a/plaso/storage/redis/redis_store.py +++ b/plaso/storage/redis/redis_store.py @@ -112,7 +112,18 @@ def _WriteExistingAttributeContainer(self, container): redis_key = identifier.CopyToString() json_dict = self._json_serializer.ConvertAttributeContainerToJSON(container) - json_string = json.dumps(json_dict) + + try: + json_string = json.dumps(json_dict) + except TypeError as exception: + raise IOError(( + f'Unable to serialize attribute container: ' + f'{container.CONTAINER_TYPE:s} with error: {exception!s}.')) + + if not json_string: + raise IOError(( + f'Unable to serialize attribute container: ' + f'{container.CONTAINER_TYPE:s}')) self._redis_client.hset(redis_hash_name, key=redis_key, value=json_string) @@ -133,7 +144,18 @@ def _WriteNewAttributeContainer(self, container): redis_key = identifier.CopyToString() json_dict = self._json_serializer.ConvertAttributeContainerToJSON(container) - json_string = json.dumps(json_dict) + + try: + json_string = json.dumps(json_dict) + except TypeError as exception: + raise IOError(( + f'Unable to serialize attribute container: ' + f'{container.CONTAINER_TYPE:s} with error: {exception!s}.')) + + if not json_string: + raise IOError(( + f'Unable to serialize attribute container: ' + f'{container.CONTAINER_TYPE:s}')) self._redis_client.hsetnx(redis_hash_name, redis_key, json_string) @@ -449,19 +471,18 @@ def GetAttributeContainerByIndex(self, container_type, index): RedisAttributeContainerStore, self).GetAttributeContainerByIndex( container_type, index) - sequence_number = index + 1 + identifier = containers_interface.AttributeContainerIdentifier( + name=container_type, sequence_number=index + 1) + redis_hash_name = self._GetRedisHashName(container_type) + redis_key = identifier.CopyToString() - serialized_data = self._redis_client.hget( - redis_hash_name, f'{container_type:s}.{sequence_number:d}') + serialized_data = self._redis_client.hget(redis_hash_name, redis_key) if not serialized_data: return None container = self._DeserializeAttributeContainer( container_type, serialized_data) - - identifier = containers_interface.AttributeContainerIdentifier( - name=container_type, sequence_number=sequence_number) container.SetIdentifier(identifier) if container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_DATA: @@ -495,16 +516,12 @@ def GetAttributeContainers(self, container_type, filter_expression=None): redis_hash_name = self._GetRedisHashName(container_type) for redis_key, serialized_data in self._redis_client.hscan_iter( redis_hash_name): - redis_key = redis_key.decode('utf-8') - container = self._DeserializeAttributeContainer( container_type, serialized_data) - _, sequence_number = redis_key.split('.') - sequence_number = int(sequence_number, 10) - identifier = containers_interface.AttributeContainerIdentifier( - name=container_type, sequence_number=sequence_number) - container.SetIdentifier(identifier) + redis_key = redis_key.decode('utf-8') + identifier = containers_interface.AttributeContainerIdentifier() + identifier.CopyFromString(redis_key) if container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_DATA: identifier = getattr(container, '_event_data_stream_identifier', None) @@ -538,11 +555,9 @@ def GetSortedEvents(self, time_range=None): for redis_key, _ in self._redis_client.zscan_iter(event_index_name): redis_key = redis_key.decode('utf-8') + identifier = containers_interface.AttributeContainerIdentifier() + identifier.CopyFromString(redis_key) - container_type, sequence_number = redis_key.split('.') - sequence_number = int(sequence_number, 10) - identifier = containers_interface.AttributeContainerIdentifier( - name=container_type, sequence_number=sequence_number) yield self.GetAttributeContainerByIdentifier( self._CONTAINER_TYPE_EVENT, identifier)