def ProcessResponses(self,
plugin_args=None,
responses=None,
process_responses_separately=False):
plugin = csv_plugin.CSVOutputPlugin(source_urn=self.results_urn,
output_base_urn=self.base_urn,
args=plugin_args,
token=self.token)
plugin.Initialize()
messages = []
for response in responses:
messages.append(
rdfvalue.GrrMessage(source=self.client_id, payload=response))
if process_responses_separately:
for message in messages:
plugin.ProcessResponses([message])
else:
plugin.ProcessResponses(messages)
plugin.Flush()
return plugin.OpenOutputStreams()
def testDeleteFlowRequestStates(self):
"""Check that we can efficiently destroy a single flow request."""
session_id = rdfvalue.SessionID("aff4:/flows/W:test3")
request = rdfvalue.RequestState(id=1,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
with queue_manager.QueueManager(token=self.token) as manager:
manager.QueueRequest(session_id, request)
manager.QueueResponse(
session_id, rdfvalue.GrrMessage(request_id=1, response_id=1))
# Check the request and responses are there.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0][0], request)
with queue_manager.QueueManager(token=self.token) as manager:
manager.DeleteFlowRequestStates(session_id, request)
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 0)
def testRDFValueCollections(self):
urn = "aff4:/test/collection"
fd = aff4.FACTORY.Create(urn,
"RDFValueCollection",
mode="w",
token=self.token)
for i in range(5):
fd.Add(rdfvalue.GrrMessage(request_id=i))
fd.Close()
fd = aff4.FACTORY.Open(urn, token=self.token)
# Make sure items are stored in order.
j = 0
for j, x in enumerate(fd):
self.assertEqual(j, x.request_id)
self.assertEqual(j, 4)
for j in range(len(fd)):
self.assertEqual(fd[j].request_id, j)
self.assertIsNone(fd[5])
def AddFileToFileStore(pathspec=None, client_id=None, token=None):
"""Adds file with given pathspec to the hash file store."""
if pathspec is None:
raise ValueError("pathspec can't be None")
if client_id is None:
raise ValueError("client_id can't be None")
urn = aff4.AFF4Object.VFSGRRClient.PathspecToURN(pathspec, client_id)
client_mock = action_mocks.ActionMock("TransferBuffer", "StatFile",
"HashBuffer")
for _ in test_lib.TestFlowHelper(
"GetFile", client_mock, token=token, client_id=client_id,
pathspec=pathspec):
pass
auth_state = rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED
flow.Events.PublishEvent(
"FileStore.AddFileToStore",
rdfvalue.GrrMessage(payload=urn, auth_state=auth_state),
token=token)
worker = test_lib.MockWorker(token=token)
worker.Simulate()
def FetchCompletedRequests(self, session_id, timestamp=None):
"""Fetch all the requests with a status message queued for them."""
subject = session_id.Add("state")
requests = {}
status = {}
if timestamp is None:
timestamp = (0, self.frozen_timestamp or rdfvalue.RDFDatetime().Now())
for predicate, serialized, _ in self.data_store.ResolveRegex(
subject, [self.FLOW_REQUEST_REGEX, self.FLOW_STATUS_REGEX],
token=self.token, limit=self.request_limit, timestamp=timestamp):
parts = predicate.split(":", 3)
request_id = parts[2]
if parts[1] == "status":
status[request_id] = serialized
else:
requests[request_id] = serialized
for request_id, serialized in sorted(requests.items()):
if request_id in status:
yield (rdfvalue.RequestState(serialized),
rdfvalue.GrrMessage(status[request_id]))
def testReSchedule(self):
"""Test the ability to re-schedule a task."""
test_queue = rdfvalue.RDFURN("fooReschedule")
task = rdfvalue.GrrMessage(queue=test_queue,
task_ttl=5,
session_id="aff4:/Test")
manager = queue_manager.QueueManager(token=self.token)
manager.Schedule([task])
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
# Record the task id
original_id = tasks[0].task_id
# If we try to get another lease on it we should fail
tasks_2 = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks_2), 0)
# Now we reschedule it
manager.Schedule(tasks)
# The id should not change
self.assertEqual(tasks[0].task_id, original_id)
# If we try to get another lease on it we should not fail
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
# But the id should not change
self.assertEqual(tasks[0].task_id, original_id)
def CallClient(self,
action_name,
request=None,
next_state=None,
client_id=None,
request_data=None,
start_time=None,
**kwargs):
"""Calls the client asynchronously.
This sends a message to the client to invoke an Action. The run
action may send back many responses. These will be queued by the
framework until a status message is sent by the client. The status
message will cause the entire transaction to be committed to the
specified state.
Args:
action_name: The function to call on the client.
request: The request to send to the client. If not specified (Or None) we
create a new RDFValue using the kwargs.
next_state: The state in this flow, that responses to this
message should go to.
client_id: rdfvalue.ClientURN to send the request to.
request_data: A dict which will be available in the RequestState
protobuf. The Responses object maintains a reference to this
protobuf for use in the execution of the state method. (so you can
access this data by responses.request). Valid values are
strings, unicode and protobufs.
start_time: Call the client at this time. This Delays the client request
for into the future.
**kwargs: These args will be used to construct the client action semantic
protobuf.
Raises:
FlowRunnerError: If next_state is not one of the allowed next states.
RuntimeError: The request passed to the client does not have the correct
type.
"""
if client_id is None:
client_id = self.args.client_id
if client_id is None:
raise FlowRunnerError(
"CallClient() is used on a flow which was not "
"started with a client.")
if not isinstance(client_id, rdfvalue.ClientURN):
# Try turning it into a ClientURN
client_id = rdfvalue.ClientURN(client_id)
# Retrieve the correct rdfvalue to use for this client action.
try:
action = actions.ActionPlugin.classes[action_name]
except KeyError:
raise RuntimeError("Client action %s not found." % action_name)
if action.in_rdfvalue is None:
if request:
raise RuntimeError("Client action %s does not expect args." %
action_name)
else:
if request is None:
# Create a new rdf request.
request = action.in_rdfvalue(**kwargs)
else:
# Verify that the request type matches the client action requirements.
if not isinstance(request, action.in_rdfvalue):
raise RuntimeError("Client action expected %s but got %s" %
(action.in_rdfvalue, type(request)))
outbound_id = self.GetNextOutboundId()
# Create a new request state
state = rdfvalue.RequestState(id=outbound_id,
session_id=self.session_id,
next_state=next_state,
client_id=client_id)
if request_data is not None:
state.data = rdfvalue.Dict(request_data)
# Send the message with the request state
msg = rdfvalue.GrrMessage(session_id=utils.SmartUnicode(
self.session_id),
name=action_name,
request_id=outbound_id,
priority=self.args.priority,
require_fastpoll=self.args.require_fastpoll,
queue=client_id.Queue(),
payload=request)
if self.context.remaining_cpu_quota:
msg.cpu_limit = int(self.context.remaining_cpu_quota)
cpu_usage = self.context.client_resources.cpu_usage
if self.context.args.cpu_limit:
msg.cpu_limit = max(
self.context.args.cpu_limit - cpu_usage.user_cpu_time -
cpu_usage.system_cpu_time, 0)
if msg.cpu_limit == 0:
raise FlowRunnerError("CPU limit exceeded.")
if self.context.args.network_bytes_limit:
msg.network_bytes_limit = max(
self.context.args.network_bytes_limit -
self.context.network_bytes_sent, 0)
if msg.network_bytes_limit == 0:
raise FlowRunnerError("Network limit exceeded.")
state.request = msg
self.QueueRequest(state, timestamp=start_time)
def CallState(self,
messages=None,
next_state="",
request_data=None,
start_time=None):
"""This method is used to schedule a new state on a different worker.
This is basically the same as CallFlow() except we are calling
ourselves. The state will be invoked in a later time and receive all the
messages we send.
Args:
messages: A list of rdfvalues to send. If the last one is not a
GrrStatus, we append an OK Status.
next_state: The state in this flow to be invoked with the responses.
request_data: Any dict provided here will be available in the
RequestState protobuf. The Responses object maintains a reference
to this protobuf for use in the execution of the state method. (so
you can access this data by responses.request).
start_time: Start the flow at this time. This Delays notification for
flow processing into the future. Note that the flow may still be
processed earlier if there are client responses waiting.
Raises:
FlowRunnerError: if the next state is not valid.
"""
if messages is None:
messages = []
# Check if the state is valid
if not getattr(self.flow_obj, next_state):
raise FlowRunnerError("Next state %s is invalid.")
# Queue the response message to the parent flow
request_state = rdfvalue.RequestState(
id=self.GetNextOutboundId(),
session_id=self.context.session_id,
client_id=self.args.client_id,
next_state=next_state)
if request_data:
request_state.data = rdfvalue.Dict().FromDict(request_data)
self.QueueRequest(request_state, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus):
messages.append(rdfvalue.GrrStatus())
# Send all the messages
for i, payload in enumerate(messages):
if isinstance(payload, rdfvalue.RDFValue):
msg = rdfvalue.GrrMessage(
session_id=self.session_id,
request_id=request_state.id,
response_id=1 + i,
auth_state=rdfvalue.GrrMessage.AuthorizationState.
AUTHENTICATED,
payload=payload,
type=rdfvalue.GrrMessage.Type.MESSAGE)
if isinstance(payload, rdfvalue.GrrStatus):
msg.type = rdfvalue.GrrMessage.Type.STATUS
else:
raise FlowRunnerError("Bad message %s of type %s." %
(payload, type(payload)))
self.QueueResponse(msg, start_time)
# Notify the worker about it.
self.QueueNotification(session_id=self.session_id,
timestamp=start_time)
def testEventNotification(self):
"""Test that events are sent to listeners."""
NoClientListener.received_events = []
worker = test_lib.MockWorker(token=self.token)
event = rdfvalue.GrrMessage(
session_id=rdfvalue.SessionID(flow_name="SomeFlow"),
name="test message",
payload=rdfvalue.PathSpec(path="foobar", pathtype="TSK"),
source="aff4:/C.0000000000000001", auth_state="AUTHENTICATED")
# Not allowed to publish a message from a client..
flow.Events.PublishEvent("TestEvent", event, token=self.token)
worker.Simulate()
self.assertEqual(NoClientListener.received_events, [])
event.source = "Source"
# First make the message unauthenticated.
event.auth_state = rdfvalue.GrrMessage.AuthorizationState.UNAUTHENTICATED
# Publish the event.
flow.Events.PublishEvent("TestEvent", event, token=self.token)
worker.Simulate()
# This should not work - the unauthenticated message is dropped.
self.assertEqual(NoClientListener.received_events, [])
# Now make the message authenticated.
event.auth_state = rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED
# Publish the event.
flow.Events.PublishEvent("TestEvent", event, token=self.token)
worker.Simulate()
# This should now work:
self.assertEqual(len(NoClientListener.received_events), 1)
# Make sure the source is correctly propagated.
self.assertEqual(NoClientListener.received_events[0][0].source,
"aff4:/Source")
self.assertEqual(NoClientListener.received_events[0][1].path, "foobar")
NoClientListener.received_events = []
# Now schedule ten events at the same time.
for i in xrange(10):
event.source = "Source%d" % i
flow.Events.PublishEvent("TestEvent", event, token=self.token)
worker.Simulate()
self.assertEqual(len(NoClientListener.received_events), 10)
# Events do not have to be delivered in order so we sort them here for
# comparison.
NoClientListener.received_events.sort(key=lambda x: x[0].source)
for i in range(10):
self.assertEqual(NoClientListener.received_events[i][0].source,
"aff4:/Source%d" % i)
self.assertEqual(NoClientListener.received_events[i][1].path, "foobar")
def SendReply(self, rdf_value,
message_type=rdfvalue.GrrMessage.Type.MESSAGE, **kw):
message = rdfvalue.GrrMessage(
type=message_type, payload=rdf_value, **kw)
self.responses.append(message)
def SendReply(self,
rdf_value=None,
request_id=None,
response_id=None,
priority=None,
session_id="W:0",
message_type=None,
name=None,
require_fastpoll=None,
ttl=None,
blocking=True,
task_id=None):
"""Send the protobuf to the server.
Args:
rdf_value: The RDFvalue to return.
request_id: The id of the request this is a response to.
response_id: The id of this response.
priority: The priority of this message, used to jump the scheduling queue.
session_id: The session id of the flow.
message_type: The contents of this message, MESSAGE, STATUS, ITERATOR or
RDF_VALUE.
name: The name of the client action that sends this response.
require_fastpoll: If set, this will set the client to fastpoll mode after
sending this message.
ttl: The time to live of this message.
blocking: If the output queue is full, block until there is space.
task_id: The task ID that the request was queued at. We send this back to
the server so it can de-queue the request.
Raises:
RuntimeError: An object other than an RDFValue was passed for sending.
"""
if not isinstance(rdf_value, rdfvalue.RDFValue):
raise RuntimeError(
"Sending objects other than RDFValues not supported.")
message = rdfvalue.GrrMessage(session_id=session_id,
task_id=task_id,
name=name,
response_id=response_id,
request_id=request_id,
priority=priority,
require_fastpoll=require_fastpoll,
ttl=ttl,
type=message_type)
if rdf_value:
message.payload = rdf_value
serialized_message = message.SerializeToString()
self.ChargeBytesToSession(session_id, len(serialized_message))
if message.type == rdfvalue.GrrMessage.Type.STATUS:
rdf_value.network_bytes_sent = self.sent_bytes_per_flow[session_id]
del self.sent_bytes_per_flow[session_id]
message.args = rdf_value.SerializeToString()
try:
self.QueueResponse(message,
priority=message.priority,
blocking=blocking)
except Queue.Full:
# In the case of a non blocking send, we reraise the exception to notify
# the caller that something went wrong.
if not blocking:
raise
# There is nothing we can do about it here - we just lose the message and
# keep going.
logging.info("Queue is full, dropping messages.")
def testQueueing(self):
"""Tests that queueing and fetching of requests and responses work."""
session_id = rdfvalue.SessionID(flow_name="test")
request = rdfvalue.RequestState(id=1,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
with queue_manager.QueueManager(token=self.token) as manager:
manager.QueueRequest(session_id, request)
# We only have one unanswered request on the queue.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0], (request, []))
# FetchCompletedRequests should return nothing now.
self.assertEqual(list(manager.FetchCompletedRequests(session_id)), [])
# Now queue more requests and responses:
with queue_manager.QueueManager(token=self.token) as manager:
# Start with request 2 - leave request 1 un-responded to.
for request_id in range(2, 5):
request = rdfvalue.RequestState(id=request_id,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
manager.QueueRequest(session_id, request)
response_id = None
for response_id in range(1, 10):
# Normal message.
manager.QueueResponse(
session_id,
rdfvalue.GrrMessage(request_id=request_id,
response_id=response_id))
# And a status message.
manager.QueueResponse(
session_id,
rdfvalue.GrrMessage(request_id=request_id,
response_id=response_id + 1,
type=rdfvalue.GrrMessage.Type.STATUS))
completed_requests = list(manager.FetchCompletedRequests(session_id))
self.assertEqual(len(completed_requests), 3)
# First completed message is request_id = 2 with 10 responses.
self.assertEqual(completed_requests[0][0].id, 2)
# Last message is the status message.
self.assertEqual(completed_requests[0][-1].type,
rdfvalue.GrrMessage.Type.STATUS)
self.assertEqual(completed_requests[0][-1].response_id, 10)
# Now fetch all the completed responses. Set the limit so we only fetch some
# of the responses.
completed_response = list(manager.FetchCompletedResponses(session_id))
self.assertEqual(len(completed_response), 3)
for i, (request, responses) in enumerate(completed_response, 2):
self.assertEqual(request.id, i)
self.assertEqual(len(responses), 10)
# Now check if the limit is enforced. The limit refers to the total number
# of responses to return. We ask for maximum 15 responses, so we should get
# a single request with 10 responses (since 2 requests will exceed the
# limit).
more_data = False
i = 0
try:
partial_response = manager.FetchCompletedResponses(session_id,
limit=15)
for i, (request, responses) in enumerate(partial_response, 2):
self.assertEqual(request.id, i)
self.assertEqual(len(responses), 10)
except queue_manager.MoreDataException:
more_data = True
# Returns the first request that is completed.
self.assertEqual(i, 3)
# Make sure the manager told us that more data is available.
self.assertTrue(more_data)
def testNannyMessage(self):
nanny_message = "Oh no!"
try:
old_send_email = email_alerts.SendEmail
self.email_message = {}
def SendEmail(address, sender, title, message, **_):
self.email_message.update(
dict(address=address,
sender=sender,
title=title,
message=message))
email_alerts.SendEmail = SendEmail
msg = rdfvalue.GrrMessage(
session_id=rdfvalue.SessionID("aff4:/flows/W:NannyMessage"),
args=rdfvalue.DataBlob(
string=nanny_message).SerializeToString(),
source=self.client_id,
auth_state=rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED
)
# This is normally done by the FrontEnd when a CLIENT_KILLED message is
# received.
flow.Events.PublishEvent("NannyMessage", msg, token=self.token)
# Now emulate a worker to process the event.
worker = test_lib.MockWorker(token=self.token)
while worker.Next():
pass
worker.pool.Join()
# We expect the email to be sent.
self.assertEqual(self.email_message.get("address"),
config_lib.CONFIG["Monitoring.alert_email"])
self.assertTrue(str(self.client_id) in self.email_message["title"])
# Make sure the message is included in the email message.
self.assertTrue(nanny_message in self.email_message["message"])
# Make sure crashes RDFValueCollections are created and written
# into proper locations. First check the per-client crashes collection.
client_crashes = list(
aff4.FACTORY.Open(self.client_id.Add("crashes"),
aff4_type="PackedVersionedCollection",
token=self.token))
self.assertEqual(len(client_crashes), 1)
crash = client_crashes[0]
self.assertEqual(crash.client_id, self.client_id)
self.assertEqual(crash.client_info.client_name, "GRR Monitor")
self.assertEqual(crash.crash_type, "aff4:/flows/W:NannyMessage")
self.assertEqual(crash.crash_message, nanny_message)
# Check global crash collection. Check that crash written there is
# equal to per-client crash.
global_crashes = list(
aff4.FACTORY.Open(aff4.ROOT_URN.Add("crashes"),
aff4_type="PackedVersionedCollection",
token=self.token))
self.assertEqual(len(global_crashes), 1)
self.assertEqual(global_crashes[0], crash)
finally:
email_alerts.SendEmail = old_send_email
def AddNewElementToCollection(*unused_args, **unused_kwargs):
with aff4.FACTORY.Create("aff4:/tmp/coll",
"PackedVersionedCollection",
mode="w",
token=self.token) as fd:
fd.Add(rdfvalue.GrrMessage(request_id=1))
def CallState(self,
messages=None,
next_state="",
client_id=None,
request_data=None,
start_time=None):
"""This method is used to asynchronously schedule a new hunt state.
The state will be invoked in a later time and receive all the messages
we send.
Args:
messages: A list of rdfvalues to send. If the last one is not a
GrrStatus, we append an OK Status.
next_state: The state in this hunt to be invoked with the responses.
client_id: ClientURN to use in scheduled requests.
request_data: Any dict provided here will be available in the
RequestState protobuf. The Responses object maintains a
reference to this protobuf for use in the execution of the
state method. (so you can access this data by
responses.request).
start_time: Schedule the state at this time. This delays notification
and messages for processing into the future.
Raises:
ValueError: on arguments error.
"""
if messages is None:
messages = []
if not next_state:
raise ValueError("next_state can't be empty.")
# Now we construct a special response which will be sent to the hunt
# flow. Randomize the request_id so we do not overwrite other messages in
# the queue.
request_state = rdfvalue.RequestState(
id=utils.PRNG.GetULong(),
session_id=self.context.session_id,
client_id=client_id,
next_state=next_state)
if request_data:
request_state.data = rdfvalue.Dict().FromDict(request_data)
self.QueueRequest(request_state, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus):
messages.append(rdfvalue.GrrStatus())
# Send all the messages
for i, payload in enumerate(messages):
if isinstance(payload, rdfvalue.RDFValue):
msg = rdfvalue.GrrMessage(
session_id=self.session_id,
request_id=request_state.id,
response_id=1 + i,
auth_state=rdfvalue.GrrMessage.AuthorizationState.
AUTHENTICATED,
payload=payload,
type=rdfvalue.GrrMessage.Type.MESSAGE)
if isinstance(payload, rdfvalue.GrrStatus):
msg.type = rdfvalue.GrrMessage.Type.STATUS
else:
raise flow_runner.FlowRunnerError(
"Bad message %s of type %s." % (payload, type(payload)))
self.QueueResponse(msg, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus):
messages.append(rdfvalue.GrrStatus())
# Notify the worker about it.
self.QueueNotification(session_id=self.session_id,
timestamp=start_time)
def _CreateResult(self, success, clientid):
success = rdfvalue.EndToEndTestResult(success=success)
return rdfvalue.GrrMessage(source=clientid,
payload=success)
def testHandleClientMessageRetransmission(self):
"""Check that requests get retransmitted but only if there is no status."""
# Make a new fake client
client_id = self.SetupClients(1)[0]
# Test the standard behavior.
base_time = 1000
msgs_recvd = []
default_ttl = rdfvalue.GrrMessage().task_ttl
with test_lib.FakeTime(base_time):
flow.GRRFlow.StartFlow(client_id=client_id,
flow_name="SendingFlow",
message_count=1,
token=self.token)
for i in range(default_ttl):
with test_lib.FakeTime(base_time + i *
(self.message_expiry_time + 1)):
tasks = self.server.DrainTaskSchedulerQueueForClient(
client_id, 100000, rdfvalue.MessageList())
msgs_recvd.append(tasks)
# Should return a client message (ttl-1) times and nothing afterwards.
self.assertEqual(map(bool, msgs_recvd), [True] *
(rdfvalue.GrrMessage().task_ttl - 1) + [False])
# Now we simulate that the workers are overloaded - the client messages
# arrive but do not get processed in time.
if default_ttl <= 3:
self.fail("TTL too low for this test.")
msgs_recvd = []
with test_lib.FakeTime(base_time):
flow_id = flow.GRRFlow.StartFlow(client_id=client_id,
flow_name="SendingFlow",
message_count=1,
token=self.token)
for i in range(default_ttl):
if i == 2:
self._ScheduleResponseAndStatus(client_id, flow_id)
with test_lib.FakeTime(base_time + i *
(self.message_expiry_time + 1)):
tasks = self.server.DrainTaskSchedulerQueueForClient(
client_id, 100000, rdfvalue.MessageList())
msgs_recvd.append(tasks)
if not tasks:
# Even if the request has not been leased ttl times yet,
# it should be dequeued by now.
new_tasks = queue_manager.QueueManager(
token=self.token).Query(
queue=rdfvalue.ClientURN(client_id).Queue(),
limit=1000)
self.assertEqual(len(new_tasks), 0)
# Should return a client message twice and nothing afterwards.
self.assertEqual(map(bool, msgs_recvd), [True] * 2 + [False] *
(rdfvalue.GrrMessage().task_ttl - 2))
def testExportWithDummyPlugin(self):
pathspec = rdfvalue.PathSpec(pathtype=rdfvalue.PathSpec.PathType.OS,
path=os.path.join(self.base_path,
"winexec_img.dd"))
pathspec.Append(path="/Ext2IFS_1_10b.exe",
pathtype=rdfvalue.PathSpec.PathType.TSK)
urn = aff4.AFF4Object.VFSGRRClient.PathspecToURN(
pathspec, self.client_id)
client_mock = test_lib.ActionMock("TransferBuffer", "StatFile",
"HashBuffer")
for _ in test_lib.TestFlowHelper("GetFile",
client_mock,
token=self.token,
client_id=self.client_id,
pathspec=pathspec):
pass
auth_state = rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED
flow.Events.PublishEvent("FileStore.AddFileToStore",
rdfvalue.GrrMessage(payload=urn,
auth_state=auth_state),
token=self.token)
worker = test_lib.MockWorker(token=self.token)
worker.Simulate()
plugin = hash_file_store_plugin.HashFileStoreExportPlugin()
parser = argparse.ArgumentParser()
plugin.ConfigureArgParser(parser)
plugin.Run(parser.parse_args(args=["--threads", "0", "dummy"]))
responses = DummyOutputPlugin.responses
self.assertEqual(len(responses), 5)
for response in responses:
self.assertTrue(isinstance(response, rdfvalue.FileStoreHash))
self.assertTrue(
rdfvalue.FileStoreHash(
fingerprint_type="pecoff",
hash_type="md5",
hash_value="a3a3259f7b145a21c7b512d876a5da06") in responses)
self.assertTrue(
rdfvalue.FileStoreHash(
fingerprint_type="pecoff",
hash_type="sha1",
hash_value="019bddad9cac09f37f3941a7f285c79d3c7e7801") in
responses)
self.assertTrue(
rdfvalue.FileStoreHash(
fingerprint_type="generic",
hash_type="md5",
hash_value="bb0a15eefe63fd41f8dc9dee01c5cf9a") in responses)
self.assertTrue(
rdfvalue.FileStoreHash(
fingerprint_type="generic",
hash_type="sha1",
hash_value="7dd6bee591dfcb6d75eb705405302c3eab65e21a") in
responses)
self.assertTrue(
rdfvalue.FileStoreHash(
fingerprint_type="generic",
hash_type="sha256",
hash_value="0e8dc93e150021bb4752029ebbff51394aa36f06"
"9cf19901578e4f06017acdb5") in responses)
def testNoValidStatusRaceIsResolved(self):
# This tests for the regression of a long standing race condition we saw
# where notifications would trigger the reading of another request that
# arrives later but wasn't completely written to the database yet.
# Timestamp based notification handling should eliminate this bug.
# We need a random flow object for this test.
session_id = flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker.GRRWorker(worker.DEFAULT_WORKER_QUEUE,
token=self.token)
manager = queue_manager.QueueManager(token=self.token)
manager.DeleteNotification(session_id)
manager.Flush()
# We have a first request that is complete (request_id 1, response_id 1).
self.SendResponse(session_id, "Response 1")
# However, we also have request #2 already coming in. The race is that
# the queue manager might write the status notification to
# session_id/state as "status:00000002" but not the status response
# itself yet under session_id/state/request:00000002
request_id = 2
response_id = 1
flow_manager = queue_manager.QueueManager(token=self.token)
flow_manager.FreezeTimestamp()
flow_manager.QueueResponse(
session_id,
rdfvalue.GrrMessage(source=self.client_id,
session_id=session_id,
payload=rdfvalue.DataBlob(string="Response 2"),
request_id=request_id,
response_id=response_id))
status = rdfvalue.GrrMessage(
source=self.client_id,
session_id=session_id,
payload=rdfvalue.GrrStatus(
status=rdfvalue.GrrStatus.ReturnedStatus.OK),
request_id=request_id,
response_id=response_id + 1,
type=rdfvalue.GrrMessage.Type.STATUS)
# Now we write half the status information.
subject = session_id.Add("state")
queue = flow_manager.to_write.setdefault(subject, {})
queue.setdefault(flow_manager.FLOW_STATUS_TEMPLATE % request_id,
[]).append((status.SerializeToString(), None))
flow_manager.Flush()
# We make the race even a bit harder by saying the new notification gets
# written right before the old one gets deleted. If we are not careful here,
# we delete the new notification as well and the flow becomes stuck.
def WriteNotification(self, arg_session_id, start=None, end=None):
if arg_session_id == session_id:
flow_manager.QueueNotification(session_id=arg_session_id)
flow_manager.Flush()
self.DeleteNotification.old_target(self,
arg_session_id,
start=start,
end=end)
with utils.Stubber(queue_manager.QueueManager, "DeleteNotification",
WriteNotification):
# This should process request 1 but not touch request 2.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
flow_obj = aff4.FACTORY.Open(session_id, token=self.token)
self.assertFalse(flow_obj.state.context.backtrace)
self.assertNotEqual(flow_obj.state.context.state,
rdfvalue.Flow.State.ERROR)
request2_data = data_store.DB.ResolveRegex(session_id.Add("state"),
".*:00000002",
token=self.token)
# Make sure the status field and the original request are still there.
self.assertEqual(len(request2_data), 2)
request1_data = data_store.DB.ResolveRegex(session_id.Add("state"),
".*:00000001",
token=self.token)
# Everything from request 1 should have been deleted.
self.assertEqual(len(request1_data), 0)
# The notification for request 2 should have survived.
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(
worker.DEFAULT_WORKER_QUEUE)
self.assertEqual(len(notifications), 1)
notification = notifications[0]
self.assertEqual(notification.session_id, session_id)
self.assertEqual(notification.timestamp,
flow_manager.frozen_timestamp)
self.assertEqual(RESULTS, ["Response 1"])
# The last missing piece of request 2 is the actual status message.
flow_manager.QueueResponse(session_id, status)
flow_manager.Flush()
# Now make sure request 2 runs as expected.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
self.assertEqual(RESULTS, ["Response 1", "Response 2"])
