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