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