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 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 Check(artifacts):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__,
args=artifact_utils.ArtifactCollectorFlowArgs(
artifact_list=artifacts)),
client_id=self.client_id),
token=self.token)
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):
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(api_flow.ApiCreateFlowArgs(
flow=api_flow.ApiFlow(
name=collectors.ArtifactCollectorFlow.__name__,
args=artifact_utils.ArtifactCollectorFlowArgs(
artifact_list=artifacts)),
client_id=self.client_id),
token=self.token)
def testPassesFlowArgsThroughIfNoOverridesSpecified(self):
h = rr.ApiRobotCreateFlowHandler(robot_id="foo")
args = api_flow.ApiCreateFlowArgs(client_id=self.client_id.Basename())
args.flow.name = file_finder.FileFinder.__name__
args.flow.args = rdf_file_finder.FileFinderArgs(paths=["foo"])
f = h.Handle(args=args, token=self.token)
self.assertEqual(f.args.paths, ["foo"])
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 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 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=test_lib.BrokenFlow.__name__),
client_id=self.client_id),
token=self.token)
def testRunnerArgsBaseSessionIdDoesNotAffectCreatedFlow(self):
"""When multiple clients match, check we run on the latest one."""
flow_runner_args = rdf_flows.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.assertFalse(utils.SmartStr(result.urn).startswith("aff4:/foo"))
def testAllClientFlowsMethodsAreAccessChecked(self):
args = api_flow.ApiListFlowsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlows, "CheckClientAccess", args=args)
args = api_flow.ApiGetFlowArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.GetFlow, "CheckClientAccess", args=args)
args = api_flow.ApiCreateFlowArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.CreateFlow, "CheckClientAccess", args=args)
self.CheckMethodIsAccessChecked(
self.router.CreateFlow, "CheckIfCanStartFlow", args=args)
args = api_flow.ApiCancelFlowArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.CancelFlow, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowRequestsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowRequests, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowResultsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowResults, "CheckClientAccess", args=args)
args = api_flow.ApiGetFlowResultsExportCommandArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.GetFlowResultsExportCommand, "CheckClientAccess", args=args)
args = api_flow.ApiGetFlowFilesArchiveArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.GetFlowFilesArchive, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowOutputPluginsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowOutputPlugins, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowOutputPluginLogsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowOutputPluginLogs, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowOutputPluginErrorsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowOutputPluginErrors, "CheckClientAccess", args=args)
args = api_flow.ApiListFlowLogsArgs(client_id=self.client_id)
self.CheckMethodIsAccessChecked(
self.router.ListFlowLogs, "CheckClientAccess", args=args)
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("CreateFlow",
args=flow_plugin.ApiCreateFlowArgs(
client_id=self.client_id.Basename(),
flow=flow_plugin.ApiFlow(
name=processes.ListProcesses.__name__,
args=processes.ListProcessesArgs(
filename_regex=".", fetch_binaries=True),
runner_args=rdf_flows.FlowRunnerArgs(
output_plugins=[],
priority="HIGH_PRIORITY",
notify_to_user=False))),
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)
def testFileFinderDownloadMaxFileSizeCanBeOverriden(self):
router = self._CreateRouter(
file_finder_flow=rr.RobotRouterFileFinderFlowParams(
enabled=True, max_file_size=42))
da = rdf_file_finder.FileFinderDownloadActionOptions()
da.max_size = 80
da.oversized_file_policy = da.OversizedFilePolicy.DOWNLOAD_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="DOWNLOAD",
download=da))),
client_id=self.client_id),
token=self.token)
da = handler.override_flow_args.action.download
self.assertEqual(da.oversized_file_policy, da.OversizedFilePolicy.SKIP)
self.assertEqual(da.max_size, 42)
def testAllGlobalFlowsMethodsAreAccessChecked(self):
args = api_flow.ApiCreateFlowArgs(flow=api_flow.ApiFlow(
name="ListProcesses"))
self.CheckMethodIsAccessChecked(self.router.CreateGlobalFlow,
"CheckIfCanStartFlow",
args=args)
def testCreateFlowIsDisabledByDefault(self):
router = self._CreateRouter()
with self.assertRaises(access_control.UnauthorizedAccess):
router.CreateFlow(
api_flow.ApiCreateFlowArgs(client_id=self.client_id),
token=self.token)
def testCreateFlowRaisesIfClientIdNotSpecified(self):
router = self._CreateRouter()
with self.assertRaises(ValueError):
router.CreateFlow(api_flow.ApiCreateFlowArgs(), token=self.token)