def testHuntIdToInt(self): self.assertEqual(db_utils.HuntIDToInt("00000001"), 1) self.assertEqual(db_utils.HuntIDToInt("1234ABCD"), 0x1234ABCD) self.assertEqual(db_utils.HuntIDToInt("FFFFFFFF"), 0xFFFFFFFF) self.assertEqual(db_utils.HuntIDToInt("0000000100000000"), 0x100000000) self.assertEqual(db_utils.HuntIDToInt("FFFFFFFFFFFFFFFF"), 0xFFFFFFFFFFFFFFFF)
def UpdateHuntObject(self, hunt_id, duration=None, client_rate=None, client_limit=None, hunt_state=None, hunt_state_comment=None, start_time=None, num_clients_at_start_time=None, cursor=None): """Updates the hunt object by applying the update function.""" vals = [] args = {} if duration is not None: vals.append("duration_micros = %(duration_micros)s") args["duration_micros"] = duration.microseconds if client_rate is not None: vals.append("client_rate = %(client_rate)s") args["client_rate"] = client_rate if client_limit is not None: vals.append("client_limit = %(client_limit)s") args["client_limit"] = client_limit if hunt_state is not None: vals.append("hunt_state = %(hunt_state)s") args["hunt_state"] = int(hunt_state) if hunt_state_comment is not None: vals.append("hunt_state_comment = %(hunt_state_comment)s") args["hunt_state_comment"] = hunt_state_comment if start_time is not None: vals.append(""" init_start_time = IFNULL(init_start_time, FROM_UNIXTIME(%(start_time)s)) """) vals.append(""" last_start_time = FROM_UNIXTIME(%(start_time)s) """) args["start_time"] = mysql_utils.RDFDatetimeToTimestamp(start_time) if num_clients_at_start_time is not None: vals.append( "num_clients_at_start_time = %(num_clients_at_start_time)s") args["num_clients_at_start_time"] = num_clients_at_start_time vals.append("last_update_timestamp = NOW(6)") query = """ UPDATE hunts SET {updates} WHERE hunt_id = %(hunt_id)s """.format(updates=", ".join(vals)) args["hunt_id"] = db_utils.HuntIDToInt(hunt_id) rows_modified = cursor.execute(query, args) if rows_modified == 0: raise db.UnknownHuntError(hunt_id)
def WriteHuntOutputPluginsStates(self, hunt_id, states, cursor=None): """Writes hunt output plugin states for a given hunt.""" columns = ", ".join(_HUNT_OUTPUT_PLUGINS_STATES_COLUMNS) placeholders = mysql_utils.Placeholders( 2 + len(_HUNT_OUTPUT_PLUGINS_STATES_COLUMNS)) hunt_id_int = db_utils.HuntIDToInt(hunt_id) for index, state in enumerate(states): query = ("INSERT INTO hunt_output_plugins_states " "(hunt_id, plugin_id, {columns}) " "VALUES {placeholders}".format(columns=columns, placeholders=placeholders)) args = [hunt_id_int, index, state.plugin_descriptor.plugin_name] if state.plugin_descriptor.plugin_args is None: args.append(None) else: args.append( state.plugin_descriptor.plugin_args.SerializeToBytes()) args.append(state.plugin_state.SerializeToBytes()) try: cursor.execute(query, args) except MySQLdb.IntegrityError as e: raise db.UnknownHuntError(hunt_id=hunt_id, cause=e)
def WriteHuntObject(self, hunt_obj, cursor=None): """Writes a hunt object to the database.""" query = """ INSERT INTO hunts (hunt_id, creator, description, duration_micros, hunt_state, client_rate, client_limit, hunt) VALUES (%(hunt_id)s, %(creator)s, %(description)s, %(duration_micros)s, %(hunt_state)s, %(client_rate)s, %(client_limit)s, %(hunt)s) """ args = { "hunt_id": db_utils.HuntIDToInt(hunt_obj.hunt_id), "creator": hunt_obj.creator, "description": hunt_obj.description, "duration_micros": hunt_obj.duration.microseconds, "hunt_state": int(rdf_hunt_objects.Hunt.HuntState.PAUSED), "client_rate": hunt_obj.client_rate, "client_limit": hunt_obj.client_limit, "hunt": hunt_obj.SerializeToBytes(), } try: cursor.execute(query, args) except MySQLdb.IntegrityError as error: raise db.DuplicatedHuntError(hunt_id=hunt_obj.hunt_id, cause=error)
def UpdateHuntOutputPluginState(self, hunt_id, state_index, update_fn, cursor=None): """Updates hunt output plugin state for a given output plugin.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = "SELECT hunt_id FROM hunts WHERE hunt_id = %s" rows_returned = cursor.execute(query, [hunt_id_int]) if rows_returned == 0: raise db.UnknownHuntError(hunt_id) columns = ", ".join(_HUNT_OUTPUT_PLUGINS_STATES_COLUMNS) query = ("SELECT {columns} FROM hunt_output_plugins_states " "WHERE hunt_id = %s AND plugin_id = %s".format( columns=columns)) rows_returned = cursor.execute(query, [hunt_id_int, state_index]) if rows_returned == 0: raise db.UnknownHuntOutputPluginStateError(hunt_id, state_index) state = self._HuntOutputPluginStateFromRow(cursor.fetchone()) modified_plugin_state = update_fn(state.plugin_state) query = ("UPDATE hunt_output_plugins_states " "SET plugin_state = %s " "WHERE hunt_id = %s AND plugin_id = %s") args = [ modified_plugin_state.SerializeToBytes(), hunt_id_int, state_index ] cursor.execute(query, args) return state
def ReadHuntCounters(self, hunt_id, cursor=None): """Reads hunt counters.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT flow_state, COUNT(*) " "FROM flows " "FORCE INDEX(flows_by_hunt) " "WHERE parent_hunt_id = %s AND parent_flow_id IS NULL " "GROUP BY flow_state") cursor.execute(query, [hunt_id_int]) counts_by_state = dict(cursor.fetchall()) num_successful_clients = counts_by_state.get( int(rdf_flow_objects.Flow.FlowState.FINISHED), 0) num_failed_clients = counts_by_state.get( int(rdf_flow_objects.Flow.FlowState.ERROR), 0) num_crashed_clients = counts_by_state.get( int(rdf_flow_objects.Flow.FlowState.CRASHED), 0) num_running_clients = counts_by_state.get( int(rdf_flow_objects.Flow.FlowState.RUNNING), 0) num_clients = sum(counts_by_state.values()) query = """ SELECT * FROM ( SELECT COUNT(client_id) FROM flows FORCE INDEX(flows_by_hunt) WHERE parent_hunt_id = %s AND parent_flow_id IS NULL AND num_replies_sent > 0) counters, ( SELECT SUM(user_cpu_time_used_micros + system_cpu_time_used_micros), SUM(network_bytes_sent), SUM(num_replies_sent) FROM flows FORCE INDEX(flows_by_hunt) WHERE parent_hunt_id = %s AND parent_flow_id IS NULL) resources """ cursor.execute(query, [hunt_id_int, hunt_id_int]) ( num_clients_with_results, total_cpu_seconds, total_network_bytes_sent, num_results, ) = cursor.fetchone() return db.HuntCounters( num_clients=num_clients, num_successful_clients=num_successful_clients, num_failed_clients=num_failed_clients, num_clients_with_results=num_clients_with_results, num_crashed_clients=num_crashed_clients, num_running_clients=num_running_clients, num_results=int(num_results or 0), total_cpu_seconds=db_utils.MicrosToSeconds(int(total_cpu_seconds or 0)), total_network_bytes_sent=int(total_network_bytes_sent or 0))
def CountHuntResultsByType(self, hunt_id, cursor=None): """Counts number of hunts results per type.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT type, COUNT(*) FROM flow_results " "WHERE hunt_id = %s GROUP BY type") cursor.execute(query, [hunt_id_int]) return dict(cursor.fetchall())
def CountHuntLogEntries(self, hunt_id, cursor=None): """Returns number of hunt log entries of a given hunt.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT COUNT(*) FROM flow_log_entries " "FORCE INDEX(flow_log_entries_by_hunt) " "WHERE hunt_id = %s AND flow_id = hunt_id") cursor.execute(query, [hunt_id_int]) return cursor.fetchone()[0]
def ReadHuntOutputPluginsStates(self, hunt_id, cursor=None): """Reads all hunt output plugins states of a given hunt.""" columns = ", ".join(_HUNT_OUTPUT_PLUGINS_STATES_COLUMNS) query = ("SELECT {columns} FROM hunt_output_plugins_states " "WHERE hunt_id = %s".format(columns=columns)) rows_returned = cursor.execute(query, [db_utils.HuntIDToInt(hunt_id)]) if rows_returned > 0: states = [] for row in cursor.fetchall(): states.append(self._HuntOutputPluginStateFromRow(row)) return states query = "SELECT hunt_id FROM hunts WHERE hunt_id = %s" rows_returned = cursor.execute(query, [db_utils.HuntIDToInt(hunt_id)]) if rows_returned == 0: raise db.UnknownHuntError(hunt_id) return []
def ReadHuntObject(self, hunt_id, cursor=None): """Reads a hunt object from the database.""" query = ("SELECT {columns} " "FROM hunts WHERE hunt_id = %s".format( columns=_HUNT_COLUMNS_SELECT)) nr_results = cursor.execute(query, [db_utils.HuntIDToInt(hunt_id)]) if nr_results == 0: raise db.UnknownHuntError(hunt_id) return self._HuntObjectFromRow(cursor.fetchone())
def DeleteHuntObject(self, hunt_id, cursor=None): """Deletes a given hunt object.""" query = "DELETE FROM hunts WHERE hunt_id = %s" hunt_id_int = db_utils.HuntIDToInt(hunt_id) rows_deleted = cursor.execute(query, [hunt_id_int]) if rows_deleted == 0: raise db.UnknownHuntError(hunt_id) query = "DELETE FROM hunt_output_plugins_states WHERE hunt_id = %s" cursor.execute(query, [hunt_id_int])
def CountHuntFlows(self, hunt_id, filter_condition=db.HuntFlowsCondition.UNSET, cursor=None): """Counts hunt flows matching given conditions.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT COUNT(*) FROM flows " "FORCE INDEX(flows_by_hunt) " "WHERE parent_hunt_id = %s AND parent_flow_id IS NULL " "{filter_condition}") filter_query, extra_args = self._HuntFlowCondition(filter_condition) args = [hunt_id_int] + extra_args query = query.format(filter_condition=filter_query) cursor.execute(query, args) return cursor.fetchone()[0]
def ReadHuntOutputPluginLogEntries(self, hunt_id, output_plugin_id, offset, count, with_type=None, cursor=None): """Reads hunt output plugin log entries.""" query = ("SELECT client_id, flow_id, log_entry_type, message, " "UNIX_TIMESTAMP(timestamp) " "FROM flow_output_plugin_log_entries " "FORCE INDEX (flow_output_plugin_log_entries_by_hunt) " "WHERE hunt_id = %s AND output_plugin_id = %s ") args = [ db_utils.HuntIDToInt(hunt_id), db_utils.OutputPluginIDToInt(output_plugin_id) ] if with_type is not None: query += "AND log_entry_type = %s " args.append(int(with_type)) query += "ORDER BY log_id ASC LIMIT %s OFFSET %s" args.append(count) args.append(offset) cursor.execute(query, args) ret = [] for (client_id_int, flow_id_int, log_entry_type, message, timestamp) in cursor.fetchall(): ret.append( rdf_flow_objects.FlowOutputPluginLogEntry( hunt_id=hunt_id, client_id=db_utils.IntToClientID(client_id_int), flow_id=db_utils.IntToFlowID(flow_id_int), output_plugin_id=output_plugin_id, log_entry_type=log_entry_type, message=message, timestamp=mysql_utils.TimestampToRDFDatetime(timestamp))) return ret
def CountHuntOutputPluginLogEntries(self, hunt_id, output_plugin_id, with_type=None, cursor=None): """Counts hunt output plugin log entries.""" query = ("SELECT COUNT(*) " "FROM flow_output_plugin_log_entries " "FORCE INDEX (flow_output_plugin_log_entries_by_hunt) " "WHERE hunt_id = %s AND output_plugin_id = %s ") args = [ db_utils.HuntIDToInt(hunt_id), db_utils.OutputPluginIDToInt(output_plugin_id) ] if with_type is not None: query += "AND log_entry_type = %s" args.append(int(with_type)) cursor.execute(query, args) return cursor.fetchone()[0]
def ReadHuntLogEntries(self, hunt_id, offset, count, with_substring=None, cursor=None): """Reads hunt log entries of a given hunt using given query options.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ( "SELECT client_id, flow_id, message, UNIX_TIMESTAMP(timestamp) " "FROM flow_log_entries " "FORCE INDEX(flow_log_entries_by_hunt) " "WHERE hunt_id = %s AND flow_id = hunt_id ") args = [hunt_id_int] if with_substring is not None: query += "AND message LIKE %s " args.append("%" + db_utils.EscapeWildcards(with_substring) + "%") query += "ORDER BY timestamp ASC LIMIT %s OFFSET %s" args.append(count) args.append(offset) cursor.execute(query, args) flow_log_entries = [] for client_id_int, flow_id_int, message, timestamp in cursor.fetchall( ): flow_log_entries.append( rdf_flow_objects.FlowLogEntry( client_id=db_utils.IntToClientID(client_id_int), flow_id=db_utils.IntToFlowID(flow_id_int), hunt_id=hunt_id, message=message, timestamp=mysql_utils.TimestampToRDFDatetime(timestamp))) return flow_log_entries
def ReadHuntFlowsStatesAndTimestamps(self, hunt_id, cursor=None): """Reads hunt flows states and timestamps.""" query = """ SELECT flow_state, UNIX_TIMESTAMP(timestamp), UNIX_TIMESTAMP(last_update) FROM flows FORCE INDEX(flows_by_hunt) WHERE parent_hunt_id = %s AND parent_flow_id IS NULL """ cursor.execute(query, [db_utils.HuntIDToInt(hunt_id)]) result = [] for fs, ct, lup in cursor.fetchall(): result.append( db.FlowStateAndTimestamps( flow_state=rdf_flow_objects.Flow.FlowState.FromInt(fs), create_time=mysql_utils.TimestampToRDFDatetime(ct), last_update_time=mysql_utils.TimestampToRDFDatetime(lup))) return result
def CountHuntResults(self, hunt_id, with_tag=None, with_type=None, cursor=None): """Counts hunt results of a given hunt using given query options.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = "SELECT COUNT(*) FROM flow_results WHERE hunt_id = %s " args = [hunt_id_int] if with_tag is not None: query += "AND tag = %s " args.append(with_tag) if with_type is not None: query += "AND type = %s " args.append(with_type) cursor.execute(query, args) return cursor.fetchone()[0]
def ReadHuntFlows(self, hunt_id, offset, count, filter_condition=db.HuntFlowsCondition.UNSET, cursor=None): """Reads hunt flows matching given conditins.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT {columns} FROM flows " "FORCE INDEX(flows_by_hunt) " "WHERE parent_hunt_id = %s AND parent_flow_id IS NULL " "{filter_condition} " "ORDER BY last_update ASC " "LIMIT %s OFFSET %s") filter_query, extra_args = self._HuntFlowCondition(filter_condition) query = query.format(columns=self.FLOW_DB_FIELDS, filter_condition=filter_query) args = [hunt_id_int] + extra_args + [count, offset] cursor.execute(query, args) return [self._FlowObjectFromRow(row) for row in cursor.fetchall()]
def ReadHuntResults(self, hunt_id, offset, count, with_tag=None, with_type=None, with_substring=None, with_timestamp=None, cursor=None): """Reads hunt results of a given hunt using given query options.""" hunt_id_int = db_utils.HuntIDToInt(hunt_id) query = ("SELECT client_id, flow_id, hunt_id, payload, type, " "UNIX_TIMESTAMP(timestamp), tag " "FROM flow_results " "FORCE INDEX(flow_results_hunt_id_flow_id_timestamp) " "WHERE hunt_id = %s ") args = [hunt_id_int] if with_tag: query += "AND tag = %s " args.append(with_tag) if with_type: query += "AND type = %s " args.append(with_type) if with_substring: query += "AND payload LIKE %s " args.append("%" + db_utils.EscapeWildcards(with_substring) + "%") if with_timestamp: query += "AND timestamp = FROM_UNIXTIME(%s) " args.append(mysql_utils.RDFDatetimeToTimestamp(with_timestamp)) query += "ORDER BY timestamp ASC LIMIT %s OFFSET %s" args.append(count) args.append(offset) cursor.execute(query, args) ret = [] for ( client_id_int, flow_id_int, hunt_id_int, serialized_payload, payload_type, timestamp, tag, ) in cursor.fetchall(): if payload_type in rdfvalue.RDFValue.classes: payload = rdfvalue.RDFValue.classes[ payload_type].FromSerializedBytes(serialized_payload) else: payload = rdf_objects.SerializedValueOfUnrecognizedType( type_name=payload_type, value=serialized_payload) result = rdf_flow_objects.FlowResult( client_id=db_utils.IntToClientID(client_id_int), flow_id=db_utils.IntToFlowID(flow_id_int), hunt_id=hunt_id, payload=payload, timestamp=mysql_utils.TimestampToRDFDatetime(timestamp)) if tag is not None: result.tag = tag ret.append(result) return ret
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