def _AddAttributeContainer(self, attribute_container):
"""Adds a single attribute container to the storage writer.
Args:
attribute_container (AttributeContainer): container
Raises:
RuntimeError: if the attribute container type is not supported.
"""
container_type = attribute_container.CONTAINER_TYPE
if container_type == 'event_source':
self._storage_writer.AddEventSource(attribute_container)
elif container_type == 'event_data':
identifier = attribute_container.GetIdentifier()
lookup_key = identifier.CopyToString()
self._storage_writer.AddEventData(attribute_container)
identifier = attribute_container.GetIdentifier()
self._event_data_identifier_mappings[lookup_key] = identifier
elif container_type == 'event':
if (hasattr(attribute_container, 'event_data_stream_number') and
hasattr(attribute_container, 'event_data_entry_index')):
event_data_identifier = identifiers.SerializedStreamIdentifier(
attribute_container.event_data_stream_number,
attribute_container.event_data_entry_index)
lookup_key = event_data_identifier.CopyToString()
event_data_identifier = self._event_data_identifier_mappings[
lookup_key]
attribute_container.SetEventDataIdentifier(
event_data_identifier)
del attribute_container.event_data_stream_number
del attribute_container.event_data_entry_index
# TODO: add event identifier mappings for event tags.
self._storage_writer.AddEvent(attribute_container)
elif container_type == 'event_tag':
event_identifier = identifiers.SerializedStreamIdentifier(
attribute_container.event_stream_number,
attribute_container.event_entry_index)
attribute_container.SetEventIdentifier(event_identifier)
self._storage_writer.AddEventTag(attribute_container)
elif container_type == 'extraction_error':
self._storage_writer.AddError(attribute_container)
elif container_type == 'analysis_report':
self._storage_writer.AddAnalysisReport(attribute_container)
elif container_type not in ('task_completion', 'task_start'):
raise RuntimeError(
'Unsupported container type: {0:s}'.format(container_type))
def _AddAttributeContainer(self, attribute_container):
"""Adds a single attribute container to the storage writer.
Args:
attribute_container (AttributeContainer): container
Raises:
RuntimeError: if the attribute container type is not supported.
"""
container_type = attribute_container.CONTAINER_TYPE
if container_type == u'event_source':
self._storage_writer.AddEventSource(attribute_container)
elif container_type == u'event':
self._storage_writer.AddEvent(attribute_container)
elif container_type == u'event_tag':
event_identifier = identifiers.SerializedStreamIdentifier(
attribute_container.event_stream_number,
attribute_container.event_entry_index)
attribute_container.SetEventIdentifier(event_identifier)
self._storage_writer.AddEventTag(attribute_container)
elif container_type == u'extraction_error':
self._storage_writer.AddError(attribute_container)
elif container_type == u'analysis_report':
self._storage_writer.AddAnalysisReport(attribute_container)
elif container_type not in (u'task_completion', u'task_start'):
raise RuntimeError(
u'Unsupported container type: {0:s}'.format(container_type))
def GetEventTags(self):
"""Retrieves the event tags.
Yields:
EventTag: event tag.
"""
for event_tag in iter(self._GetAttributeContainerList(u'event_tag')):
event_identifier = identifiers.SerializedStreamIdentifier(
event_tag.event_stream_number, event_tag.event_entry_index)
event_tag.SetEventIdentifier(event_identifier)
yield event_tag
def _WriteAttributeContainer(self, attribute_container):
"""Writes an attribute container.
Args:
attribute_container (AttributeContainer): attribute container.
"""
identifier = identifiers.SerializedStreamIdentifier(
1, self._number_of_containers)
attribute_container.SetIdentifier(identifier)
serialized_data = self._SerializeAttributeContainer(
attribute_container)
self._number_of_containers += 1
self._gzip_file.write(serialized_data)
self._gzip_file.write(b'\n')
def _WriteAttributeContainer(self, attribute_container):
"""Writes an attribute container.
Args:
attribute_container (AttributeContainer): attribute container.
Raises:
IOError: when the storage file is closed or read-only.
"""
if not self._is_open:
raise IOError(u'Unable to write to closed storage file.')
if self._read_only:
raise IOError(u'Unable to write to read-only storage file.')
attribute_container_identifier = identifiers.SerializedStreamIdentifier(
1, len(self._attribute_containers))
attribute_container.SetIdentifier(attribute_container_identifier)
attribute_container_data = self._SerializeAttributeContainer(
attribute_container)
self._gzip_file.write(attribute_container_data)
self._gzip_file.write(b'\n')
def MergeAttributeContainers(self, maximum_number_of_containers=0):
"""Reads attribute containers from a task storage file into the writer.
Args:
maximum_number_of_containers (Optional[int]): maximum number of
containers to merge, where 0 represent no limit.
Returns:
bool: True if the entire task storage file has been merged.
Raises:
OSError: if the task storage file cannot be deleted.
"""
if not self._data_buffer:
# Do not use gzip.readlines() here since it can consume a large amount
# of memory.
self._data_buffer = self._gzip_file.read(self._DATA_BUFFER_SIZE)
self._number_of_containers = 0
number_of_containers = 0
while self._data_buffer:
lines = self._data_buffer.splitlines(True)
self._data_buffer = b''
for index, line in enumerate(lines):
if not line.endswith(b'\n'):
self._data_buffer = b''.join(lines[index:])
continue
identifier = identifiers.SerializedStreamIdentifier(
1, self._number_of_containers)
attribute_container = self._DeserializeAttributeContainer(
'attribute_container', line)
attribute_container.SetIdentifier(identifier)
self._AddAttributeContainer(attribute_container)
self._number_of_containers += 1
number_of_containers += 1
if (maximum_number_of_containers > 0 and
number_of_containers >= maximum_number_of_containers):
self._data_buffer = b''.join(lines[index + 1:])
return False
additional_data_buffer = self._gzip_file.read(
self._DATA_BUFFER_SIZE)
self._data_buffer = b''.join(
[self._data_buffer, additional_data_buffer])
self._gzip_file.close()
self._gzip_file = None
# On Windows the file can sometimes be in use and we have to wait.
for attempt in range(1, self._MAXIMUM_NUMBER_OF_LOCKED_FILE_ATTEMPTS):
try:
os.remove(self._path)
break
except OSError:
if attempt == (self._MAXIMUM_NUMBER_OF_LOCKED_FILE_ATTEMPTS -
1):
raise
time.sleep(self._LOCKED_FILE_SLEEP_TIME)
return True