def Run(self): # Fix the time to avoid regressions. with test_lib.FakeTime(42): self.SetupClients(1) start_handler = flow_plugin.ApiStartRobotGetFilesOperationHandler() start_args = flow_plugin.ApiStartRobotGetFilesOperationArgs( hostname="Host", paths=["/test"]) start_result = start_handler.Handle(start_args, token=self.token) # Exploit the fact that 'get files' operation id is effectively a flow # URN. flow_urn = rdfvalue.RDFURN(start_result.operation_id) # Put something in the output collection flow_obj = aff4.FACTORY.Open(flow_urn, aff4_type=flow.GRRFlow, token=self.token) with aff4.FACTORY.Create( flow_obj.GetRunner().output_urn, aff4_type=sequential_collection.GeneralIndexedCollection, token=self.token) as collection: collection.Add(rdf_client.ClientSummary()) self.Check("GET", "/api/robot-actions/get-files/%s" % start_result.operation_id, replace={flow_urn.Basename(): "F:ABCDEF12"})
def testClientLookup(self): """When multiple clients match, check we run on the latest one.""" args = flow_plugin.ApiStartRobotGetFilesOperationArgs(hostname="Host", paths=["/test"]) result = self.handler.Handle(args, token=self.token) # Here we exploit the fact that operation_id is effectively a flow URN. self.assertIn("C.1000000000000003", result.operation_id)
def testThrottle(self): """Calling the same flow should raise.""" args = flow_plugin.ApiStartRobotGetFilesOperationArgs(hostname="Host", paths=["/test"]) self.handler.Handle(args, token=self.token) with self.assertRaises(throttle.ErrorFlowDuplicate): self.handler.Handle(args, token=self.token)
def Run(self): def ReplaceFlowId(): flows_dir_fd = aff4.FACTORY.Open(self.client_id.Add("flows"), token=self.token) flow_urn = list(flows_dir_fd.ListChildren())[0] return {flow_urn.Basename(): "W:ABCDEF"} with test_lib.FakeTime(42): self.Check("StartRobotGetFilesOperation", args=flow_plugin.ApiStartRobotGetFilesOperationArgs( hostname=self.client_id.Basename(), paths=["/tmp/test"]), replace=ReplaceFlowId)
def Run(self): # Fix the time to avoid regressions. with test_lib.FakeTime(42): self.SetupClients(1) start_handler = flow_plugin.ApiStartRobotGetFilesOperationHandler() start_args = flow_plugin.ApiStartRobotGetFilesOperationArgs( hostname="Host", paths=["/test"]) start_result = start_handler.Handle(start_args, token=self.token) # Exploit the fact that 'get files' operation id is effectively a flow # URN. flow_urn = rdfvalue.RDFURN(start_result.operation_id) # Put something in the output collection collection = flow.GRRFlow.ResultCollectionForFID(flow_urn, token=self.token) collection.Add(rdf_client.ClientSummary()) self.Check("GetRobotGetFilesOperationState", args=flow_plugin.ApiGetRobotGetFilesOperationStateArgs( operation_id=start_result.operation_id), replace={flow_urn.Basename(): "F:ABCDEF12"})