"""Helper to write data streams."""
import os
from dfimagetools import definitions
[docs]
class DataStreamWriter:
"""Data stream writer."""
# Class constant that defines the default read buffer size.
_READ_BUFFER_SIZE = 16 * 1024 * 1024
_INVALID_PATH_CHARACTERS = [
os.path.sep,
"!",
"$",
"%",
"&",
"*",
"+",
":",
";",
"<",
">",
"?",
"@",
"|",
"~",
]
_INVALID_PATH_CHARACTERS.extend(definitions.NON_PRINTABLE_CHARACTERS.keys())
[docs]
def __init__(self):
"""Initializes a data stream writer."""
super().__init__()
self._invalid_path_characters = str.maketrans(
{value: "_" for value in self._INVALID_PATH_CHARACTERS}
)
[docs]
def GetDisplayPath(self, source_path_segments, source_data_stream_name):
"""Retrieves a path to display.
Args:
source_path_segments (list[str]): path segment of the source file.
source_data_stream_name (str): name of the data stream of the source file.
Returns:
str: display path.
"""
path_segments = [
path_segment.translate(
definitions.NON_PRINTABLE_CHARACTER_TRANSLATION_TABLE
)
for path_segment in source_path_segments
]
display_path = "/".join(path_segments)
if source_data_stream_name:
display_path = ":".join([display_path, source_data_stream_name])
return display_path
[docs]
def GetSanitizedPath(
self, source_path_segments, source_data_stream_name, target_path
):
"""Retrieves santized a path.
This function replaces non-printable and other invalid path characters with
an underscore "_".
Args:
source_path_segments (list[str]): path segment of the source file.
source_data_stream_name (str): name of the data stream of the source file.
target_path (str): path of the target directory.
Returns:
str: sanitized path.
"""
path_segments = [
path_segment.translate(self._invalid_path_characters)
for path_segment in source_path_segments
]
destination_path = os.path.join(target_path, *path_segments)
if source_data_stream_name:
source_data_stream_name = source_data_stream_name.translate(
self._invalid_path_characters
)
destination_path = "_".join([destination_path, source_data_stream_name])
return destination_path
[docs]
def WriteDataStream(self, file_entry, data_stream_name, destination_path):
"""Writes the contents of the source data stream to a destination file.
Note that this function will overwrite an existing file.
Args:
file_entry (dfvfs.FileEntry): file entry whose content is to be written.
data_stream_name (str): name of the data stream whose content is to be
written.
destination_path (str): path of the destination file.
"""
source_file_object = file_entry.GetFileObject(data_stream_name=data_stream_name)
if source_file_object:
with open(destination_path, "wb") as destination_file_object:
source_file_object.seek(0, os.SEEK_SET)
data = source_file_object.read(self._READ_BUFFER_SIZE)
while data:
destination_file_object.write(data)
data = source_file_object.read(self._READ_BUFFER_SIZE)