def _FlowStatusToClientResources(flow_obj, status_msg): return rdf_client_stats.ClientResources( client_id=flow_obj.client_id, session_id=flow_obj.flow_id, cpu_usage=rdf_client_stats.CpuSeconds( user_cpu_time=status_msg.cpu_time_used.user_cpu_time, system_cpu_time=status_msg.cpu_time_used.system_cpu_time), network_bytes_sent=status_msg.network_bytes_sent)
def ReadHuntClientResourcesStats(self, hunt_id): """Read/calculate hunt client resources stats.""" result = rdf_stats.ClientResourcesStats() for f in self._GetHuntFlows(hunt_id): cr = rdf_client_stats.ClientResources( session_id="%s/%s" % (f.client_id, f.flow_id), client_id=f.client_id, cpu_usage=f.cpu_time_used, network_bytes_sent=f.network_bytes_sent) result.RegisterResources(cr) # TODO(user): remove this hack when compatibility with AFF4 is not # important. return rdf_stats.ClientResourcesStats.FromSerializedString( result.SerializeToString())
def ProcessHuntFlowDone(flow_obj, status_msg=None): """Notifis hunt about a given hunt-induced flow completion.""" if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id): resources = rdf_client_stats.ClientResources( client_id=flow_obj.client_id, session_id=flow_obj.flow_id, cpu_usage=rdf_client_stats.CpuSeconds( user_cpu_time=status_msg.cpu_time_used.user_cpu_time, system_cpu_time=status_msg.cpu_time_used.system_cpu_time), network_bytes_sent=status_msg.network_bytes_sent) def UpdateFn(hunt_obj): hunt_obj.num_successful_clients += 1 if flow_obj.num_replies_sent: hunt_obj.num_clients_with_results += 1 hunt_obj.client_resources_stats.RegisterResources(resources) return hunt_obj hunt_obj = data_store.REL_DB.UpdateHuntObject(flow_obj.parent_hunt_id, UpdateFn) hunt_obj = hunt.StopHuntIfAverageLimitsExceeded(hunt_obj) hunt.CompleteHuntIfExpirationTimeReached(hunt_obj) return hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id) client_urn = rdf_client.ClientURN(flow_obj.client_id) with aff4.FACTORY.OpenWithLock(hunt_urn, lease_time=_HUNT_LEASE_TIME, blocking=True) as fd: # Legacy AFF4 code expects token to be set. fd.token = access_control.ACLToken(username=fd.creator) fd.RegisterCompletedClient(client_urn) if flow_obj.num_replies_sent: fd.RegisterClientWithResults(client_urn) fd.context.clients_with_results_count += 1 fd.context.completed_clients_count += 1 fd.context.results_count += flow_obj.num_replies_sent fd.GetRunner().SaveResourceUsage(flow_obj.client_id, status_msg) fd.StopHuntIfAverageLimitsExceeded()
def InitFromFlowObject(self, flow_obj, with_args=True, with_progress=False, with_state_and_context=False): try: self.flow_id = flow_obj.flow_id self.client_id = flow_obj.client_id # TODO(amoser): Get rid of all urns. self.urn = flow_obj.long_flow_id self.name = flow_obj.flow_class_name self.started_at = flow_obj.create_time self.last_active_at = flow_obj.last_update_time self.creator = flow_obj.creator if flow_obj.client_crash_info: self.state = "CLIENT_CRASHED" elif flow_obj.pending_termination: self.state = "ERROR" self.status = ("Pending termination: %s" % flow_obj.pending_termination.reason) else: context_state_map = {1: "RUNNING", 2: "TERMINATED", 3: "ERROR"} self.state = context_state_map[int(flow_obj.flow_state)] if with_state_and_context: outstanding_requests = (flow_obj.next_outbound_id - flow_obj.next_request_to_process) self.context = rdf_flow_runner.FlowContext( # TODO(amoser): No need to set this in all cases once the legacy API # is removed. client_resources=rdf_client_stats.ClientResources( cpu_usage=rdf_client_stats.CpuSeconds()), create_time=flow_obj.create_time, creator=flow_obj.creator, current_state=flow_obj.current_state, next_outbound_id=flow_obj.next_outbound_id, outstanding_requests=outstanding_requests, state=self.state, # TODO(amoser): Get rid of all urns. session_id=flow_obj.long_flow_id, ) if flow_obj.output_plugins_states: self.context.output_plugins_states = flow_obj.output_plugins_states if flow_obj.network_bytes_sent: self.context.network_bytes_sent = flow_obj.network_bytes_sent self.context.client_resources.network_bytes_sent = ( flow_obj.network_bytes_sent) if flow_obj.cpu_time_used: self.context.client_resources.cpu_time_used = flow_obj.cpu_time_used if flow_obj.error_message: self.context.status = flow_obj.error_message if flow_obj.backtrace: self.context.backtrace = flow_obj.backtrace if with_args: try: self.args = flow_obj.args except ValueError: # If args class name has changed, ValueError will be raised. Handling # this gracefully - we should still try to display some useful info # about the flow. pass if with_progress: flow_cls = self._GetFlowClass() if flow_cls: self.progress = flow_cls(flow_obj).GetProgress() self.runner_args = rdf_flow_runner.FlowRunnerArgs( client_id=flow_obj.client_id, flow_name=flow_obj.flow_class_name, notify_to_user=flow_base.FlowBase( flow_obj).ShouldSendNotifications()) if flow_obj.output_plugins: self.runner_args.output_plugins = flow_obj.output_plugins if flow_obj.HasField("cpu_limit"): self.runner_args.cpu_limit = flow_obj.cpu_limit if flow_obj.HasField("network_bytes_limit"): self.runner_args.cpu_limit = flow_obj.network_bytes_limit if flow_obj.original_flow.flow_id: self.original_flow = ApiFlowReference().FromFlowReference( flow_obj.original_flow) if with_state_and_context and flow_obj.persistent_data.ToDict(): self.state_data = ( api_call_handler_utils.ApiDataObject().InitFromDataObject( flow_obj.persistent_data)) except Exception as e: # pylint: disable=broad-except self.internal_error = "Error while opening flow: %s" % str(e) return self
def ReadHuntClientResourcesStats(self, hunt_id, cursor=None): """Read/calculate hunt client resources stats.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = """ SELECT COUNT(*), SUM(user_cpu_time_used_micros), SUM((user_cpu_time_used_micros) * (user_cpu_time_used_micros)), SUM(system_cpu_time_used_micros), SUM((system_cpu_time_used_micros) * (system_cpu_time_used_micros)), SUM(network_bytes_sent), SUM(network_bytes_sent * network_bytes_sent), """ scaled_bins = [ int(1000000 * b) for b in rdf_stats.ClientResourcesStats.CPU_STATS_BINS ] query += self._BinsToQuery(scaled_bins, "(user_cpu_time_used_micros)") query += "," query += self._BinsToQuery(scaled_bins, "(system_cpu_time_used_micros)") query += "," query += self._BinsToQuery( rdf_stats.ClientResourcesStats.NETWORK_STATS_BINS, "network_bytes_sent") query += " FROM flows " query += "FORCE INDEX(flows_by_hunt) " query += "WHERE parent_hunt_id = %s AND parent_flow_id IS NULL" cursor.execute(query, [hunt_id_int]) response = cursor.fetchone() (count, user_sum, user_sq_sum, system_sum, system_sq_sum, network_sum, network_sq_sum) = response[:7] stats = rdf_stats.ClientResourcesStats( user_cpu_stats=rdf_stats.RunningStats( num=count, sum=db_utils.MicrosToSeconds(int(user_sum or 0)), sum_sq=int(user_sq_sum or 0) / 1e12, ), system_cpu_stats=rdf_stats.RunningStats( num=count, sum=db_utils.MicrosToSeconds(int(system_sum or 0)), sum_sq=int(system_sq_sum or 0) / 1e12, ), network_bytes_sent_stats=rdf_stats.RunningStats( num=count, sum=float(network_sum or 0), sum_sq=float(network_sq_sum or 0), ), ) offset = 7 stats.user_cpu_stats.histogram = rdf_stats.StatsHistogram() for b_num, b_max_value in zip( response[offset:], rdf_stats.ClientResourcesStats.CPU_STATS_BINS): stats.user_cpu_stats.histogram.bins.append( rdf_stats.StatsHistogramBin(range_max_value=b_max_value, num=b_num)) offset += len(rdf_stats.ClientResourcesStats.CPU_STATS_BINS) stats.system_cpu_stats.histogram = rdf_stats.StatsHistogram() for b_num, b_max_value in zip( response[offset:], rdf_stats.ClientResourcesStats.CPU_STATS_BINS): stats.system_cpu_stats.histogram.bins.append( rdf_stats.StatsHistogramBin(range_max_value=b_max_value, num=b_num)) offset += len(rdf_stats.ClientResourcesStats.CPU_STATS_BINS) stats.network_bytes_sent_stats.histogram = rdf_stats.StatsHistogram() for b_num, b_max_value in zip( response[offset:], rdf_stats.ClientResourcesStats.NETWORK_STATS_BINS): stats.network_bytes_sent_stats.histogram.bins.append( rdf_stats.StatsHistogramBin(range_max_value=b_max_value, num=b_num)) query = """ SELECT client_id, flow_id, user_cpu_time_used_micros, system_cpu_time_used_micros, network_bytes_sent FROM flows FORCE INDEX(flows_by_hunt) WHERE parent_hunt_id = %s AND parent_flow_id IS NULL AND (user_cpu_time_used_micros > 0 OR system_cpu_time_used_micros > 0 OR network_bytes_sent > 0) ORDER BY (user_cpu_time_used_micros + system_cpu_time_used_micros) DESC LIMIT 10 """ cursor.execute(query, [hunt_id_int]) for cid, fid, ucpu, scpu, nbs in cursor.fetchall(): client_id = db_utils.IntToClientID(cid) flow_id = db_utils.IntToFlowID(fid) stats.worst_performers.append( rdf_client_stats.ClientResources( client_id=client_id, session_id=rdfvalue.RDFURN(client_id).Add(flow_id), cpu_usage=rdf_client_stats.CpuSeconds( user_cpu_time=db_utils.MicrosToSeconds(ucpu), system_cpu_time=db_utils.MicrosToSeconds(scpu), ), network_bytes_sent=nbs)) return stats