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
