def GenerateToken(self, username, reason): return access_control.ACLToken(username=username, reason=reason)
def CheckCronJobAccess(self, username, cron_job_id): token = access_control.ACLToken(username=username) self.legacy_manager.CheckCronJobAccess(token, cron_job_id.ToURN())
def TestHuntHelperWithMultipleMocks(client_mocks, token=None, iteration_limit=None, worker=None): """Runs a hunt with a given set of clients mocks. Args: client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock objects are used to handle client actions. Methods names of a client mock object correspond to client actions names. For an example of a client mock object, see SampleHuntMock. token: An instance of access_control.ACLToken security token. iteration_limit: If None, hunt will run until it's finished. Otherwise, worker_mock.Next() will be called iteration_limit number of times. Every iteration processes worker's message queue. If new messages are sent to the queue during the iteration processing, they will be processed on next iteration. worker: flow_test_lib.TestWorker object to use. Returns: A number of iterations complete. """ if token is None: token = access_control.ACLToken(username="******") # Worker always runs with absolute privileges, therefore making the token # SetUID(). token = token.SetUID() client_mocks = [ flow_test_lib.MockClient(client_id, client_mock, token=token) for client_id, client_mock in client_mocks.items() ] if worker is None: rel_db_worker = flow_test_lib.TestWorker() data_store.REL_DB.RegisterFlowProcessingHandler( rel_db_worker.ProcessFlow) else: rel_db_worker = worker num_iterations = 0 try: # Run the clients and worker until nothing changes any more. while iteration_limit is None or num_iterations < iteration_limit: data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess(timeout=10) worker_processed = rel_db_worker.ResetProcessedFlows() client_processed = 0 for client_mock in client_mocks: client_processed += int(client_mock.Next()) num_iterations += 1 if client_processed == 0 and not worker_processed: break finally: if worker is None: data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60) rel_db_worker.Shutdown() return num_iterations
def testTokenWithUsernameAndReasonIsValid(self): token = access_control.ACLToken(username="******", reason="For testing") user_managers.ValidateToken(token, "aff4:/C.0000000000000001")
def CheckClientAccess(self, username, client_id): token = access_control.ACLToken(username=username) self.legacy_manager.CheckClientAccess(token, client_id.ToClientURN())
def __init__(self, run_state, job): self.run_state = run_state self.job = job self.token = access_control.ACLToken(username="******")
def testRaisesIfTraitsSetInRequest(self): user = user_plugin.ApiGrrUser( interface_traits=user_plugin.ApiGrrUserInterfaceTraits()) with self.assertRaises(ValueError): self.handler.Handle(user, token=access_control.ACLToken(username=u"foo"))
def __init__(self, run_state, job): self.run_state = run_state self.job = job self.token = access_control.ACLToken(username=CRON_JOB_USERNAME)
def main(argv): del argv # Unused. if flags.FLAGS.version: print("GRR FUSE {}".format(config_server.VERSION["packageversion"])) return config.CONFIG.AddContext(contexts.COMMAND_LINE_CONTEXT, "Context applied for all command line tools") server_startup.Init() if fuse is None: logging.fatal("""Could not start! fusepy must be installed to run fuse_mount.py! Try to run: pip install fusepy inside your virtualenv. """) sys.exit(1) if not flags.FLAGS.mountpoint: Usage() sys.exit(1) # We multiple inherit from GRRFuse and fuse.Operations. In the # case that fuse is present, we run the actual FUSE layer, since we have # fuse.Operations. In the case that fuse is not present, we have already # exited by now if we were run from the command line, and if we were not run # from the command line, we've been imported, and we run the tests using a # mock fuse. class FuseOperation(GRRFuse, fuse.Operations): pass root = flags.FLAGS.aff4path username = flags.FLAGS.username or getpass.getuser() data_store.default_token = access_control.ACLToken( username=username, reason=flags.FLAGS.reason or "fusemount") logging.info("fuse_mount.py is mounting %s at %s....", root, flags.FLAGS.mountpoint) refresh_policy = flags.FLAGS.refresh_policy if refresh_policy == "always": max_age_before_refresh = datetime.timedelta(0) elif refresh_policy == "never": # Set the max age to be the maximum possible time difference. max_age_before_refresh = datetime.timedelta.max elif refresh_policy == "if_older_than_max_age": max_age_before_refresh = datetime.timedelta( seconds=flags.FLAGS.max_age_before_refresh) else: # Otherwise, a flag outside the enum was given and the flag validator threw # an execption. pass fuse_operation = FuseOperation( root=root, token=data_store.default_token, max_age_before_refresh=max_age_before_refresh, ignore_cache=flags.FLAGS.ignore_cache, force_sparse_image=flags.FLAGS.force_sparse_image, sparse_image_threshold=flags.FLAGS.sparse_image_threshold, timeout=flags.FLAGS.timeout) fuse.FUSE( fuse_operation, flags.FLAGS.mountpoint, foreground=not flags.FLAGS.background)
def testUsesCollectionTimeFiles(self, db: abstract_db.Database): context = _CreateContext(db) client_id = db_test_utils.InitializeClient(db) snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.os = "redox" db.WriteClientSnapshot(snapshot) with temp.AutoTempFilePath() as temp_filepath: fake_artifact_source = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.FILE, attributes={ "paths": [temp_filepath], }) fake_artifact = rdf_artifacts.Artifact( name="FakeArtifact", doc="Lorem ipsum.", sources=[fake_artifact_source]) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [fake_artifact.name] flow_args.apply_parsers = False with io.open(temp_filepath, mode="wb") as temp_filedesc: temp_filedesc.write(b"OLD") with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(fake_artifact) # First, we run the artifact collector to collect the old file and save # the flow id to parse the results later. flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, action_mocks.FileFinderClientMock(), client_id=client_id, args=flow_args, token=access_control.ACLToken(username=context.username)) flow_test_lib.FinishAllFlowsOnClient(client_id) with io.open(temp_filepath, mode="wb") as temp_filedesc: temp_filedesc.write(b"NEW") with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(fake_artifact) # Now, we run the artifact collector again to collect the new file to # update to this version on the server. The parsing should be performed # against the previous flow. flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, action_mocks.FileFinderClientMock(), client_id=client_id, args=flow_args, token=access_control.ACLToken(username=context.username)) flow_test_lib.FinishAllFlowsOnClient(client_id) class FakeFileParser( abstract_parser.SingleFileParser[rdfvalue.RDFBytes]): supported_artifacts = [fake_artifact.name] def ParseFile( self, knowledge_base: rdf_client.KnowledgeBase, pathspec: rdf_paths.PathSpec, filedesc: file_store.BlobStream, ) -> Iterable[rdfvalue.RDFBytes]: del knowledge_base, pathspec # Unused. return [rdfvalue.RDFBytes(filedesc.Read())] with parser_test_lib._ParserContext("FakeFile", FakeFileParser): args = flow_plugin.ApiListParsedFlowResultsArgs( client_id=client_id, flow_id=flow_id, offset=0, count=1024) result = self.handler.Handle(args, context=context) self.assertEmpty(result.errors) self.assertLen(result.items, 1) response = result.items[0].payload self.assertEqual(response, b"OLD")
def TerminateFlow(cls, flow_id, reason=None, status=None, token=None, force=False): """Terminate a flow. Args: flow_id: The flow session_id to terminate. reason: A reason to log. status: Status code used in the generated status message. token: The access token to be used for this request. force: If True then terminate locked flows hard. Raises: FlowError: If the flow can not be found. """ if not force: flow_obj = aff4.FACTORY.OpenWithLock(flow_id, aff4_type=GRRFlow, blocking=True, token=token) else: flow_obj = aff4.FACTORY.Open(flow_id, aff4_type=GRRFlow, mode="rw", token=token) if not flow_obj: raise FlowError("Could not terminate flow %s" % flow_id) with flow_obj: runner = flow_obj.GetRunner() if not runner.IsRunning(): return if token is None: token = access_control.ACLToken() if reason is None: reason = "Manual termination by console." # This calls runner.Terminate to kill the flow runner.Error(reason, status_code=status) flow_obj.Log("Terminated by user {0}. Reason: {1}".format( token.username, reason)) # From now on we run with supervisor access super_token = token.SetUID() # Also terminate its children children_to_kill = aff4.FACTORY.MultiOpen(flow_obj.ListChildren(), token=super_token, aff4_type=GRRFlow) for child_obj in children_to_kill: cls.TerminateFlow(child_obj.urn, reason="Parent flow terminated.", token=super_token, force=force)
def testUsesKnowledgebaseFromFlow(self, db: abstract_db.Database): context = _CreateContext(db) client_id = db_test_utils.InitializeClient(db) # This is the snapshot that is visible to the flow and should be used for # parsing results. snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.os = "redox" db.WriteClientSnapshot(snapshot) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(self.ECHO1337_ARTIFACT) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name] flow_args.apply_parsers = False flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, token=access_control.ACLToken(username=context.username)) class FakeParser( abstract_parser.SingleResponseParser[ rdf_client_action.ExecuteResponse], ): supported_artifacts = [self.ECHO1337_ARTIFACT.name] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdf_client_action.ExecuteResponse, ) -> Iterable[rdf_client_action.ExecuteResponse]: precondition.AssertType(response, rdf_client_action.ExecuteResponse) parsed_response = rdf_client_action.ExecuteResponse() parsed_response.stdout = response.stdout parsed_response.stderr = knowledge_base.os.encode("utf-8") return [parsed_response] # This is a snapshot written to the database after the responses were # collected, so this should not be used for parsing. snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.os = "linux" db.WriteClientSnapshot(snapshot) with parser_test_lib._ParserContext("Fake", FakeParser): args = flow_plugin.ApiListParsedFlowResultsArgs( client_id=client_id, flow_id=flow_id, offset=0, count=1024) result = self.handler.Handle(args, context=context) self.assertEmpty(result.errors) self.assertLen(result.items, 1) response = result.items[0].payload self.assertIsInstance(response, rdf_client_action.ExecuteResponse) self.assertEqual(response.stdout, b"1337") self.assertEqual(response.stderr.decode("utf-8"), "redox")
def StartFlow(client_id=None, cpu_limit=7200, creator=None, flow_args=None, flow_cls=None, network_bytes_limit=None, original_flow=None, output_plugins=None, parent_flow_obj=None, **kwargs): """The main factory function for creating and executing a new flow. Args: client_id: ID of the client this flow should run on. cpu_limit: CPU limit in seconds for this flow. creator: Username that requested this flow. flow_args: An arg protocol buffer which is an instance of the required flow's args_type class attribute. flow_cls: Class of the flow that should be started. network_bytes_limit: Limit on the network traffic this flow can generated. original_flow: A FlowReference object in case this flow was copied from another flow. output_plugins: An OutputPluginDescriptor object indicating what output plugins should be used for this flow. parent_flow_obj: A parent flow object. None if this is a top level flow. **kwargs: If args or runner_args are not specified, we construct these protobufs from these keywords. Returns: the flow id of the new flow. Raises: ValueError: Unknown or invalid parameters were provided. """ # Is the required flow a known flow? try: registry.FlowRegistry.FlowClassByName(flow_cls.__name__) except ValueError: stats_collector_instance.Get().IncrementCounter( "grr_flow_invalid_flow_count") raise ValueError("Unable to locate flow %s" % flow_cls.__name__) if not client_id: raise ValueError("Client_id is needed to start a flow.") # Now parse the flow args into the new object from the keywords. if flow_args is None: flow_args = flow_cls.args_type() FilterArgsFromSemanticProtobuf(flow_args, kwargs) # At this point we should exhaust all the keyword args. If any are left # over, we do not know what to do with them so raise. if kwargs: raise type_info.UnknownArg("Unknown parameters to StartFlow: %s" % kwargs) # Check that the flow args are valid. flow_args.Validate() rdf_flow = rdf_flow_objects.Flow(client_id=client_id, flow_class_name=flow_cls.__name__, args=flow_args, create_time=rdfvalue.RDFDatetime.Now(), creator=creator, output_plugins=output_plugins, original_flow=original_flow, flow_state="RUNNING") rdf_flow.flow_id = "%08X" % random.PositiveUInt32() if parent_flow_obj: parent_rdf_flow = parent_flow_obj.rdf_flow rdf_flow.long_flow_id = "%s/%s" % (parent_rdf_flow.long_flow_id, rdf_flow.flow_id) rdf_flow.parent_flow_id = parent_rdf_flow.flow_id rdf_flow.parent_request_id = parent_flow_obj.GetCurrentOutboundId() if parent_rdf_flow.creator: rdf_flow.creator = parent_rdf_flow.creator else: rdf_flow.long_flow_id = "%s/%s" % (client_id, rdf_flow.flow_id) if output_plugins: rdf_flow.output_plugins_states = GetOutputPluginStates( output_plugins, rdf_flow.long_flow_id, token=access_control.ACLToken(username=rdf_flow.creator)) if network_bytes_limit is not None: rdf_flow.network_bytes_limit = network_bytes_limit if cpu_limit is not None: rdf_flow.cpu_limit = cpu_limit logging.info(u"Scheduling %s(%s) on %s", rdf_flow.long_flow_id, rdf_flow.flow_class_name, client_id) flow_obj = flow_cls(rdf_flow) # Just run the first state inline. NOTE: Running synchronously means # that this runs on the thread that starts the flow. The advantage is # that that Start method can raise any errors immediately. flow_obj.Start() flow_obj.PersistState() # The flow does not need to actually remain running. if not flow_obj.outstanding_requests: flow_obj.RunStateMethod("End") flow_obj.MarkDone() data_store.REL_DB.WriteFlowObject(flow_obj.rdf_flow) if parent_flow_obj is not None: # We can optimize here and not write requests/responses to the database # since we have to do this for the parent flow at some point anyways. parent_flow_obj.MergeQueuedMessages(flow_obj) else: flow_obj.FlushQueuedMessages() # Publish an audit event, only for top level flows. # TODO(amoser): split urn field into dedicated strings. events.Events.PublishEvent( "Audit", rdf_events.AuditEvent(user=creator, action="RUN_FLOW", flow_name=rdf_flow.flow_class_name, urn=rdf_flow.long_flow_id, client=client_id)) return rdf_flow.flow_id
def GetApprovalForObject(object_urn, token=None, username=""): """Looks for approvals for an object and returns available valid tokens. Args: object_urn: Urn of the object we want access to. token: The token to use to lookup the ACLs. username: The user to get the approval for, if "" we get it from the token. Returns: A token for access to the object on success, otherwise raises. Raises: UnauthorizedAccess: If there are no valid approvals available. """ if token is None: raise access_control.UnauthorizedAccess( "No token given, cannot authenticate.") if not username: username = token.username approvals_root_urn = aff4.ROOT_URN.Add("ACL").Add( object_urn.Path()).Add(username) children_urns = list(aff4.FACTORY.ListChildren(approvals_root_urn)) if not children_urns: raise access_control.UnauthorizedAccess( "No approval found for user %s" % utils.SmartStr(username), subject=object_urn) last_error = None approvals = aff4.FACTORY.MultiOpen( children_urns, mode="r", aff4_type=Approval, age=aff4.ALL_TIMES, token=token) for approval in approvals: try: test_token = access_control.ACLToken( username=username, reason=approval.Get(approval.Schema.REASON)) approval.CheckAccess(test_token) return test_token except access_control.UnauthorizedAccess as e: last_error = e if last_error: # We tried all possible approvals, but got no usable results. raise access_control.UnauthorizedAccess(last_error, subject=object_urn) else: # If last error is None, means that none of the URNs in children_urns # could be opened. This shouldn't really happen ever, but we have # to make sure to provide a meaningful error message. raise access_control.UnauthorizedAccess( "Couldn't open any of %d approvals " "for user %s" % (len(children_urns), utils.SmartStr(username)), subject=object_urn)
def ListCronJobApprovals(self, requestor=None): requestor = requestor or self.token.username handler = api_user.ApiListCronJobApprovalsHandler() return handler.Handle( api_user.ApiListCronJobApprovalsArgs(), token=access_control.ACLToken(username=requestor)).items
def InitGRRRootAPI(): return api.GrrApi(connector=api_shell_raw_access_lib.RawConnector( token=access_control.ACLToken(username="******"), page_size=_GRR_API_PAGE_SIZE)).root
def SecurityCheck(self, func, request, *args, **kwargs): """A decorator applied to protected web handlers.""" request.user = self.username request.token = access_control.ACLToken(username="******", reason="Just a test") return func(request, *args, **kwargs)
def testRendersObjectForNonExistingUser(self): result = self.handler.Handle( None, token=access_control.ACLToken(username="******")) self.assertEqual(result.username, "foo")
def GetToken(): # Extend for user authorization # SetUID is required to create and write to various aff4 paths when updating # config. return access_control.ACLToken(username="******").SetUID()
def StartFlow(client_id=None, cpu_limit=None, creator=None, flow_args=None, flow_cls=None, network_bytes_limit=None, original_flow=None, output_plugins=None, start_at=None, parent_flow_obj=None, parent_hunt_id=None, **kwargs): """The main factory function for creating and executing a new flow. Args: client_id: ID of the client this flow should run on. cpu_limit: CPU limit in seconds for this flow. creator: Username that requested this flow. flow_args: An arg protocol buffer which is an instance of the required flow's args_type class attribute. flow_cls: Class of the flow that should be started. network_bytes_limit: Limit on the network traffic this flow can generated. original_flow: A FlowReference object in case this flow was copied from another flow. output_plugins: An OutputPluginDescriptor object indicating what output plugins should be used for this flow. start_at: If specified, flow will be started not immediately, but at a given time. parent_flow_obj: A parent flow object. None if this is a top level flow. parent_hunt_id: String identifying parent hunt. Can't be passed together with parent_flow_obj. **kwargs: If args or runner_args are not specified, we construct these protobufs from these keywords. Returns: the flow id of the new flow. Raises: ValueError: Unknown or invalid parameters were provided. """ if parent_flow_obj is not None and parent_hunt_id is not None: raise ValueError( "parent_flow_obj and parent_hunt_id are mutually exclusive.") # Is the required flow a known flow? try: registry.FlowRegistry.FlowClassByName(flow_cls.__name__) except ValueError: stats_collector_instance.Get().IncrementCounter( "grr_flow_invalid_flow_count") raise ValueError("Unable to locate flow %s" % flow_cls.__name__) if not client_id: raise ValueError("Client_id is needed to start a flow.") # Now parse the flow args into the new object from the keywords. if flow_args is None: flow_args = flow_cls.args_type() FilterArgsFromSemanticProtobuf(flow_args, kwargs) # At this point we should exhaust all the keyword args. If any are left # over, we do not know what to do with them so raise. if kwargs: raise type_info.UnknownArg("Unknown parameters to StartFlow: %s" % kwargs) # Check that the flow args are valid. flow_args.Validate() rdf_flow = rdf_flow_objects.Flow(client_id=client_id, flow_class_name=flow_cls.__name__, args=flow_args, create_time=rdfvalue.RDFDatetime.Now(), creator=creator, output_plugins=output_plugins, original_flow=original_flow, flow_state="RUNNING") if parent_hunt_id is not None and parent_flow_obj is None: rdf_flow.flow_id = parent_hunt_id if IsLegacyHunt(parent_hunt_id): rdf_flow.flow_id = rdf_flow.flow_id[2:] else: rdf_flow.flow_id = RandomFlowId() # For better performance, only do conflicting IDs check for top-level flows. if not parent_flow_obj: try: data_store.REL_DB.ReadFlowObject(client_id, rdf_flow.flow_id) raise CanNotStartFlowWithExistingIdError(client_id, rdf_flow.flow_id) except db.UnknownFlowError: pass if parent_flow_obj: # A flow is a nested flow. parent_rdf_flow = parent_flow_obj.rdf_flow rdf_flow.long_flow_id = "%s/%s" % (parent_rdf_flow.long_flow_id, rdf_flow.flow_id) rdf_flow.parent_flow_id = parent_rdf_flow.flow_id rdf_flow.parent_hunt_id = parent_rdf_flow.parent_hunt_id rdf_flow.parent_request_id = parent_flow_obj.GetCurrentOutboundId() if parent_rdf_flow.creator: rdf_flow.creator = parent_rdf_flow.creator elif parent_hunt_id: # A flow is a root-level hunt-induced flow. rdf_flow.long_flow_id = "%s/%s" % (client_id, rdf_flow.flow_id) rdf_flow.parent_hunt_id = parent_hunt_id else: # A flow is a root-level non-hunt flow. rdf_flow.long_flow_id = "%s/%s" % (client_id, rdf_flow.flow_id) if output_plugins: rdf_flow.output_plugins_states = GetOutputPluginStates( output_plugins, rdf_flow.long_flow_id, token=access_control.ACLToken(username=rdf_flow.creator)) if network_bytes_limit is not None: rdf_flow.network_bytes_limit = network_bytes_limit if cpu_limit is not None: rdf_flow.cpu_limit = cpu_limit logging.info(u"Scheduling %s(%s) on %s (%s)", rdf_flow.long_flow_id, rdf_flow.flow_class_name, client_id, start_at or "now") rdf_flow.current_state = "Start" flow_obj = flow_cls(rdf_flow) if start_at is None: # Store an initial version of the flow straight away. This is needed so the # database doesn't raise consistency errors due to missing parent keys when # writing logs / errors / results which might happen in Start(). data_store.REL_DB.WriteFlowObject(flow_obj.rdf_flow) # Just run the first state inline. NOTE: Running synchronously means # that this runs on the thread that starts the flow. The advantage is # that that Start method can raise any errors immediately. flow_obj.Start() # The flow does not need to actually remain running. if not flow_obj.outstanding_requests: flow_obj.RunStateMethod("End") # Additional check for the correct state in case the End method raised and # terminated the flow. if flow_obj.IsRunning(): flow_obj.MarkDone() else: flow_obj.CallState("Start", start_time=start_at) flow_obj.PersistState() data_store.REL_DB.WriteFlowObject(flow_obj.rdf_flow) if parent_flow_obj is not None: # We can optimize here and not write requests/responses to the database # since we have to do this for the parent flow at some point anyways. parent_flow_obj.MergeQueuedMessages(flow_obj) else: flow_obj.FlushQueuedMessages() # Publish an audit event, only for top level flows. # TODO(amoser): split urn field into dedicated strings. events.Events.PublishEvent( "Audit", rdf_events.AuditEvent(user=creator, action="RUN_FLOW", flow_name=rdf_flow.flow_class_name, urn=rdf_flow.long_flow_id, client=client_id)) return rdf_flow.flow_id
def __init__(self): self.token = access_control.ACLToken(username="******", reason="Implied.") self.token.supervisor = True
def setUp(self): super(GRRBaseTest, self).setUp() test_user = u"test" system_users_patcher = mock.patch.object( access_control, "SYSTEM_USERS", frozenset(itertools.chain(access_control.SYSTEM_USERS, [test_user]))) system_users_patcher.start() self.addCleanup(system_users_patcher.stop) self.token = access_control.ACLToken(username=test_user, reason="Running tests") self.temp_dir = temp.TempDirPath() config.CONFIG.SetWriteBack( os.path.join(self.temp_dir, "writeback.yaml")) self.addCleanup( lambda: shutil.rmtree(self.temp_dir, ignore_errors=True)) # Each datastore is wrapped with DatabaseValidationWrapper, so we have # to access the delegate directly (assuming it's an InMemoryDB # implementation). data_store.REL_DB.delegate.ClearTestDB() email_alerts.InitializeEmailAlerterOnce() # Stub out the email function self.emails_sent = [] def SendEmailStub(to_user, from_user, subject, message, **unused_kwargs): self.emails_sent.append((to_user, from_user, subject, message)) self.mail_stubber = utils.MultiStubber( (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub), (email.utils, "make_msgid", lambda: "<message id stub>")) self.mail_stubber.Start() self.addCleanup(self.mail_stubber.Stop) # We don't want to send actual email in our tests self.smtp_patcher = mock.patch("smtplib.SMTP") self.mock_smtp = self.smtp_patcher.start() self.addCleanup(self.smtp_patcher.stop) def DisabledSet(*unused_args, **unused_kw): raise NotImplementedError( "Usage of Set() is disabled, please use a configoverrider in tests." ) self.config_set_disable = utils.Stubber(config.CONFIG, "Set", DisabledSet) self.config_set_disable.Start() self.addCleanup(self.config_set_disable.Stop) self._SetupFakeStatsContext() # Turn off WithLimitedCallFrequency-based caching in tests. Tests that need # to test caching behavior explicitly, should turn it on explicitly. with_limited_call_frequency_stubber = utils.Stubber( cache, "WITH_LIMITED_CALL_FREQUENCY_PASS_THROUGH", True) with_limited_call_frequency_stubber.Start() self.addCleanup(with_limited_call_frequency_stubber.Stop)
def testTokenWithoutUsernameIsNotValid(self): token = access_control.ACLToken(reason="For testing") with self.assertRaises(access_control.UnauthorizedAccess): user_managers.ValidateToken(token, "aff4:/C.0000000000000001")
def TestHuntHelperWithMultipleMocks(client_mocks, check_flow_errors=False, token=None, iteration_limit=None, worker=None): """Runs a hunt with a given set of clients mocks. Args: client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock objects are used to handle client actions. Methods names of a client mock object correspond to client actions names. For an example of a client mock object, see SampleHuntMock. check_flow_errors: If True, raises when one of hunt-initiated flows fails. token: An instance of access_control.ACLToken security token. iteration_limit: If None, hunt will run until it's finished. Otherwise, worker_mock.Next() will be called iteration_limit number of times. Every iteration processes worker's message queue. If new messages are sent to the queue during the iteration processing, they will be processed on next iteration. worker: flow_test_lib.TestWorker object to use. Returns: A number of iterations complete. """ if token is None: token = access_control.ACLToken(username="******") total_flows = set() # Worker always runs with absolute privileges, therefore making the token # SetUID(). token = token.SetUID() client_mocks = [ flow_test_lib.MockClient(client_id, client_mock, token=token) for client_id, client_mock in iteritems(client_mocks) ] if worker is None: rel_db_worker = flow_test_lib.TestWorker(threadpool_size=0, token=True) data_store.REL_DB.RegisterFlowProcessingHandler( rel_db_worker.ProcessFlow) else: rel_db_worker = worker num_iterations = 0 try: worker_mock = worker_test_lib.MockWorker( check_flow_errors=check_flow_errors, token=token) # Run the clients and worker until nothing changes any more. while iteration_limit is None or num_iterations < iteration_limit: worker_processed = [] if data_store.RelationalDBEnabled(): data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess( timeout=10) worker_processed = rel_db_worker.ResetProcessedFlows() client_processed = 0 for client_mock in client_mocks: client_processed += client_mock.Next() flows_run = [] for flow_run in worker_mock.Next(): total_flows.add(flow_run) flows_run.append(flow_run) flows_run.extend(worker_processed) num_iterations += 1 if client_processed == 0 and not flows_run and not worker_processed: break if check_flow_errors: flow_test_lib.CheckFlowErrors(total_flows, token=token) finally: if worker is None: data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60) rel_db_worker.Shutdown() return num_iterations
def CheckHuntAccess(self, username, hunt_id): token = access_control.ACLToken(username=username) self.legacy_manager.CheckHuntAccess(token, hunt_id.ToURN())
def GetToken(): user = getpass.getuser() return access_control.ACLToken(username=user)
def CheckIfCanStartClientFlow(self, username, flow_name): token = access_control.ACLToken(username=username) self.legacy_manager.CheckIfCanStartFlow(token, flow_name)
def StartFlow(client_id=None, cpu_limit=None, creator=None, flow_args=None, flow_cls=None, network_bytes_limit=None, original_flow=None, output_plugins=None, start_at=None, parent=None, runtime_limit=None, **kwargs): """The main factory function for creating and executing a new flow. Args: client_id: ID of the client this flow should run on. cpu_limit: CPU limit in seconds for this flow. creator: Username that requested this flow. flow_args: An arg protocol buffer which is an instance of the required flow's args_type class attribute. flow_cls: Class of the flow that should be started. network_bytes_limit: Limit on the network traffic this flow can generated. original_flow: A FlowReference object in case this flow was copied from another flow. output_plugins: An OutputPluginDescriptor object indicating what output plugins should be used for this flow. start_at: If specified, flow will be started not immediately, but at a given time. parent: A FlowParent referencing the parent, or None for top-level flows. runtime_limit: Runtime limit as Duration for all ClientActions. **kwargs: If args or runner_args are not specified, we construct these protobufs from these keywords. Returns: the flow id of the new flow. Raises: ValueError: Unknown or invalid parameters were provided. """ # Is the required flow a known flow? try: registry.FlowRegistry.FlowClassByName(flow_cls.__name__) except ValueError: GRR_FLOW_INVALID_FLOW_COUNT.Increment() raise ValueError("Unable to locate flow %s" % flow_cls.__name__) if not client_id: raise ValueError("Client_id is needed to start a flow.") # Now parse the flow args into the new object from the keywords. if flow_args is None: flow_args = flow_cls.args_type() FilterArgsFromSemanticProtobuf(flow_args, kwargs) # At this point we should exhaust all the keyword args. If any are left # over, we do not know what to do with them so raise. if kwargs: raise type_info.UnknownArg("Unknown parameters to StartFlow: %s" % kwargs) # Check that the flow args are valid. flow_args.Validate() rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_class_name=flow_cls.__name__, args=flow_args, create_time=rdfvalue.RDFDatetime.Now(), creator=creator, output_plugins=output_plugins, original_flow=original_flow, flow_state="RUNNING") if parent is None: parent = FlowParent.FromRoot() if parent.is_hunt or parent.is_scheduled_flow: # When starting a flow from a hunt or ScheduledFlow, re-use the parent's id # to make it easy to find flows. For hunts, every client has a top-level # flow with the hunt's id. rdf_flow.flow_id = parent.id else: # For new top-level and child flows, assign a random ID. rdf_flow.flow_id = RandomFlowId() # For better performance, only do conflicting IDs check for top-level flows. if not parent.is_flow: try: data_store.REL_DB.ReadFlowObject(client_id, rdf_flow.flow_id) raise CanNotStartFlowWithExistingIdError(client_id, rdf_flow.flow_id) except db.UnknownFlowError: pass if parent.is_flow: # A flow is a nested flow. parent_rdf_flow = parent.flow_obj.rdf_flow rdf_flow.long_flow_id = "%s/%s" % (parent_rdf_flow.long_flow_id, rdf_flow.flow_id) rdf_flow.parent_flow_id = parent_rdf_flow.flow_id rdf_flow.parent_hunt_id = parent_rdf_flow.parent_hunt_id rdf_flow.parent_request_id = parent.flow_obj.GetCurrentOutboundId() if parent_rdf_flow.creator: rdf_flow.creator = parent_rdf_flow.creator elif parent.is_hunt: # Root-level hunt-induced flow. rdf_flow.long_flow_id = "%s/%s" % (client_id, rdf_flow.flow_id) rdf_flow.parent_hunt_id = parent.id elif parent.is_root or parent.is_scheduled_flow: # A flow is a root-level non-hunt flow. rdf_flow.long_flow_id = "%s/%s" % (client_id, rdf_flow.flow_id) else: raise ValueError(f"Unknown flow parent type {parent}") if output_plugins: rdf_flow.output_plugins_states = GetOutputPluginStates( output_plugins, rdf_flow.long_flow_id, token=access_control.ACLToken(username=rdf_flow.creator)) if network_bytes_limit is not None: rdf_flow.network_bytes_limit = network_bytes_limit if cpu_limit is not None: rdf_flow.cpu_limit = cpu_limit if runtime_limit is not None: rdf_flow.runtime_limit_us = runtime_limit logging.info(u"Starting %s(%s) on %s (%s)", rdf_flow.long_flow_id, rdf_flow.flow_class_name, client_id, start_at or "now") rdf_flow.current_state = "Start" flow_obj = flow_cls(rdf_flow) # Prevent a race condition, where a flow is scheduled twice, because one # worker inserts the row and another worker silently updates the existing row. allow_update = False if start_at is None: # Store an initial version of the flow straight away. This is needed so the # database doesn't raise consistency errors due to missing parent keys when # writing logs / errors / results which might happen in Start(). try: data_store.REL_DB.WriteFlowObject(flow_obj.rdf_flow, allow_update=False) except db.FlowExistsError: raise CanNotStartFlowWithExistingIdError(client_id, rdf_flow.flow_id) allow_update = True try: # Just run the first state inline. NOTE: Running synchronously means # that this runs on the thread that starts the flow. The advantage is # that that Start method can raise any errors immediately. flow_obj.Start() # The flow does not need to actually remain running. if not flow_obj.outstanding_requests: flow_obj.RunStateMethod("End") # Additional check for the correct state in case the End method raised # and terminated the flow. if flow_obj.IsRunning(): flow_obj.MarkDone() except Exception as e: # pylint: disable=broad-except # We catch all exceptions that happen in Start() and mark the flow as # failed. msg = compatibility.NativeStr(e) if compatibility.PY2: msg = msg.decode("utf-8", "replace") flow_obj.Error(error_message=msg, backtrace=traceback.format_exc()) else: flow_obj.CallState("Start", start_time=start_at) flow_obj.PersistState() try: data_store.REL_DB.WriteFlowObject( flow_obj.rdf_flow, allow_update=allow_update) except db.FlowExistsError: raise CanNotStartFlowWithExistingIdError(client_id, rdf_flow.flow_id) if parent.is_flow: # We can optimize here and not write requests/responses to the database # since we have to do this for the parent flow at some point anyways. parent.flow_obj.MergeQueuedMessages(flow_obj) else: flow_obj.FlushQueuedMessages() return rdf_flow.flow_id