def testResponsesForUnknownRequest(self):
client_id, flow_id = self._SetupClientAndFlow()
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1)
self.db.WriteFlowRequests([request])
# Write two responses at a time, one request exists, the other doesn't.
with test_lib.SuppressLogs():
self.db.WriteFlowResponses([
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1),
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=2,
response_id=1)
])
# We should have one response in the db.
read = self.db.ReadAllFlowRequestsAndResponses(client_id, flow_id)
self.assertEqual(len(read), 1)
request, responses = read[0]
self.assertEqual(len(responses), 1)
def testDeleteFlowRequests(self):
client_id, flow_id = self._SetupClientAndFlow()
requests = []
responses = []
for request_id in range(1, 4):
requests.append(
rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=request_id))
responses.append(
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=request_id,
response_id=1))
self.db.WriteFlowRequests(requests)
self.db.WriteFlowResponses(responses)
request_list = self.db.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertItemsEqual([req.request_id for req, _ in request_list],
[req.request_id for req in requests])
random.shuffle(requests)
while requests:
request = requests.pop()
self.db.DeleteFlowRequests([request])
request_list = self.db.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertItemsEqual([req.request_id for req, _ in request_list],
[req.request_id for req in requests])
def SendReply(self, response, tag=None):
"""Allows this flow to send a message to its parent flow.
If this flow does not have a parent, the message is ignored.
Args:
response: An RDFValue() instance to be sent to the parent.
tag: If specified, tag the result with this tag.
Raises:
ValueError: If responses is not of the correct type.
"""
if not isinstance(response, rdfvalue.RDFValue):
raise ValueError("SendReply can only send RDFValues")
if self.rdf_flow.parent_flow_id:
response = rdf_flow_objects.FlowResponse(
client_id=self.rdf_flow.client_id,
request_id=self.rdf_flow.parent_request_id,
response_id=self.GetNextResponseId(),
payload=response,
flow_id=self.rdf_flow.parent_flow_id,
tag=tag)
self.flow_responses.append(response)
else:
reply = rdf_flow_objects.FlowResult(payload=response, tag=tag)
self.replies_to_write.append(reply)
self.replies_to_process.append(reply)
self.rdf_flow.num_replies_sent += 1
def _StartFlow(self, client_id, flow_cls, **kw):
if data_store.RelationalDBEnabled():
flow_id = flow.StartFlow(flow_cls=flow_cls,
client_id=client_id,
**kw)
# Lease the client message.
data_store.REL_DB.LeaseClientActionRequests(
client_id, lease_time=rdfvalue.Duration("10000s"))
# Write some responses. In the relational db, the client queue will be
# cleaned up as soon as all responses are available. Therefore we cheat
# here and make it look like the request needs more responses so it's not
# considered complete.
# Write the status first. This will mark the request as waiting for 2
# responses.
status = rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=2)
data_store.REL_DB.WriteFlowResponses([status])
# Now we read the request, adjust the number, and write it back.
reqs = data_store.REL_DB.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
req = reqs[0][0]
req.nr_responses_expected = 99
data_store.REL_DB.WriteFlowRequests([req])
# This response now won't trigger any deletion of client messages.
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1,
payload=rdf_client.Process(name="test_process"))
data_store.REL_DB.WriteFlowResponses([response])
# This is not strictly needed as we don't display this information in the
# UI.
req.nr_responses_expected = 2
data_store.REL_DB.WriteFlowRequests([req])
return flow_id
else:
flow_id = flow.StartAFF4Flow(
flow_name=compatibility.GetName(flow_cls),
client_id=client_id,
token=self.token,
**kw).Basename()
# Have the client write some responses.
test_process = rdf_client.Process(name="test_process")
mock = flow_test_lib.MockClient(client_id,
action_mocks.ListProcessesMock(
[test_process]),
token=self.token)
mock.Next()
return flow_id
def testResponseWriting(self):
client_id, flow_id = self._SetupClientAndFlow()
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1,
needs_processing=False)
self.db.WriteFlowRequests([request])
responses = [
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=i) for i in range(3)
]
self.db.WriteFlowResponses(responses)
all_requests = self.db.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertEqual(len(all_requests), 1)
read_request, read_responses = all_requests[0]
self.assertEqual(read_request, request)
self.assertEqual(list(read_responses), [0, 1, 2])
for response_id, response in iteritems(read_responses):
self.assertEqual(response.response_id, response_id)
def SendReply(self,
response: rdfvalue.RDFValue,
tag: Optional[str] = None) -> None:
"""Allows this flow to send a message to its parent flow.
If this flow does not have a parent, the message is saved to the database
as flow result.
Args:
response: An RDFValue() instance to be sent to the parent.
tag: If specified, tag the result with this tag.
Raises:
ValueError: If responses is not of the correct type.
"""
if not isinstance(response, rdfvalue.RDFValue):
raise ValueError("SendReply can only send RDFValues")
if not any(isinstance(response, t) for t in self.result_types):
logging.warning("Flow %s sends response of unexpected type %s.",
type(self).__name__,
type(response).__name__)
reply = rdf_flow_objects.FlowResult(
client_id=self.rdf_flow.client_id,
flow_id=self.rdf_flow.flow_id,
hunt_id=self.rdf_flow.parent_hunt_id,
payload=response,
tag=tag)
if self.rdf_flow.parent_flow_id:
response = rdf_flow_objects.FlowResponse(
client_id=self.rdf_flow.client_id,
request_id=self.rdf_flow.parent_request_id,
response_id=self.GetNextResponseId(),
payload=response,
flow_id=self.rdf_flow.parent_flow_id,
tag=tag)
self.flow_responses.append(response)
# For nested flows we want the replies to be written,
# but not to be processed by output plugins.
self.replies_to_write.append(reply)
else:
self.replies_to_write.append(reply)
self.replies_to_process.append(reply)
self.rdf_flow.num_replies_sent += 1
# Keeping track of result types/tags in a plain Python
# _num_replies_per_type_tag dict. In RDFValues/proto2 we have to represent
# dictionaries as lists of key-value pairs (i.e. there's no library
# support for dicts as data structures). Hence, updating a key would require
# iterating over the pairs - which might get expensive for hundreds of
# thousands of results. To avoid the issue we keep a non-serialized Python
# dict to be later accumulated into a serializable FlowResultCount
# in PersistState().
key = (type(response).__name__, tag or "")
self._num_replies_per_type_tag[key] += 1
def testResponsesForUnknownFlow(self):
client_id = u"C.1234567890123456"
flow_id = u"1234ABCD"
with self.assertRaises(db.UnknownFlowError):
self.db.WriteFlowResponses([
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1)
])
def _ResponsesAndStatus(self, client_id, flow_id, request_id,
num_responses):
return [
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=request_id,
response_id=i)
for i in range(1, num_responses + 1)
] + [
rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=request_id,
response_id=num_responses + 1)
]
def testResponsesForUnknownFlow(self):
client_id = u"C.1234567890123456"
flow_id = u"1234ABCD"
# This will not raise but also not write anything.
with test_lib.SuppressLogs():
self.db.WriteFlowResponses([
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1)
])
read = self.db.ReadAllFlowRequestsAndResponses(client_id, flow_id)
self.assertEqual(read, [])
def testPathSpecCasingIsCorrected(self):
flow = memory.DumpProcessMemory(rdf_flow_objects.Flow())
flow.SendReply = mock.Mock(spec=flow.SendReply)
request = rdf_flow_objects.FlowRequest(
request_data={
"YaraProcessDumpResponse":
rdf_memory.YaraProcessDumpResponse(dumped_processes=[
rdf_memory.YaraProcessDumpInformation(memory_regions=[
rdf_memory.ProcessMemoryRegion(
start=1,
size=1,
file=rdf_paths.PathSpec.Temp(
path="/C:/grr/x_1_0_1.tmp")),
rdf_memory.ProcessMemoryRegion(
start=1,
size=1,
file=rdf_paths.PathSpec.Temp(
path="/C:/GRR/x_1_1_2.tmp"))
])
])
})
pathspecs = [
rdf_paths.PathSpec.Temp(path="/C:/Grr/x_1_0_1.tmp"),
rdf_paths.PathSpec.Temp(path="/C:/Grr/x_1_1_2.tmp")
]
responses = flow_responses.Responses.FromResponses(
request, [
rdf_flow_objects.FlowResponse(payload=rdf_client_fs.StatEntry(
pathspec=pathspec)) for pathspec in pathspecs
])
flow.ProcessMemoryRegions(responses)
flow.SendReply.assert_any_call(
rdf_memory.YaraProcessDumpResponse(dumped_processes=[
rdf_memory.YaraProcessDumpInformation(memory_regions=[
rdf_memory.ProcessMemoryRegion(
start=1,
size=1,
file=rdf_paths.PathSpec.Temp(
path="/C:/Grr/x_1_0_1.tmp")),
rdf_memory.ProcessMemoryRegion(
start=1,
size=1,
file=rdf_paths.PathSpec.Temp(
path="/C:/Grr/x_1_1_2.tmp"))
])
]))
def SendReply(self,
response: rdfvalue.RDFValue,
tag: Optional[str] = None) -> None:
"""Allows this flow to send a message to its parent flow.
If this flow does not have a parent, the message is saved to the database
as flow result.
Args:
response: An RDFValue() instance to be sent to the parent.
tag: If specified, tag the result with this tag.
Raises:
ValueError: If responses is not of the correct type.
"""
if not isinstance(response, rdfvalue.RDFValue):
raise ValueError("SendReply can only send RDFValues")
if not any(isinstance(response, t) for t in self.result_types):
logging.warning("Flow %s sends response of unexpected type %s.",
type(self).__name__,
type(response).__name__)
reply = rdf_flow_objects.FlowResult(
client_id=self.rdf_flow.client_id,
flow_id=self.rdf_flow.flow_id,
hunt_id=self.rdf_flow.parent_hunt_id,
payload=response,
tag=tag)
if self.rdf_flow.parent_flow_id:
response = rdf_flow_objects.FlowResponse(
client_id=self.rdf_flow.client_id,
request_id=self.rdf_flow.parent_request_id,
response_id=self.GetNextResponseId(),
payload=response,
flow_id=self.rdf_flow.parent_flow_id,
tag=tag)
self.flow_responses.append(response)
# For nested flows we want the replies to be written,
# but not to be processed by output plugins.
self.replies_to_write.append(reply)
else:
self.replies_to_write.append(reply)
self.replies_to_process.append(reply)
self.rdf_flow.num_replies_sent += 1
def _WriteRequestAndResponses(self, client_id, flow_id):
rdf_flow = rdf_flow_objects.Flow(client_id=client_id, flow_id=flow_id)
self.db.WriteFlowObject(rdf_flow)
for request_id in range(1, 4):
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=request_id)
self.db.WriteFlowRequests([request])
for response_id in range(1, 3):
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=request_id,
response_id=response_id)
self.db.WriteFlowResponses([response])
def _StartFlow(self, client_id, flow_cls, **kw):
flow_id = flow.StartFlow(flow_cls=flow_cls, client_id=client_id, **kw)
# Lease the client message.
data_store.REL_DB.LeaseClientActionRequests(
client_id,
lease_time=rdfvalue.Duration.From(10000, rdfvalue.SECONDS))
# Write some responses. In the relational db, the client queue will be
# cleaned up as soon as all responses are available. Therefore we cheat
# here and make it look like the request needs more responses so it's not
# considered complete.
# Write the status first. This will mark the request as waiting for 2
# responses.
status = rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=2)
data_store.REL_DB.WriteFlowResponses([status])
# Now we read the request, adjust the number, and write it back.
reqs = data_store.REL_DB.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
req = reqs[0][0]
req.nr_responses_expected = 99
data_store.REL_DB.WriteFlowRequests([req])
# This response now won't trigger any deletion of client messages.
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1,
payload=rdf_client.Process(name="test_process"))
data_store.REL_DB.WriteFlowResponses([response])
# This is not strictly needed as we don't display this information in the
# UI.
req.nr_responses_expected = 2
data_store.REL_DB.WriteFlowRequests([req])
return flow_id
def testReadFlowRequestsReadyForProcessing(self):
client_id = u"C.1234567890000000"
flow_id = u"12344321"
requests_for_processing = self.db.ReadFlowRequestsReadyForProcessing(
client_id, flow_id, next_needed_request=1)
self.assertEqual(requests_for_processing, {})
client_id, flow_id = self._SetupClientAndFlow(
next_request_to_process=3)
for request_id in [1, 3, 4, 5, 7]:
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=request_id,
needs_processing=True)
self.db.WriteFlowRequests([request])
# Request 4 has some responses.
responses = [
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=4,
response_id=i) for i in range(3)
]
self.db.WriteFlowResponses(responses)
requests_for_processing = self.db.ReadFlowRequestsReadyForProcessing(
client_id, flow_id, next_needed_request=3)
# We expect three requests here. Req #1 is old and should not be there, req
# #7 can't be processed since we are missing #6 in between. That leaves
# requests #3, #4 and #5.
self.assertEqual(len(requests_for_processing), 3)
self.assertEqual(list(requests_for_processing), [3, 4, 5])
for request_id in requests_for_processing:
request, _ = requests_for_processing[request_id]
self.assertEqual(request_id, request.request_id)
self.assertEqual(requests_for_processing[4][1], responses)
def testStatusMessagesCanBeWrittenAndRead(self):
client_id, flow_id = self._SetupClientAndFlow()
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1,
needs_processing=False)
self.db.WriteFlowRequests([request])
responses = [
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=i) for i in range(3)
]
# Also store an Iterator, why not.
responses.append(
rdf_flow_objects.FlowIterator(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=3))
responses.append(
rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=4))
self.db.WriteFlowResponses(responses)
all_requests = self.db.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertEqual(len(all_requests), 1)
_, read_responses = all_requests[0]
self.assertEqual(list(read_responses), [0, 1, 2, 3, 4])
for i in range(3):
self.assertIsInstance(read_responses[i],
rdf_flow_objects.FlowResponse)
self.assertIsInstance(read_responses[3], rdf_flow_objects.FlowIterator)
self.assertIsInstance(read_responses[4], rdf_flow_objects.FlowStatus)
def _StartFlow(self, client_id, flow_cls, **kw):
if data_store.RelationalDBFlowsEnabled():
flow_id = flow.StartFlow(flow_cls=flow_cls,
client_id=client_id,
**kw)
# Lease the client message.
data_store.REL_DB.LeaseClientMessages(
client_id, lease_time=rdfvalue.Duration("10000s"))
# Write some responses.
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1,
payload=rdf_client.Process(name="test_process"))
status = rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=2)
data_store.REL_DB.WriteFlowResponses([response, status])
return flow_id
else:
flow_id = flow.StartAFF4Flow(
flow_name=compatibility.GetName(flow_cls),
client_id=client_id,
token=self.token,
**kw).Basename()
# Have the client write some responses.
test_process = rdf_client.Process(name="test_process")
mock = flow_test_lib.MockClient(client_id,
action_mocks.ListProcessesMock(
[test_process]),
token=self.token)
mock.Next()
return flow_id
def testUninitializedDynamicValueIsNonePerDefault(self): response = rdf_flow_objects.FlowResponse() # Do not set payload. self.assertIsNone(response.payload)
