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)
