def testUnscheduleFlowRemovesScheduledFlow(self, db: abstract_db.Database):
context = _CreateContext(db)
client_id = db_test_utils.InitializeClient(db)
handler = flow_plugin.ApiScheduleFlowHandler()
sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs(
client_id=client_id,
flow=flow_plugin.ApiFlow(
name=file.CollectSingleFile.__name__,
args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))),
context=context)
sf2 = handler.Handle(flow_plugin.ApiCreateFlowArgs(
client_id=client_id,
flow=flow_plugin.ApiFlow(
name=file.CollectSingleFile.__name__,
args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))),
context=context)
handler = flow_plugin.ApiUnscheduleFlowHandler()
args = flow_plugin.ApiUnscheduleFlowArgs(
client_id=client_id, scheduled_flow_id=sf1.scheduled_flow_id)
handler.Handle(args, context=context)
handler = flow_plugin.ApiListScheduledFlowsHandler()
args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id,
creator=context.username)
results = handler.Handle(args, context=context)
self.assertEqual(results.scheduled_flows, [sf2])
def testFlowWithoutFlowProgressTypeDoesNotReportProgress(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=flow_test_lib.DummyFlow)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertIsNone(flow_api_obj.progress)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(
flow_obj, with_progress=True)
self.assertIsNone(flow_api_obj.progress)
def testWithFlowProgressTypeReportsProgressCorrectly(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=flow_test_lib.DummyFlowWithProgress)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertIsNotNone(flow_api_obj.progress)
# An empty proto is created by default.
self.assertFalse(flow_api_obj.progress.HasField("status"))
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(
flow_obj, with_progress=True)
self.assertIsNotNone(flow_api_obj.progress)
self.assertEqual(flow_api_obj.progress.status, "Progress.")
def testFileFinderFlowNameCanBeOverriden(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(
enabled=True,
file_finder_flow_name=AnotherFileFinder.__name__))
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=file_finder.FileFinder.__name__),
client_id=self.client_id),
token=self.token)
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=AnotherFileFinder.__name__),
client_id=self.client_id),
token=self.token)
def Handle(self, args, token=None):
if not args.client_id:
raise RuntimeError("Client id has to be specified.")
if not args.flow.name:
raise RuntimeError("Flow name is not specified.")
if data_store.RelationalDBEnabled():
delegate = api_flow.ApiCreateFlowHandler()
# Note that runner_args are dropped. From all the arguments We use only
# the flow name and the arguments.
delegate_args = api_flow.ApiCreateFlowArgs(
client_id=args.client_id)
delegate_args.flow.name = self.override_flow_name or args.flow.name
delegate_args.flow.args = self.override_flow_args or args.flow.args
return delegate.Handle(delegate_args, token=token)
else:
flow_id = flow.StartAFF4Flow(
client_id=args.client_id.ToClientURN(),
flow_name=self.override_flow_name or args.flow.name,
token=token,
args=self.override_flow_args or args.flow.args)
with aff4.FACTORY.Open(flow_id,
aff4_type=flow.GRRFlow,
mode="rw",
token=token) as fd:
fd.AddLabel(LABEL_NAME_PREFIX + self.robot_id)
return api_flow.ApiFlow().InitFromAff4Object(
fd, flow_id=flow_id.Basename())
def Check(path):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=file_finder.FileFinder.__name__,
args=rdf_file_finder.FileFinderArgs(paths=[path])),
client_id=self.client_id),
token=self.token)
def Run(self):
client_urn = self.SetupClient(0)
client_id = client_urn.Basename()
def ReplaceFlowId():
if data_store.RelationalDBFlowsEnabled():
flows = data_store.REL_DB.ReadAllFlowObjects(
client_id=client_id)
self.assertNotEmpty(flows)
flow_id = flows[0].flow_id
else:
flows_dir_fd = aff4.FACTORY.Open(client_urn.Add("flows"),
token=self.token)
flow_id = list(flows_dir_fd.ListChildren())[0].Basename()
return api_regression_test_lib.GetFlowTestReplaceDict(
client_id, flow_id)
with test_lib.FakeTime(42):
self.Check("CreateFlow",
args=flow_plugin.ApiCreateFlowArgs(
client_id=client_id,
flow=flow_plugin.ApiFlow(
name=processes.ListProcesses.__name__,
args=processes.ListProcessesArgs(
filename_regex=".", fetch_binaries=True),
runner_args=rdf_flow_runner.FlowRunnerArgs(
output_plugins=[], notify_to_user=True))),
replace=ReplaceFlowId)
def InitFromDatabaseObject(self, db_obj, approval_subject_obj=None):
_InitApiApprovalFromDatabaseObject(self, db_obj)
if not approval_subject_obj:
approval_subject_obj = data_store.REL_DB.ReadHuntObject(
db_obj.subject_id)
approval_subject_counters = data_store.REL_DB.ReadHuntCounters(
db_obj.subject_id)
self.subject = api_hunt.ApiHunt().InitFromHuntObject(
approval_subject_obj,
hunt_counters=approval_subject_counters,
with_full_summary=True)
original_object = approval_subject_obj.original_object
if original_object.object_type == "FLOW_REFERENCE":
original_flow = data_store.REL_DB.ReadFlowObject(
original_object.flow_reference.client_id,
original_object.flow_reference.flow_id)
self.copied_from_flow = api_flow.ApiFlow().InitFromFlowObject(
original_flow)
elif original_object.object_type == "HUNT_REFERENCE":
original_hunt = data_store.REL_DB.ReadHuntObject(
original_object.hunt_reference.hunt_id)
original_hunt_counters = data_store.REL_DB.ReadHuntCounters(
original_object.hunt_reference.hunt_id)
self.copied_from_hunt = api_hunt.ApiHunt().InitFromHuntObject(
original_hunt,
hunt_counters=original_hunt_counters,
with_full_summary=True)
return self
def Check(path):
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=file_finder.FileFinder.__name__,
args=rdf_file_finder.FileFinderArgs(paths=[path])),
client_id=self.client_id),
token=self.token)
def Check(artifacts):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__,
args=rdf_artifacts.ArtifactCollectorFlowArgs(
artifact_list=artifacts)),
client_id=self.client_id),
token=self.token)
def testInitializesClientIdForClientBasedFlows(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=processes.ListProcesses)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertEqual(flow_api_obj.client_id,
client_plugin.ApiClientId(client_id))
def testFlowWithoutResultsCorrectlyReportsEmptyResultMetadata(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=flow_test_lib.DummyFlow)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertIsNotNone(flow_api_obj.result_metadata)
self.assertEmpty(flow_api_obj.result_metadata.num_results_per_type_tag)
def Check(artifacts):
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__,
args=rdf_artifacts.ArtifactCollectorFlowArgs(
artifact_list=artifacts)),
client_id=self.client_id),
token=self.token)
def testArtifactCollectorFlowNameCanBeOverriden(self):
router = self._CreateRouter(
artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams(
enabled=True,
artifact_collector_flow_name=AnotherArtifactCollector.__name__)
)
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__),
client_id=self.client_id),
token=self.token)
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=AnotherArtifactCollector.__name__),
client_id=self.client_id),
token=self.token)
def _CreateFlowWithRobotId(self, flow_name=None, flow_args=None):
flow_name = flow_name or file_finder.FileFinder.__name__
handler = rr.ApiRobotCreateFlowHandler(robot_id=self.robot_id)
flow_result = handler.Handle(api_flow.ApiCreateFlowArgs(
client_id=self.client_id,
flow=api_flow.ApiFlow(name=flow_name, args=flow_args)),
token=self.token)
return flow_result.flow_id
def testFlowWithoutFlowProgressTypeReportsDefaultFlowProgress(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=flow_test_lib.DummyFlow)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertIsNotNone(flow_api_obj.progress)
self.assertIsInstance(flow_api_obj.progress,
rdf_flow_objects.DefaultFlowProgress)
def testLeavesClientIdEmptyForNonClientBasedFlows(self):
client_id = self.SetupClient(0)
flow_urn = flow.StartFlow(client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=self.token)
flow_obj = aff4.FACTORY.Open(flow_urn, token=self.token)
flow_api_obj = flow_plugin.ApiFlow().InitFromAff4Object(
flow_obj, flow_id=flow_urn.Basename())
self.assertEquals(flow_api_obj.client_id,
client_plugin.ApiClientId(client_id))
def testRunnerArgsBaseSessionIdDoesNotAffectCreatedFlow(self):
"""When multiple clients match, check we run on the latest one."""
flow_runner_args = rdf_flow_runner.FlowRunnerArgs(
base_session_id="aff4:/foo")
args = flow_plugin.ApiCreateFlowArgs(
client_id=self.client_id.Basename(),
flow=flow_plugin.ApiFlow(name=processes.ListProcesses.__name__,
runner_args=flow_runner_args))
result = self.handler.Handle(args, token=self.token)
self.assertNotStartsWith(str(result.urn), "aff4:/foo")
def Handle(self, args, token=None):
# Note: this is a legacy AFF4 implementation.
flow_urn = args.run_id.ToURN(args.cron_job_id)
flow_obj = aff4.FACTORY.Open(flow_urn,
aff4_type=flow.GRRFlow,
mode="r",
token=token)
f = api_plugins_flow.ApiFlow().InitFromAff4Object(
flow_obj, with_state_and_context=True)
return ApiCronJobRun().InitFromApiFlow(f)
def testOnlyFileFinderAndArtifactCollectorFlowsAreAllowed(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(enabled=True),
artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams(
enabled=True))
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=flow_test_lib.BrokenFlow.__name__),
client_id=self.client_id),
token=self.token)
def testOverriddenFileFinderFlowCanBeCreatedUsingOriginalFileFinderName(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(
enabled=True, file_finder_flow_name=AnotherFileFinder.__name__)) # pylint: disable=undefined-variable
handler = router.CreateFlow(
api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=file_finder.FileFinder.__name__),
client_id=self.client_id),
token=self.token)
self.assertEqual(handler.override_flow_name, AnotherFileFinder.__name__) # pylint: disable=undefined-variable
def testFileFinderFlowNameCanBeOverridden(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(
enabled=True,
file_finder_flow_name=AnotherFileFinder.__name__)) # pylint: disable=undefined-variable
handler = router.CreateFlow(
api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=AnotherFileFinder.__name__), # pylint: disable=undefined-variable
client_id=self.client_id),
context=self.context)
self.assertEqual(handler.override_flow_name,
AnotherFileFinder.__name__) # pylint: disable=undefined-variable
def testArtifactCollectorFlowNameCanBeOverridden(self):
router = self._CreateRouter(
artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams(
enabled=True,
artifact_collector_flow_name=AnotherArtifactCollector.__name__)) # pylint: disable=undefined-variable
handler = router.CreateFlow(
api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=AnotherArtifactCollector.__name__), # pylint: disable=undefined-variable
client_id=self.client_id),
token=self.token)
self.assertEqual(handler.override_flow_name,
AnotherArtifactCollector.__name__) # pylint: disable=undefined-variable
def testInitializesClientIdForClientBasedFlows(self):
client_id = self.SetupClient(0)
flow_urn = flow.StartAFF4Flow(
# Override base session id, so that the flow URN looks
# like: aff4:/F:112233
base_session_id="aff4:/",
client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=self.token)
flow_obj = aff4.FACTORY.Open(flow_urn, token=self.token)
flow_api_obj = flow_plugin.ApiFlow().InitFromAff4Object(
flow_obj, flow_id=flow_urn.Basename())
self.assertIsNone(flow_api_obj.client_id)
def testOverriddenArtifactCollectorFlowCanBeCreatedUsingOriginalName(self):
router = self._CreateRouter(
artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams(
enabled=True,
artifact_collector_flow_name=AnotherArtifactCollector.__name__)
) # pylint: disable=undefined-variable
handler = router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__),
client_id=self.client_id),
context=self.context)
self.assertEqual(handler.override_flow_name,
AnotherArtifactCollector.__name__) # pylint: disable=undefined-variable
def Handle(self, args, token=None):
if data_store.RelationalDBEnabled():
run = cronjobs.CronManager().ReadJobRun(
str(args.cron_job_id), str(args.run_id))
if not run:
raise CronJobRunNotFoundError(
"Cron job run with id %s could not be found" % args.run_id)
return ApiCronJobRun().InitFromRunObject(run)
else:
# Note: this is a legacy AFF4 implementation.
flow_urn = args.run_id.ToURN(args.cron_job_id)
flow_obj = aff4.FACTORY.Open(
flow_urn, aff4_type=flow.GRRFlow, mode="r", token=token)
f = api_plugins_flow.ApiFlow().InitFromAff4Object(
flow_obj, with_state_and_context=True)
return ApiCronJobRun().InitFromApiFlow(f, cron_job_id=args.cron_job_id)
def testScheduleFlow(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
handler = flow_plugin.ApiScheduleFlowHandler()
args = flow_plugin.ApiCreateFlowArgs(
client_id=client_id,
flow=flow_plugin.ApiFlow(
name=file.CollectSingleFile.__name__,
args=rdf_file_finder.CollectSingleFileArgs(path="/foo"),
runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60)))
sf = handler.Handle(args, token=token)
self.assertEqual(sf.client_id, client_id)
self.assertEqual(sf.creator, token.username)
self.assertNotEmpty(sf.scheduled_flow_id)
self.assertEqual(sf.flow_name, file.CollectSingleFile.__name__)
self.assertEqual(sf.flow_args.path, "/foo")
self.assertEqual(sf.runner_args.cpu_limit, 60)
def Run(self):
client_id = self.SetupClient(0)
def ReplaceFlowId():
flows_dir_fd = aff4.FACTORY.Open(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("CreateFlow",
args=flow_plugin.ApiCreateFlowArgs(
client_id=client_id.Basename(),
flow=flow_plugin.ApiFlow(
name=processes.ListProcesses.__name__,
args=processes.ListProcessesArgs(
filename_regex=".", fetch_binaries=True),
runner_args=rdf_flow_runner.FlowRunnerArgs(
output_plugins=[], notify_to_user=False))),
replace=ReplaceFlowId)
def Run(self):
client_id = self.SetupClient(0)
def ReplaceFlowId():
flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id)
self.assertNotEmpty(flows)
flow_id = flows[0].flow_id
return api_regression_test_lib.GetFlowTestReplaceDict(client_id, flow_id)
with test_lib.FakeTime(42):
self.Check(
"CreateFlow",
args=flow_plugin.ApiCreateFlowArgs(
client_id=client_id,
flow=flow_plugin.ApiFlow(
name=processes.ListProcesses.__name__,
args=processes.ListProcessesArgs(
filename_regex=".", fetch_binaries=True),
runner_args=rdf_flow_runner.FlowRunnerArgs(
output_plugins=[], notify_to_user=True))),
replace=ReplaceFlowId)
def testFileFinderHashMaxFileSizeCanBeOverriden(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(
enabled=True, max_file_size=42))
ha = rdf_file_finder.FileFinderHashActionOptions()
ha.max_size = 80
ha.oversized_file_policy = ha.OversizedFilePolicy.HASH_TRUNCATED
path = "/foo/bar"
handler = router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(name=file_finder.FileFinder.__name__,
args=rdf_file_finder.FileFinderArgs(
paths=[path],
action=rdf_file_finder.FileFinderAction(
action_type="HASH", hash=ha))),
client_id=self.client_id),
token=self.token)
ha = handler.override_flow_args.action.hash
self.assertEqual(ha.oversized_file_policy, ha.OversizedFilePolicy.SKIP)
self.assertEqual(ha.max_size, 42)
