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)