def testUpdatingHuntOutputStateWorksCorrectly(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) state_1 = rdf_flow_runner.OutputPluginState( plugin_descriptor=rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin1"), plugin_state={}) state_2 = rdf_flow_runner.OutputPluginState( plugin_descriptor=rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin2"), plugin_state={}) self.db.WriteHuntOutputPluginsStates(hunt_obj.hunt_id, [state_1, state_2]) def Update(s): s["foo"] = "bar" return s self.db.UpdateHuntOutputPluginState(hunt_obj.hunt_id, 0, Update) states = self.db.ReadHuntOutputPluginsStates(hunt_obj.hunt_id) self.assertEqual(states[0].plugin_state, {"foo": "bar"}) self.assertEqual(states[1].plugin_state, {}) self.db.UpdateHuntOutputPluginState(hunt_obj.hunt_id, 1, Update) states = self.db.ReadHuntOutputPluginsStates(hunt_obj.hunt_id) self.assertEqual(states[0].plugin_state, {"foo": "bar"}) self.assertEqual(states[1].plugin_state, {"foo": "bar"})
def GetOutputPluginStates(output_plugins, source=None, token=None): """Initializes state for a list of output plugins.""" output_plugins_states = [] for plugin_descriptor in output_plugins: plugin_class = plugin_descriptor.GetPluginClass() plugin = plugin_class(source_urn=source, args=plugin_descriptor.plugin_args, token=token) try: plugin.InitializeState() except Exception as e: # pylint: disable=broad-except logging.warning( "Plugin %s failed to initialize (%s), ignoring it.", plugin, e) continue # TODO(amoser): Those do not need to be inside the state, they # could be part of the plugin descriptor. plugin.state["logs"] = [] plugin.state["errors"] = [] output_plugins_states.append( rdf_flow_runner.OutputPluginState( plugin_state=plugin.state, plugin_descriptor=plugin_descriptor)) return output_plugins_states
def GetOutputPluginStates(output_plugins, source=None, token=None): """Initializes state for a list of output plugins.""" output_plugins_states = [] for plugin_descriptor in output_plugins: plugin_class = plugin_descriptor.GetPluginClass() try: _, plugin_state = plugin_class.CreatePluginAndDefaultState( source_urn=source, args=plugin_descriptor.plugin_args, token=token) except Exception as e: # pylint: disable=broad-except raise ValueError("Plugin %s failed to initialize (%s)" % (plugin_class, e)) # TODO(amoser): Those do not need to be inside the state, they # could be part of the plugin descriptor. plugin_state["logs"] = [] plugin_state["errors"] = [] output_plugins_states.append( rdf_flow_runner.OutputPluginState( plugin_state=plugin_state, plugin_descriptor=plugin_descriptor)) return output_plugins_states
def testWritingHuntOutputStatesForUnknownHuntRaises(self): state = rdf_flow_runner.OutputPluginState( plugin_descriptor=rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin1"), plugin_state={}) with self.assertRaises(db.UnknownHuntError): self.db.WriteHuntOutputPluginsStates(rdf_hunt_objects.RandomHuntId(), [state])
def testWritingAndReadingHuntOutputPluginsStatesWorks(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=compatibility.GetName(email_plugin.EmailOutputPlugin), plugin_args=email_plugin.EmailOutputPluginArgs(emails_limit=42)) state_1 = rdf_flow_runner.OutputPluginState( plugin_descriptor=plugin_descriptor, plugin_state={}) plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=compatibility.GetName(email_plugin.EmailOutputPlugin), plugin_args=email_plugin.EmailOutputPluginArgs(emails_limit=43)) state_2 = rdf_flow_runner.OutputPluginState( plugin_descriptor=plugin_descriptor, plugin_state={}) written_states = [state_1, state_2] self.db.WriteHuntOutputPluginsStates(hunt_obj.hunt_id, written_states) read_states = self.db.ReadHuntOutputPluginsStates(hunt_obj.hunt_id) self.assertEqual(read_states, written_states)
def InitializeContext(self, args): """Initializes the context of this flow.""" if args is None: args = rdf_flow_runner.FlowRunnerArgs() output_plugins_states = [] for plugin_descriptor in args.output_plugins: if not args.client_id: self.Log( "Not initializing output plugin %s as flow does not run on " "the client.", plugin_descriptor.plugin_name) continue output_base_urn = self.session_id.Add(OUTPUT_PLUGIN_BASE_SUFFIX) plugin_class = plugin_descriptor.GetPluginClass() plugin = plugin_class(self.flow_obj.output_urn, args=plugin_descriptor.plugin_args, output_base_urn=output_base_urn, token=self.token) try: plugin.InitializeState() # TODO(amoser): Those do not need to be inside the state, they # could be part of the plugin descriptor. plugin.state["logs"] = [] plugin.state["errors"] = [] output_plugins_states.append( rdf_flow_runner.OutputPluginState( plugin_state=plugin.state, plugin_descriptor=plugin_descriptor)) except Exception as e: # pylint: disable=broad-except logging.info( "Plugin %s failed to initialize (%s), ignoring it.", plugin, e) parent_creator = None if self.parent_runner: parent_creator = self.parent_runner.context.creator context = rdf_flow_runner.FlowContext( create_time=rdfvalue.RDFDatetime.Now(), creator=parent_creator or self.token.username, current_state="Start", output_plugins_states=output_plugins_states, state=rdf_flow_runner.FlowContext.State.RUNNING, ) return context
def testReadingHuntOutputPluginsReturnsThemInOrderOfWriting(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) states = [] for i in range(100): states.append( rdf_flow_runner.OutputPluginState( plugin_descriptor=rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin_%d" % i), plugin_state={})) random.shuffle(states) self.db.WriteHuntOutputPluginsStates(hunt_obj.hunt_id, states) read_states = self.db.ReadHuntOutputPluginsStates(hunt_obj.hunt_id) self.assertEqual(read_states, states)
def _HuntOutputPluginStateFromRow(self, row): """Builds OutputPluginState object from a DB row.""" plugin_name, plugin_args_bytes, plugin_state_bytes = row plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=plugin_name) if plugin_args_bytes is not None: plugin_args_cls = plugin_descriptor.GetPluginArgsClass() # If plugin_args_cls is None, we have no clue what class plugin args # should be and therefore no way to deserialize it. This can happen if # a plugin got renamed or removed, for example. In this case we # still want to get plugin's definition and state back and not fail hard, # so that all other plugins can be read. if plugin_args_cls is not None: plugin_descriptor.plugin_args = plugin_args_cls.FromSerializedBytes( plugin_args_bytes) plugin_state = rdf_protodict.AttributedDict.FromSerializedBytes( plugin_state_bytes) return rdf_flow_runner.OutputPluginState( plugin_descriptor=plugin_descriptor, plugin_state=plugin_state)