def testStuckNotificationGetsDeletedAfterTheFlowIsTerminated(self): worker_obj = self._TestWorker() initial_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100) stuck_flows_timeout = flow_runner.FlowRunner.stuck_flows_timeout try: with test_lib.FakeTime(initial_time.AsSecondsSinceEpoch()): session_id = flow.StartAFF4Flow( flow_name=WorkerStuckableTestFlow.__name__, client_id=self.client_id, token=self.token, sync=False) # Process all messages worker_obj.RunOnce() # Wait until worker thread starts processing the flow. WorkerStuckableTestFlow.WaitUntilWorkerStartsProcessing() # Set the time to max worker flow duration + 1 minute. The flow is # currently blocked because of the way how semaphores are set up. # Worker should consider the flow to be stuck and terminate it. future_time = (initial_time + rdfvalue.Duration("1m") + stuck_flows_timeout) with test_lib.FakeTime(future_time.AsSecondsSinceEpoch()): worker_obj.RunOnce() killed_flow = aff4.FACTORY.Open(session_id, token=self.token) self.assertEqual(killed_flow.context.state, rdf_flow_runner.FlowContext.State.ERROR) self.assertEqual( killed_flow.context.status, "Terminated by user test. Reason: Stuck in the worker") # Check that stuck notification has been removed. qm = queue_manager.QueueManager(token=self.token) notifications = qm.GetNotifications(queues.FLOWS) for n in notifications: self.assertFalse(n.in_progress) finally: # Release the semaphore so that worker thread unblocks and finishes # processing the flow. WorkerStuckableTestFlow.StopFlow() WorkerStuckableTestFlow.LetWorkerFinishProcessing() worker_obj.thread_pool.Join()
def Fn(*args, **kwargs): """Wrapper around the decorated function.""" if WITH_LIMITED_CALL_FREQUENCY_PASS_THROUGH: # This effectively turns off the caching. min_time = rdfvalue.Duration(0) else: min_time = min_time_between_calls key = (args, tuple(sorted(kwargs.items()))) now = rdfvalue.RDFDatetime.Now() with lock: for k, prev_time in list(prev_times.items()): if now - prev_time >= min_time: prev_times.pop(k) prev_results.pop(k, None) result_locks.pop(k, None) try: # We eliminated all the old entries, so if the key is present # in the cache, it means that the data is fresh enough to be used. return prev_results[key] except KeyError: prev_time = None try: result_lock = result_locks[key] except KeyError: result_lock = threading.RLock() result_locks[key] = result_lock with result_lock: t = prev_times.get(key) if t == prev_time: result = f(*args, **kwargs) with lock: prev_times[key] = rdfvalue.RDFDatetime.Now() prev_results[key] = result return result else: return prev_results[key]
def RunOnce(self, names=None, token=None): """Tries to lock and run cron jobs. Args: names: List of cron jobs to run. If unset, run them all. token: security token. """ del token leased_jobs = data_store.REL_DB.LeaseCronJobs( cronjob_ids=names, lease_time=rdfvalue.Duration("10m")) if not leased_jobs: return for job in leased_jobs: logging.info("Running cron job: %s", job.cron_job_id) self.RunJob(job) data_store.REL_DB.ReturnLeasedCronJobs(leased_jobs)
def QueryAndOwn(self, queue, lease_seconds=10, limit=1): """Returns a list of Tasks leased for a certain time. Args: queue: The queue to query from. lease_seconds: The tasks will be leased for this long. limit: Number of values to fetch. Returns: A list of GrrMessage() objects leased. """ if data_store.RelationalDBReadEnabled(category="client_messages"): return data_store.REL_DB.LeaseClientMessages( queue.Split()[0], lease_time=rdfvalue.Duration("%ds" % lease_seconds), limit=limit) with self.data_store.GetMutationPool() as mutation_pool: return mutation_pool.QueueQueryAndOwn(queue, lease_seconds, limit, self.frozen_timestamp)
def ProcessFlow(self, flow_processing_request): """The callback for the flow processing queue.""" client_id = flow_processing_request.client_id flow_id = flow_processing_request.flow_id data_store.REL_DB.AckFlowProcessingRequests([flow_processing_request]) try: rdf_flow = data_store.REL_DB.LeaseFlowForProcessing( client_id, flow_id, processing_time=rdfvalue.Duration("6h")) except db.ParentHuntIsNotRunningError: flow_base.TerminateFlow(client_id, flow_id, "Parent hunt stopped.") return logging.info("Processing Flow %s/%s/%d (%s).", client_id, flow_id, rdf_flow.next_request_to_process, rdf_flow.flow_class_name) flow_cls = registry.FlowRegistry.FlowClassByName( rdf_flow.flow_class_name) flow_obj = flow_cls(rdf_flow) if not flow_obj.IsRunning(): logging.info( "Received a request to process flow %s on client %s that is not " "running.", flow_id, client_id) return processed = flow_obj.ProcessAllReadyRequests() if processed == 0: raise ValueError( "Unable to process any requests for flow %s on client %s." % (flow_id, client_id)) while not self._ReleaseProcessedFlow(flow_obj): processed = flow_obj.ProcessAllReadyRequests() if processed == 0: raise ValueError( "%s/%s: ReleaseProcessedFlow returned false but no " "request could be processed (next req: %d)." % (client_id, flow_id, flow_obj.rdf_flow.next_request_to_process))
class OSBreakDown(AbstractClientStatsCronFlow): """Records relative ratios of OS versions in 7 day actives.""" recency_window = rdfvalue.Duration("30d") def BeginProcessing(self): self.counters = [ _ActiveCounter(rdf_stats.ClientGraphSeries.ReportType.OS_TYPE), _ActiveCounter(rdf_stats.ClientGraphSeries.ReportType.OS_RELEASE), ] def FinishProcessing(self): # Write all the counter attributes. for counter in self.counters: counter.Save(self.token) def _Process(self, labels, ping, system, uname): if not ping: return for label in labels: # Windows, Linux, Darwin self.counters[0].Add(system, label, ping) # Windows-2008ServerR2-6.1.7601SP1, Linux-Ubuntu-12.04, # Darwin-OSX-10.9.3 self.counters[1].Add(uname, label, ping) def ProcessLegacyClient(self, ping, client): """Update counters for system, version and release attributes.""" labels = self._GetClientLabelsList(client) system = client.Get(client.Schema.SYSTEM, "Unknown") uname = client.Get(client.Schema.UNAME, "Unknown") self._Process(labels, ping, system, uname) def ProcessClientFullInfo(self, client_full_info): labels = self._GetClientLabelsList(client_full_info) ping = client_full_info.metadata.ping system = client_full_info.last_snapshot.knowledge_base.os uname = client_full_info.last_snapshot.Uname() self._Process(labels, ping, system, uname)
def testCronApprovalsReportPluginWithNoActivityToReport(self): report = report_plugins.GetReportByName( server_report_plugins.CronApprovalsReportPlugin.__name__) now = rdfvalue.RDFDatetime().Now() month_duration = rdfvalue.Duration("30d") api_report_data = report.GetReportData( stats_api.ApiGetReportArgs( name=report.__class__.__name__, start_time=now - month_duration, duration=month_duration), token=self.token) self.assertEqual(api_report_data.representation_type, RepresentationType.AUDIT_CHART) self.assertCountEqual(api_report_data.audit_chart.used_fields, ["action", "timestamp", "user", "urn"]) self.assertEmpty(api_report_data.audit_chart.rows)
def testUserFlowsReportPluginWithNoActivityToReport(self): report = report_plugins.GetReportByName( server_report_plugins.UserFlowsReportPlugin.__name__) now = rdfvalue.RDFDatetime().Now() month_duration = rdfvalue.Duration("30d") api_report_data = report.GetReportData( stats_api.ApiGetReportArgs( name=report.__class__.__name__, start_time=now - month_duration, duration=month_duration), token=self.token) self.assertEqual( api_report_data, rdf_report_plugins.ApiReportData( representation_type=RepresentationType.STACK_CHART, stack_chart=rdf_report_plugins.ApiStackChartReportData(x_ticks=[])))
def testMostActiveUsersReportPluginWithNoActivityToReport(self): report = report_plugins.GetReportByName( server_report_plugins.MostActiveUsersReportPlugin.__name__) now = rdfvalue.RDFDatetime().Now() month_duration = rdfvalue.Duration("30d") api_report_data = report.GetReportData(stats_api.ApiGetReportArgs( name=report.__class__.__name__, start_time=now - month_duration, duration=month_duration), token=self.token) self.assertEqual( api_report_data, rdf_report_plugins.ApiReportData( representation_type=rdf_report_plugins.ApiReportData. RepresentationType.PIE_CHART, pie_chart=rdf_report_plugins.ApiPieChartReportData(data=[])))
def _AnalyzeKeywords(self, keywords): """Extracts a start time from a list of keywords if present.""" start_time = rdfvalue.RDFDatetime.Now() - rdfvalue.Duration("180d") filtered_keywords = [] for k in keywords: if k.startswith(self.START_TIME_PREFIX): try: start_time = rdfvalue.RDFDatetime.FromHumanReadable( k[self.START_TIME_PREFIX_LEN:]) except ValueError: pass else: filtered_keywords.append(k) if not filtered_keywords: filtered_keywords.append(".") return start_time, filtered_keywords
def testKeepsHuntsWithRetainLabel(self): exception_label_name = config.CONFIG[ "DataRetention.hunts_ttl_exception_label"] for hunt_urn in self.hunts_urns[:3]: with aff4.FACTORY.Open(hunt_urn, mode="rw", token=self.token) as fd: fd.AddLabel(exception_label_name) with test_lib.ConfigOverrider( {"DataRetention.hunts_ttl": rdfvalue.Duration("10s")}): with test_lib.FakeTime(40 + 60 * self.NUM_HUNTS): self._RunCleanup() hunts_urns = list( aff4.FACTORY.Open("aff4:/hunts", token=self.token).ListChildren()) self.assertEqual(len(hunts_urns), 3)
def testHuntTermination(self): """This tests that hunts with a client limit terminate correctly.""" with test_lib.FakeTime(1000, increment=1e-6): with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)), client_rule_set=self._CreateForemanClientRuleSet(), client_limit=5, client_rate=0, expiry_time=rdfvalue.Duration("1000s"), token=self.token) as hunt: hunt.Run() # Pretend to be the foreman now and dish out hunting jobs to all the # clients (Note we have 10 clients here). self.AssignTasksToClients() # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper( client_mock, self.client_ids, check_flow_errors=False, token=self.token) hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) started, finished, errors = hunt_obj.GetClientsCounts() self.assertEqual(started, 5) self.assertEqual(finished, 5) self.assertEqual(errors, 2) hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) # Hunts are automatically paused when they reach the client limit. self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "PAUSED")
def testHuntActionsReportPluginWithNoActivityToReport(self): report = report_plugins.GetReportByName( server_report_plugins.HuntActionsReportPlugin.__name__) now = rdfvalue.RDFDatetime().Now() month_duration = rdfvalue.Duration("30d") api_report_data = report.GetReportData(stats_api.ApiGetReportArgs( name=report.__class__.__name__, start_time=now - month_duration, duration=month_duration), token=self.token) self.assertEqual( api_report_data, rdf_report_plugins.ApiReportData( representation_type=RepresentationType.AUDIT_CHART, audit_chart=rdf_report_plugins.ApiAuditChartReportData( used_fields=["action", "timestamp", "user"], rows=[])))
def setUp(self): super(SystemCronTestMixin, self).setUp() # This is not optimal, we create clients 0-19 with Linux, then # overwrite clients 0-9 with Windows, leaving 10-19 for Linux. client_ping_time = rdfvalue.RDFDatetime.Now() - rdfvalue.Duration("8d") self.SetupClients(20, system="Linux", ping=client_ping_time) self.SetupClients(10, system="Windows", ping=client_ping_time) for i in range(0, 10): client_id = "C.1%015x" % i with aff4.FACTORY.Open(client_id, mode="rw", token=self.token) as client: client.AddLabels([u"Label1", u"Label2"], owner="GRR") client.AddLabel(u"UserLabel", owner="jim") data_store.REL_DB.AddClientLabels(client_id, "GRR", [u"Label1", u"Label2"]) data_store.REL_DB.AddClientLabels(client_id, "jim", [u"UserLabel"])
class ClientStats(rdf_structs.RDFProtoStruct): """A client stat object.""" protobuf = jobs_pb2.ClientStats rdf_deps = [ CpuSample, IOSample, rdfvalue.RDFDatetime, ] DEFAULT_SAMPLING_INTERVAL = rdfvalue.Duration("60s") @classmethod def Downsampled(cls, stats, interval=None): """Constructs a copy of given stats but downsampled to given interval. Args: stats: A `ClientStats` instance. interval: A downsampling interval. Returns: A downsampled `ClientStats` instance. """ interval = interval or cls.DEFAULT_SAMPLING_INTERVAL result = cls(stats) result.cpu_samples = cls._Downsample(kind=CpuSample, samples=stats.cpu_samples, interval=interval) result.io_samples = cls._Downsample(kind=IOSample, samples=stats.io_samples, interval=interval) return result @classmethod def _Downsample(cls, kind, samples, interval): buckets = {} for sample in samples: bucket = buckets.setdefault(sample.timestamp.Floor(interval), []) bucket.append(sample) for bucket in itervalues(buckets): yield kind.FromMany(bucket)
def Next(self): """Grab tasks for us from the server's queue.""" with queue_manager.QueueManager(token=self.token) as manager: if data_store.RelationalDBEnabled(): request_tasks = data_store.REL_DB.LeaseClientActionRequests( self.client_id.Basename(), lease_time=rdfvalue.Duration("10000s"), limit=1) request_tasks = [ rdf_flow_objects.GRRMessageFromClientActionRequest(r) for r in request_tasks ] else: request_tasks = manager.QueryAndOwn(self.client_id.Queue(), limit=1, lease_seconds=10000) request_tasks.extend(self._mock_task_queue) self._mock_task_queue[:] = [] # Clear the referenced list. for message in request_tasks: try: responses = self.client_mock.HandleMessage(message) logging.info( "Called client action %s generating %s responses", message.name, len(responses) + 1) except Exception as e: # pylint: disable=broad-except logging.exception("Error %s occurred in client", e) responses = [ self.client_mock.GenerateStatusMessage( message, 1, status="GENERIC_ERROR") ] # Now insert those on the flow state queue for response in responses: self.PushToStateQueue(manager, response) # Additionally schedule a task for the worker manager.QueueNotification(session_id=message.session_id) return len(request_tasks)
def testFlowStateUpdateUsingReturnProcessedFlow(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) hunt_id = hunt_obj.hunt_id client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id) flow_obj = self.db.ReadFlowForProcessing(client_id, flow_id, rdfvalue.Duration("1m")) self.assertEqual(flow_obj.flow_state, rdf_flow_objects.Flow.FlowState.UNSET) flow_obj.flow_state = rdf_flow_objects.Flow.FlowState.ERROR self.db.ReturnProcessedFlow(flow_obj) results = self.db.ReadHuntFlows( hunt_id, 0, 10, filter_condition=db.HuntFlowsCondition.FAILED_FLOWS_ONLY) self.assertLen(results, 1)
def ProcessHuntOutputPlugins(self): if data_store.RelationalDBReadEnabled("hunts"): # No processing needed for new style hunts. return if data_store.RelationalDBFlowsEnabled(): job = rdf_cronjobs.CronJob( cron_job_id="some/id", lifetime=rdfvalue.Duration("1h")) run_state = rdf_cronjobs.CronJobRun( cron_job_id="some/id", status="RUNNING", started_at=rdfvalue.RDFDatetime.Now()) process_results.ProcessHuntResultCollectionsCronJob(run_state, job).Run() else: flow_urn = flow.StartAFF4Flow( flow_name=process_results.ProcessHuntResultCollectionsCronFlow .__name__, token=self.token) flow_test_lib.TestFlowHelper(flow_urn, token=self.token) return flow_urn
def GetReportData(self, get_report_args, token): """Filter the last week of user actions.""" ret = rdf_report_plugins.ApiReportData( representation_type=RepresentationType.STACK_CHART) week_duration = rdfvalue.Duration.From(7, rdfvalue.DAYS) num_weeks = int( math.ceil( rdfvalue.Duration(get_report_args.duration).ToFractional( rdfvalue.SECONDS) / week_duration.ToFractional(rdfvalue.SECONDS))) weeks = range(0, num_weeks) start_time = get_report_args.start_time end_time = start_time + num_weeks * week_duration user_activity = collections.defaultdict( lambda: {week: 0 for week in weeks}) entries = self._LoadUserActivity(start_time=get_report_args.start_time, end_time=end_time, token=token) for username, timestamp, count in entries: week = (timestamp - start_time).ToInt( rdfvalue.SECONDS) // week_duration.ToInt(rdfvalue.SECONDS) if week in user_activity[username]: user_activity[username][week] += count user_activity = sorted(iteritems(user_activity)) user_activity = [(user, data) for user, data in user_activity if user not in access_control.SYSTEM_USERS] ret.stack_chart.data = [ rdf_report_plugins.ApiReportDataSeries2D( label=user, points=(rdf_report_plugins.ApiReportDataPoint2D(x=x, y=y) for x, y in sorted(data.items()))) for user, data in user_activity ] return ret
def testIndexedReads(self): spacing = 10 with utils.Stubber(sequential_collection.IndexedSequentialCollection, "INDEX_SPACING", spacing): urn = "aff4:/sequential_collection/testIndexedReads" collection = self._TestCollection(urn) data_size = 4 * spacing # TODO(amoser): Without using a mutation pool, this test is really # slow on MySQL data store. with data_store.DB.GetMutationPool() as pool: for i in range(data_size): collection.StaticAdd( rdfvalue.RDFURN(urn), rdfvalue.RDFInteger(i), mutation_pool=pool) with test_lib.FakeTime(rdfvalue.RDFDatetime.Now() + rdfvalue.Duration("10m")): for i in range(data_size - 1, data_size - 20, -1): self.assertEqual(collection[i], i) for i in [spacing - 1, spacing, spacing + 1]: self.assertEqual(collection[i], i) for i in range(data_size - spacing + 5, data_size - spacing - 5, -1): self.assertEqual(collection[i], i)
def testReturnProcessedFlow(self): client_id, flow_id = self._SetupClientAndFlow( next_request_to_process=1) processing_time = rdfvalue.Duration("60s") processed_flow = self.db.ReadFlowForProcessing(client_id, flow_id, processing_time) # Let's say we processed one request on this flow. processed_flow.next_request_to_process = 2 # There are some requests ready for processing but not #2. self.db.WriteFlowRequests([ rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=1, needs_processing=True), rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=4, needs_processing=True) ]) self.assertTrue(self.db.ReturnProcessedFlow(processed_flow)) processed_flow = self.db.ReadFlowForProcessing(client_id, flow_id, processing_time) # And another one. processed_flow.next_request_to_process = 3 # But in the meantime, request 3 is ready for processing. self.db.WriteFlowRequests([ rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=3, needs_processing=True) ]) self.assertFalse(self.db.ReturnProcessedFlow(processed_flow))
class PurgeClientStatsCronJob(cronjobs.CronJobBase): """Deletes outdated client statistics.""" frequency = rdfvalue.Duration("1w") # Keep stats for one month. MAX_AGE = 31 * 24 * 3600 def Run(self): self.start = 0 self.end = int(1e6 * (time.time() - self.MAX_AGE)) client_urns = export_utils.GetAllClients(token=self.token) for batch in utils.Grouper(client_urns, 10000): with data_store.DB.GetMutationPool() as mutation_pool: for client_urn in batch: mutation_pool.DeleteAttributes(client_urn.Add("stats"), [u"aff4:stats"], start=self.start, end=self.end) self.HeartBeat()