class AgentWorker(common.AbstractControllerCommand):
name = "worker"
__args = [
dict(name="batches", type="ChoiceArray", positional=True,
choices=utils.JITIteratorCallable(
utils.get_all_subclasses, BatchTicket),
help="One or more batch jobs to run."),
dict(name="loop", type="IntParser", default=0,
help="If specified we loop forever and sleep that "
"many seconds between loops."),
]
table_header = [
dict(name="Message")
]
def collect(self):
batches = [BatchRunner(session=self.session, batch_name=x)
for x in (self.plugin_args.batches or
utils.get_all_subclasses(BatchTicket))]
while 1:
for batch_runner in batches:
yield dict(Message="Running batch %s" % batch_runner.batch_name)
try:
batch_runner.run()
except Exception:
self.session.logging.exception("Error running batch")
if not self.plugin_args.loop:
break
time.sleep(self.plugin_args.loop)
class AgentControllerExplainFlows(common.AbstractControllerCommand):
"""Explain all the parameters a flow may take."""
name = "explain_flow"
__args = [
dict(name="flow",
type="Choices",
positional=True,
required=True,
choices=utils.JITIteratorCallable(utils.get_all_subclasses,
flow.Flow),
help="The flow to explain"),
dict(name="verbosity",
type="IntParser",
default=0,
help="If non zero show all fields."),
dict(name="recursive",
type="Bool",
help="Show recursively nested fields."),
]
table_header = [
dict(name="Field", type="TreeNode", max_depth=5),
dict(name="Type"),
dict(name="Description")
]
table_options = dict(auto_widths=True, )
def _explain(self, flow_cls, depth=0):
for descriptor in flow_cls.get_descriptors():
user_accessible = descriptor.get("user")
if self.plugin_args.verbosity < 1 and not user_accessible:
continue
field = descriptor["name"]
field_type = descriptor.get("type", "string")
field_description = field_type
if isinstance(field_type, type):
field_description = "(%s)" % field_type.__name__
yield dict(Field=field,
Type=field_description,
Description=descriptor.get("doc", ""),
depth=depth)
if (self.plugin_args.recursive and isinstance(field_type, type)
and issubclass(field_type, serializer.SerializedObject)):
for row in self._explain(field_type, depth=depth + 1):
#row["Field"] = "%s.%s" % (field, row["Field"])
yield row
def collect(self):
flow_cls = flow.Flow.ImplementationByClass(self.plugin_args.flow)
for x in self._explain(flow_cls):
yield x
class AgentControllerRunFlow(SerializedObjectInspectorMixin,
common.AbstractControllerCommand):
name = "launch_flow"
__args = [
dict(name="flow", type="Choices", positional=True, required=True,
choices=utils.JITIteratorCallable(
utils.get_all_subclasses, flow.Flow),
help="The flow to launch"),
dict(name="args", type="Any", positional=True, default={},
help="Arguments to the flow (use explain_flow to see valid args)."
"This may also be a JSON encoded string"),
dict(name="queue",
help="Which queue to schedule the hunt on."),
dict(name="condition",
help="An EFilter query to evaluate if the flow should be run."),
# This should only be set if no condition is specified.
dict(name="canned_condition", type="Choices", default=None,
choices=CANNED_CONDITIONS,
help="Canned conditions for the hunt."),
dict(name="live", type="Choices", default="API",
choices=["API", "Memory"],
help="Live mode to use"),
dict(name="quota", type="IntParser", default=3600,
help="Total number of CPU seconds allowed for this flow."),
]
def make_flow_object(self):
flow_cls = flow.Flow.ImplementationByClass(self.plugin_args.flow)
if not flow_cls:
raise plugin.PluginError("Unknown flow %s" % self.plugin_args.flow)
args = self.plugin_args.args
if isinstance(args, basestring):
try:
args = json.loads(args)
except Exception as e:
raise plugin.PluginError(
"args should be a JSON string of a dict: %s" % e)
if not isinstance(args, dict):
raise plugin.PluginError("args should be a dict")
flow_obj = flow_cls.from_primitive(args, session=self.session)
flow_obj.client_id = self.client_id
flow_obj.queue = self.plugin_args.queue
flow_obj.session.live = self.plugin_args.live
# If a canned condition was specified automatically add it.
if self.plugin_args.canned_condition:
flow_obj.condition = CANNED_CONDITIONS[
self.plugin_args.canned_condition]
elif self.plugin_args.condition:
flow_obj.condition = self.plugin_args.condition
# Specify flow quota.
flow_obj.quota.user_time = self.plugin_args.quota
return flow_obj
def collect(self):
# Now launch the flow.
flow_obj = self.make_flow_object()
flow_obj.start()
for x in self._explain(flow_obj):
yield x
