def _WrapContentGenerator(self, generator, args, username):
if args.file_path:
path_type, components = rdf_objects.ParseCategorizedPath(
args.file_path)
vfs_file_ref = rdf_objects.VfsFileReference(
client_id=args.client_id,
path_type=path_type,
path_components=components)
else:
vfs_file_ref = rdf_objects.VfsFileReference(
client_id=args.client_id)
object_reference = rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=vfs_file_ref)
try:
for item in generator:
yield item
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_FILE_ARCHIVE_GENERATED,
"Downloaded an archive of folder %s from client %s." %
(args.file_path, args.client_id), object_reference)
except Exception as e:
notification.Notify(
username, rdf_objects.UserNotification.Type.
TYPE_FILE_ARCHIVE_GENERATION_FAILED,
"Archive generation failed for folder %s on client %s: %s" %
(args.file_path, args.client_id, e), object_reference)
raise
def NotifyAboutEnd(self):
"""Sends a notification that this flow is done."""
if not self.state.urn:
super(ListDirectory, self).NotifyAboutEnd()
return
st = self.state.stat
ps_path_type = st.pathspec.last.pathtype
path_type = rdf_objects.PathInfo.PathTypeFromPathspecPathType(
ps_path_type)
full_path = st.pathspec.CollapsePath()
path_components = full_path.strip("/").split("/")
file_ref = rdf_objects.VfsFileReference(
client_id=self.client_id,
path_type=path_type,
path_components=path_components)
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_LIST_DIRECTORY_COMPLETED, "Listed {0}".format(full_path),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def testRegistryPathIsConvertedToURNCorrectly(self):
v = rdf_objects.VfsFileReference(client_id=self.client_id,
path_type="REGISTRY",
path_components=["a", "b", "c"])
self.assertEqual(
v.ToURN(),
rdfvalue.RDFURN("aff4:/%s/registry/a/b/c" % self.client_id))
def NotifyAboutEnd(self):
status_text = "Recursive Directory Listing complete %d nodes, %d dirs"
urn = self.state.first_directory
if not urn:
try:
urn = self.args.pathspec.AFF4Path(self.client_urn)
except ValueError:
pass
if urn:
components = urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_RECURSIVE_LIST_DIRECTORY_COMPLETED,
status_text % (self.state.file_count, self.state.dir_count),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def NotifyAboutEnd(self):
super(GetFile, self).NotifyAboutEnd()
stat_entry = self.state.stat_entry
if not stat_entry:
stat_entry = rdf_client_fs.StatEntry(pathspec=self.args.pathspec)
urn = stat_entry.AFF4Path(self.client_urn)
components = urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
if not self.state.get("success"):
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_FILE_COLLECTION_FAILED, "File transfer failed.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
else:
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File transferred successfully.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def GenerateNotifications(cls, client_id, token):
"""Generates fake notifications of different notification types."""
session_id = flow_test_lib.StartFlow(discovery.Interrogate,
client_id=client_id,
creator=token.username)
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"Fake discovery message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=client_id.Basename())))
# ViewObject: VirtualFileSystem
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File fetch completed",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["proc", "10", "exe"])))
gui_test_lib.CreateFileVersion(client_id,
"fs/os/proc/10/exe",
b"",
timestamp=gui_test_lib.TIME_0,
token=token)
# ViewObject: Flow
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_COMPLETED,
"Fake view flow message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=rdf_objects.FlowReference(client_id=client_id.Basename(),
flow_id=session_id)))
# FlowError
if data_store.RelationalDBEnabled():
flow_base.TerminateFlow(client_id.Basename(), session_id,
"Fake flow error")
else:
with aff4.FACTORY.Open(client_id.Add("flows").Add(session_id),
mode="rw",
token=token) as flow_obj:
flow_obj.GetRunner().Error("Fake flow error")
return session_id
def testVfsNotificationWithInvalidReferenceIsParsedDefensively(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=self.client_id,
# UNSET path type is an invalid value here:
# it make it impossible to find the file.
path_type=rdf_objects.PathInfo.PathType.UNSET,
path_components=["foo", "bar"])))
self.assertEqual(n.reference.type, "UNSET")
def testVfsListDirectoryCompletedIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_VFS_LIST_DIRECTORY_COMPLETED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=self.client_id,
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["foo", "bar"])))
self.assertEqual(n.reference.type, "VFS")
self.assertEqual(n.reference.vfs.client_id.ToString(), self.client_id)
self.assertEqual(n.reference.vfs.vfs_path, "fs/os/foo/bar")
def testVfsNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=self.client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["foo", "bar"])))
self.assertEqual(n.reference.type, "VFS")
self.assertEqual(n.reference.vfs.client_id.ToClientURN(), self.client_id)
self.assertEqual(n.reference.vfs.vfs_path, "fs/os/foo/bar")
def GenerateNotifications(cls, client_id, username):
"""Generates fake notifications of different notification types."""
session_id = flow_test_lib.StartFlow(discovery.Interrogate,
client_id=client_id,
creator=username)
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"Fake discovery message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(client_id=client_id)))
# ViewObject: VirtualFileSystem
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File fetch completed",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=client_id,
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["proc", "10", "exe"])))
gui_test_lib.CreateFileVersion(client_id,
"fs/os/proc/10/exe",
b"",
timestamp=gui_test_lib.TIME_0)
# ViewObject: Flow
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_COMPLETED,
"Fake view flow message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=rdf_objects.FlowReference(client_id=client_id,
flow_id=session_id)))
# FlowError
flow_base.TerminateFlow(client_id, session_id, "Fake flow error")
return session_id
def _SendNotifications(username, client_id):
with test_lib.FakeTime(42):
notification.Notify(
username, rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"<some message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(client_id=client_id.Basename())))
with test_lib.FakeTime(44):
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTION_FAILED,
"<some other message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["foo"])))
def NotifyAboutEnd(self):
"""Sends a notification that this flow is done."""
if not self.state.urn:
super(ListDirectoryMixin, self).NotifyAboutEnd()
return
components = self.state.urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_LIST_DIRECTORY_COMPLETED,
"Listed {0}".format(self.args.pathspec.CollapsePath()),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def NotifyAboutEnd(self):
super().NotifyAboutEnd()
stat_entry = self.state.stat_entry
if not stat_entry:
stat_entry = rdf_client_fs.StatEntry(pathspec=self.state.target_pathspec)
urn = stat_entry.AFF4Path(self.client_urn)
components = urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
if self.state.num_bytes_collected >= self.state.file_size:
notification.Notify(
self.creator,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File transferred successfully.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
elif self.state.num_bytes_collected > 0:
notification.Notify(
self.creator,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File transferred partially (%d bytes out of %d)." %
(self.state.num_bytes_collected, self.state.file_size),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
else:
notification.Notify(
self.creator,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTION_FAILED,
"File transfer failed.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def _WrapContentGenerator(self, generator, args, username):
try:
for item in generator:
yield item
except Exception as e:
path_type, components = rdf_objects.ParseCategorizedPath(
args.file_path)
vfs_file_ref = rdf_objects.VfsFileReference(
client_id=args.client_id,
path_type=path_type,
path_components=components)
object_reference = rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=vfs_file_ref)
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_FILE_BLOB_FETCH_FAILED,
"File blob fetch failed for path %s on client %s: %s" %
(args.client_id, args.file_path, e), object_reference)
raise
def NotifyAboutEnd(self):
if not self.runner.ShouldSendNotifications():
return
if self.state.urn:
components = self.state.urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
notification.Notify(
self.token.username,
notification.UserNotification.Type.TYPE_VFS_LIST_DIRECTORY_COMPLETED,
"List of {0} completed.".format(utils.SmartStr(self.args.pathspec)),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
else:
super(IteratedListDirectory, self).NotifyAboutEnd()
def testTempPathIsConvertedVfsPathStringCorrectly(self):
v = rdf_objects.VfsFileReference(
client_id=self.client_id,
path_type="TEMP",
path_components=["a", "b", "c"])
self.assertEqual(v.ToPath(), "temp/a/b/c")
def testConvertingPathVfsPathStringWithUnknownTypeRaises(self):
with self.assertRaises(ValueError):
rdf_objects.VfsFileReference().ToPath()
def testRegistryPathIsConvertedVfsPathStringCorrectly(self):
v = rdf_objects.VfsFileReference(
client_id=self.client_id,
path_type="REGISTRY",
path_components=["a", "b", "c"])
self.assertEqual(v.ToPath(), "registry/a/b/c")
