def _HandleRelational(self, args, token=None): iop_cls = instant_output_plugin.InstantOutputPlugin plugin_cls = iop_cls.GetPluginClassByPluginName(args.plugin_name) # TODO(user): Instant output plugins shouldn't depend on tokens # and URNs. flow_urn = args.flow_id.ResolveClientFlowURN(args.client_id, token=token) plugin = plugin_cls(source_urn=flow_urn, token=token) client_id = str(args.client_id) flow_id = str(args.flow_id) types = data_store.REL_DB.CountFlowResultsByType(client_id, flow_id) def FetchFn(type_name): for r in data_store.REL_DB.ReadFlowResults( client_id, flow_id, offset=0, count=db.MAX_COUNT, with_type=type_name): msg = r.AsLegacyGrrMessage() msg.source = client_id yield msg content_generator = instant_output_plugin.ApplyPluginToTypedCollection( plugin, types, FetchFn) return api_call_handler_base.ApiBinaryStream( plugin.output_file_name, content_generator=content_generator)
def Handle(self, args, token=None): iop_cls = instant_output_plugin.InstantOutputPlugin plugin_cls = iop_cls.GetPluginClassByPluginName(args.plugin_name) # TODO(user): Instant output plugins shouldn't depend on tokens # and URNs. flow_urn = rdfvalue.RDFURN("{}/flows/{}".format( args.client_id, args.flow_id)) plugin = plugin_cls(source_urn=flow_urn, token=token) client_id = str(args.client_id) flow_id = str(args.flow_id) types = data_store.REL_DB.CountFlowResultsByType(client_id, flow_id) def FetchFn(type_name): """Fetches all flow results of a given type.""" offset = 0 while True: results = data_store.REL_DB.ReadFlowResults( client_id, flow_id, offset=offset, count=self._RESULTS_PAGE_SIZE, with_type=type_name) if not results: break for r in results: msg = r.AsLegacyGrrMessage() msg.source = client_id yield msg offset += self._RESULTS_PAGE_SIZE content_generator = instant_output_plugin.ApplyPluginToTypedCollection( plugin, types, FetchFn) return api_call_handler_base.ApiBinaryStream( plugin.output_file_name, content_generator=content_generator)
def Handle(self, args, context=None): hunt_id = str(args.hunt_id) source_urn = rdfvalue.RDFURN("hunts").Add(hunt_id) iop_cls = instant_output_plugin.InstantOutputPlugin plugin_cls = iop_cls.GetPluginClassByPluginName(args.plugin_name) # TODO(user): Instant output plugins shouldn't depend on contexts # and URNs. plugin = plugin_cls( source_urn=source_urn, token=access_control.ACLToken(username=context.username)) types = data_store.REL_DB.CountHuntResultsByType(hunt_id) def FetchFn(type_name): """Fetches all hunt results of a given type.""" offset = 0 while True: results = data_store.REL_DB.ReadHuntResults( hunt_id, offset=offset, count=self._RESULTS_PAGE_SIZE, with_type=type_name) if not results: break for r in results: msg = r.AsLegacyGrrMessage() msg.source_urn = source_urn yield msg offset += self._RESULTS_PAGE_SIZE content_generator = instant_output_plugin.ApplyPluginToTypedCollection( plugin, types, FetchFn) return api_call_handler_base.ApiBinaryStream( plugin.output_file_name, content_generator=content_generator)