def UpdateFlow(self,
client_id,
flow_id,
flow_obj=db.Database.unchanged,
flow_state=db.Database.unchanged,
client_crash_info=db.Database.unchanged,
pending_termination=db.Database.unchanged,
processing_on=db.Database.unchanged,
processing_since=db.Database.unchanged,
processing_deadline=db.Database.unchanged):
"""Updates flow objects in the database."""
try:
flow = self.flows[(client_id, flow_id)]
except KeyError:
raise db.UnknownFlowError(client_id, flow_id)
if flow_obj != db.Database.unchanged:
self.flows[(client_id, flow_id)] = flow_obj
flow = flow_obj
if flow_state != db.Database.unchanged:
flow.flow_state = flow_state
if client_crash_info != db.Database.unchanged:
flow.client_crash_info = client_crash_info
if pending_termination != db.Database.unchanged:
flow.pending_termination = pending_termination
if processing_on != db.Database.unchanged:
flow.processing_on = processing_on
if processing_since != db.Database.unchanged:
flow.processing_since = processing_since
if processing_deadline != db.Database.unchanged:
flow.processing_deadline = processing_deadline
flow.last_update_time = rdfvalue.RDFDatetime.Now()
def UpdateIncrementalFlowRequests(
self, client_id: str, flow_id: str,
next_response_id_updates: Dict[int, int]) -> None:
"""Updates incremental flow requests."""
if (client_id, flow_id) not in self.flows:
raise db.UnknownFlowError(client_id, flow_id)
request_dict = self.flow_requests[(client_id, flow_id)]
for request_id, next_response_id in next_response_id_updates.items():
request_dict[request_id].next_response_id = next_response_id
request_dict[request_id].timestamp = rdfvalue.RDFDatetime.Now()
def WriteFlowLogEntry(self, entry: rdf_flow_objects.FlowLogEntry) -> None:
"""Writes a single flow log entry to the database."""
key = (entry.client_id, entry.flow_id)
if key not in self.flows:
raise db.UnknownFlowError(entry.client_id, entry.flow_id)
entry = entry.Copy()
entry.timestamp = rdfvalue.RDFDatetime.Now()
self.flow_log_entries.setdefault(key, []).append(entry)
def DeleteAllFlowRequestsAndResponses(self, client_id, flow_id):
"""Deletes all requests and responses for a given flow from the database."""
flow_key = (client_id, flow_id)
try:
self.flows[flow_key]
except KeyError:
raise db.UnknownFlowError(client_id, flow_id)
try:
del self.flow_requests[flow_key]
except KeyError:
pass
try:
del self.flow_responses[flow_key]
except KeyError:
pass
def UpdateFlow(self,
client_id,
flow_id,
flow_obj=db.Database.unchanged,
flow_state=db.Database.unchanged,
client_crash_info=db.Database.unchanged,
pending_termination=db.Database.unchanged,
processing_on=db.Database.unchanged,
processing_since=db.Database.unchanged,
processing_deadline=db.Database.unchanged):
"""Updates flow objects in the database."""
try:
flow = self.flows[(client_id, flow_id)]
except KeyError:
raise db.UnknownFlowError(client_id, flow_id)
if flow_obj != db.Database.unchanged:
new_flow = flow_obj.Copy()
# Some fields cannot be updated.
new_flow.client_id = flow.client_id
new_flow.flow_id = flow.flow_id
new_flow.long_flow_id = flow.long_flow_id
new_flow.parent_flow_id = flow.parent_flow_id
new_flow.parent_hunt_id = flow.parent_hunt_id
new_flow.flow_class_name = flow.flow_class_name
new_flow.creator = flow.creator
self.flows[(client_id, flow_id)] = new_flow
flow = new_flow
if flow_state != db.Database.unchanged:
flow.flow_state = flow_state
if client_crash_info != db.Database.unchanged:
flow.client_crash_info = client_crash_info
if pending_termination != db.Database.unchanged:
flow.pending_termination = pending_termination
if processing_on != db.Database.unchanged:
flow.processing_on = processing_on
if processing_since != db.Database.unchanged:
flow.processing_since = processing_since
if processing_deadline != db.Database.unchanged:
flow.processing_deadline = processing_deadline
flow.last_update_time = rdfvalue.RDFDatetime.Now()
def DeleteFlowRequests(self, requests):
"""Deletes a list of flow requests from the database."""
for request in requests:
if (request.client_id, request.flow_id) not in self.flows:
raise db.UnknownFlowError(request.client_id, request.flow_id)
for request in requests:
key = (request.client_id, request.flow_id)
request_dict = self.flow_requests.get(key, {})
try:
del request_dict[request.request_id]
except KeyError:
raise db.UnknownFlowRequestError(request.client_id, request.flow_id,
request.request_id)
response_dict = self.flow_responses.get(key, {})
try:
del response_dict[request.request_id]
except KeyError:
pass
def ReadFlowObject(self, client_id, flow_id):
"""Reads a flow object from the database."""
try:
return self.flows[(client_id, flow_id)].Copy()
except KeyError:
raise db.UnknownFlowError(client_id, flow_id)
