def WriteClientMessages(self, messages):
"""Writes messages that should go to the client to the db."""
client_ids = [db_utils.ClientIdFromGrrMessage(msg) for msg in messages]
for client_id in client_ids:
if client_id not in self.metadatas:
raise db.AtLeastOneUnknownClientError(client_ids=client_ids)
for m in messages:
client_id = db_utils.ClientIdFromGrrMessage(m)
self.client_messages.setdefault(client_id, {})[m.task_id] = m
def WriteClientMessages(self, messages, cursor=None):
"""Writes messages that should go to the client to the db."""
query = ("INSERT IGNORE INTO client_messages "
"(client_id, message_id, timestamp, message) "
"VALUES %s ON DUPLICATE KEY UPDATE "
"timestamp=VALUES(timestamp), message=VALUES(message)")
now = mysql_utils.RDFDatetimeToMysqlString(rdfvalue.RDFDatetime.Now())
client_ids = set()
value_templates = []
args = []
for m in messages:
cid = db_utils.ClientIdFromGrrMessage(m)
client_ids.add(cid)
client_id_int = mysql_utils.ClientIDToInt(cid)
args.extend([client_id_int, m.task_id, now, m.SerializeToString()])
value_templates.append("(%s, %s, %s, %s)")
query %= ",".join(value_templates)
try:
cursor.execute(query, args)
except MySQLdb.IntegrityError as e:
raise db.AtLeastOneUnknownClientError(client_ids=client_ids, cause=e)
def LeaseClientMessages(self,
client_id,
lease_time=None,
limit=sys.maxsize):
"""Leases available client messages for the client with the given id."""
leased_messages = []
now = rdfvalue.RDFDatetime.Now()
expiration_time = now + lease_time
process_id_str = utils.ProcessIdString()
leases = self.client_message_leases
for msgs_by_id in itervalues(self.client_messages):
for msg in sorted(itervalues(msgs_by_id), key=lambda m: m.task_id):
if db_utils.ClientIdFromGrrMessage(msg) != client_id:
continue
existing_lease = leases.get(msg.task_id)
if not existing_lease or existing_lease[0] < now:
leases[msg.task_id] = (expiration_time, process_id_str)
msg.leased_until = expiration_time
msg.leased_by = process_id_str
leased_messages.append(msg)
if len(leased_messages) >= limit:
break
return leased_messages
def LeaseClientMessages(self, client_id, lease_time=None, limit=sys.maxsize):
"""Leases available client messages for the client with the given id."""
leased_messages = []
now = rdfvalue.RDFDatetime.Now()
expiration_time = now + lease_time
process_id_str = utils.ProcessIdString()
leases = self.client_message_leases
for msgs_by_id in itervalues(self.client_messages):
for msg in sorted(itervalues(msgs_by_id), key=lambda m: m.task_id):
if db_utils.ClientIdFromGrrMessage(msg) != client_id:
continue
existing_lease = leases.get(msg.task_id)
if not existing_lease or existing_lease[0] < now:
if existing_lease:
lease_count = existing_lease[-1] + 1
# >= comparison since this check happens before the lease.
if lease_count >= db.Database.CLIENT_MESSAGES_TTL:
self._DeleteClientMessage(client_id, msg.task_id)
continue
leases[msg.task_id] = (expiration_time, process_id_str, lease_count)
else:
leases[msg.task_id] = (expiration_time, process_id_str, 0)
msg.leased_until = expiration_time
msg.leased_by = process_id_str
leased_messages.append(msg)
if len(leased_messages) >= limit:
break
return leased_messages
def DeleteClientMessages(self, messages):
"""Deletes a list of client messages from the db."""
to_delete = []
for m in messages:
client_id = db_utils.ClientIdFromGrrMessage(m)
to_delete.append((client_id, m.task_id))
if len(set(to_delete)) != len(to_delete):
raise ValueError(
"Received multiple copies of the same message to delete.")
for client_id, task_id in to_delete:
self._DeleteClientMessage(client_id, task_id)
def ReadClientMessages(self, client_id):
"""Reads all client messages available for a given client_id."""
res = []
for msgs_by_id in itervalues(self.client_messages):
for orig_msg in sorted(itervalues(msgs_by_id), key=lambda m: m.task_id):
if db_utils.ClientIdFromGrrMessage(orig_msg) != client_id:
continue
msg = orig_msg.Copy()
current_lease = self.client_message_leases.get(msg.task_id)
if current_lease:
msg.leased_until, msg.leased_by = current_lease
res.append(msg)
return res
def DeleteClientMessages(self, messages, cursor=None):
"""Deletes a list of client messages from the db."""
if not messages:
return
to_delete = []
for m in messages:
to_delete.append((db_utils.ClientIdFromGrrMessage(m), m.task_id))
if len(set(to_delete)) != len(to_delete):
raise ValueError(
"Received multiple copies of the same message to delete.")
self._DeleteClientMessages(to_delete, cursor=cursor)
def DeleteClientMessages(self, messages):
"""Deletes a list of client messages from the db."""
to_delete = []
for m in messages:
client_id = db_utils.ClientIdFromGrrMessage(m)
to_delete.append((client_id, m.task_id))
if len(set(to_delete)) != len(to_delete):
raise ValueError(
"Received multiple copies of the same message to delete.")
for client_id, task_id in to_delete:
tasks = self.client_messages.get(client_id)
if not tasks or task_id not in tasks:
# TODO(amoser): Once new flows are in, reevaluate if we can raise on
# deletion request for unknown messages.
continue
del tasks[task_id]
if task_id in self.client_message_leases:
del self.client_message_leases[task_id]
def DeleteClientMessages(self, messages, cursor=None):
"""Deletes a list of client messages from the db."""
if not messages:
return
args = []
conditions = ["(client_id=%s and message_id=%s)"] * len(messages)
query = "DELETE FROM client_messages WHERE " + " OR ".join(conditions)
to_delete = []
for m in messages:
client_id = mysql_utils.ClientIDToInt(db_utils.ClientIdFromGrrMessage(m))
to_delete.append((client_id, m.task_id))
if len(set(to_delete)) != len(to_delete):
raise ValueError(
"Received multiple copies of the same message to delete.")
for client_id, task_id in to_delete:
args.append(client_id)
args.append(task_id)
cursor.execute(query, args)
def WriteClientMessages(self, messages):
"""Writes messages that should go to the client to the db."""
for m in messages:
client_id = db_utils.ClientIdFromGrrMessage(m)
self.client_messages.setdefault(client_id, {})[m.task_id] = m