Пример #1
0
    def testListScheduledFlows(self, db: abstract_db.Database):
        context = _CreateContext(db)
        client_id1 = db_test_utils.InitializeClient(db)
        client_id2 = db_test_utils.InitializeClient(db)

        handler = flow_plugin.ApiScheduleFlowHandler()
        sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs(
            client_id=client_id1,
            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_id1,
            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.Handle(flow_plugin.ApiCreateFlowArgs(
            client_id=client_id2,
            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.ApiListScheduledFlowsHandler()
        args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id1,
                                                     creator=context.username)
        results = handler.Handle(args, context=context)

        self.assertEqual(results.scheduled_flows, [sf1, sf2])
Пример #2
0
    def testUnscheduleFlowRemovesScheduledFlow(self, db: abstract_db.Database):
        token = _CreateToken(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))),
                             token=token)
        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))),
                             token=token)

        handler = flow_plugin.ApiUnscheduleFlowHandler()
        args = flow_plugin.ApiUnscheduleFlowArgs(
            client_id=client_id, scheduled_flow_id=sf1.scheduled_flow_id)
        handler.Handle(args, token=token)

        handler = flow_plugin.ApiListScheduledFlowsHandler()
        args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id,
                                                     creator=token.username)
        results = handler.Handle(args, token=token)

        self.assertEqual(results.scheduled_flows, [sf2])
Пример #3
0
    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)
Пример #4
0
 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)
Пример #5
0
    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)
Пример #6
0
    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())
Пример #7
0
 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)
Пример #8
0
 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)
Пример #9
0
 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)
Пример #10
0
    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)
Пример #11
0
    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
Пример #12
0
    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"])
Пример #13
0
    def testOverridesFlowArgsThroughIfOverridesSpecified(self):
        override_flow_args = rdf_file_finder.FileFinderArgs(paths=["bar"])
        h = rr.ApiRobotCreateFlowHandler(override_flow_args=override_flow_args)

        args = api_flow.ApiCreateFlowArgs(client_id=self.client_id)
        args.flow.name = file_finder.FileFinder.__name__
        args.flow.args = rdf_file_finder.FileFinderArgs(paths=["foo"])

        f = h.Handle(args=args, context=self.context)
        self.assertEqual(f.args.paths, ["bar"])
Пример #14
0
    def testOverridesFlowNameIfOverrideArgIsSpecified(self):
        h = rr.ApiRobotCreateFlowHandler(
            override_flow_name=AnotherFileFinder.__name__)  # pylint: disable=undefined-variable

        args = api_flow.ApiCreateFlowArgs(client_id=self.client_id)
        args.flow.name = file_finder.FileFinder.__name__
        args.flow.args = rdf_file_finder.FileFinderArgs(paths=["foo"])

        f = h.Handle(args=args, context=self.context)
        self.assertEqual(f.name, AnotherFileFinder.__name__)  # pylint: disable=undefined-variable
Пример #15
0
  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, "CheckIfCanStartClientFlow", 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.ApiGetExportedFlowResultsArgs(client_id=self.client_id)
    self.CheckMethodIsAccessChecked(
        self.router.GetExportedFlowResults, "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)
Пример #16
0
    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")
Пример #17
0
    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)
Пример #18
0
  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
Пример #19
0
  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
Пример #20
0
    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
Пример #21
0
    def Handle(self, args, context=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.")

        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, context=context)
Пример #22
0
    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
Пример #23
0
    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)
Пример #24
0
    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)
Пример #26
0
    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)
Пример #27
0
    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)
Пример #28
0
 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)
Пример #29
0
 def testCreateFlowRaisesIfClientIdNotSpecified(self):
     router = self._CreateRouter()
     with self.assertRaises(ValueError):
         router.CreateFlow(api_flow.ApiCreateFlowArgs(), token=self.token)