class NannyMessageHandler(ClientCrashEventListener):
"""A listener for nanny messages."""
EVENTS = ["NannyMessage"]
well_known_session_id = rdfvalue.SessionID(flow_name="NannyMessage")
mail_template = """
<html><body><h1>GRR nanny message received.</h1>
The nanny for client %(client_id)s (%(hostname)s) just sent a message:<br>
<br>
%(message)s
<br>
Click <a href='%(admin_ui)s/#%(urn)s'> here </a> to access this machine.
<p>%(signature)s</p>
</body></html>"""
subject = "GRR nanny message received from %s."
logline = "Nanny for client %s sent: %s"
@flow.EventHandler(allow_client_access=True)
def ProcessMessage(self, message=None, event=None):
"""Processes this event."""
_ = event
client_id = message.source
message = message.payload.string
logging.info(self.logline, client_id, message)
# Write crash data to AFF4.
client = aff4.FACTORY.Open(client_id, token=self.token)
client_info = client.Get(client.Schema.CLIENT_INFO)
crash_details = rdf_client.ClientCrash(
client_id=client_id,
client_info=client_info,
crash_message=message,
timestamp=long(time.time() * 1e6),
crash_type=self.well_known_session_id)
self.WriteAllCrashDetails(client_id, crash_details)
# Also send email.
if config_lib.CONFIG["Monitoring.alert_email"]:
client = aff4.FACTORY.Open(client_id, token=self.token)
hostname = client.Get(client.Schema.HOSTNAME)
url = urllib.urlencode(
(("c", client_id), ("main", "HostInformation")))
email_alerts.EMAIL_ALERTER.SendEmail(
config_lib.CONFIG["Monitoring.alert_email"],
"GRR server",
self.subject % client_id,
self.mail_template %
dict(client_id=client_id,
admin_ui=config_lib.CONFIG["AdminUI.url"],
hostname=hostname,
signature=config_lib.CONFIG["Email.signature"],
urn=url,
message=message),
is_html=True)
class ClientCrashHandler(ClientCrashEventListener):
"""A listener for client crashes."""
EVENTS = ["ClientCrash"]
well_known_session_id = rdfvalue.SessionID(flow_name="CrashHandler")
mail_template = """
<html><body><h1>GRR client crash report.</h1>
Client %(client_id)s (%(hostname)s) just crashed while executing an action.
Click <a href='%(admin_ui)s/#%(urn)s'> here </a> to access this machine.
<p>Thanks,</p>
<p>%(signature)s</p>
<p>
P.S. The state of the failing flow was:
%(state)s
%(nanny_msg)s
</body></html>"""
@flow.EventHandler(allow_client_access=True)
def ProcessMessage(self, message=None, event=None):
"""Processes this event."""
_ = event
client_id = message.source
nanny_msg = ""
flow_obj = aff4.FACTORY.Open(message.session_id, token=self.token)
# Log.
logging.info("Client crash reported, client %s.", client_id)
# Export.
stats.STATS.IncrementCounter("grr_client_crashes")
# Write crash data to AFF4.
client = aff4.FACTORY.Open(client_id, token=self.token)
client_info = client.Get(client.Schema.CLIENT_INFO)
status = rdf_flows.GrrStatus(message.payload)
crash_details = rdf_client.ClientCrash(
client_id=client_id,
session_id=message.session_id,
client_info=client_info,
crash_message=status.error_message,
timestamp=rdfvalue.RDFDatetime().Now(),
crash_type=self.well_known_session_id)
self.WriteAllCrashDetails(client_id,
crash_details,
flow_session_id=message.session_id)
# Also send email.
to_send = []
try:
hunt_session_id = self._ExtractHuntId(message.session_id)
if hunt_session_id and hunt_session_id != message.session_id:
hunt_obj = aff4.FACTORY.Open(hunt_session_id,
aff4_type="GRRHunt",
token=self.token)
email = hunt_obj.GetRunner().args.crash_alert_email
if email:
to_send.append(email)
except aff4.InstantiationError:
logging.error("Failed to open hunt %s." % hunt_session_id)
email = config_lib.CONFIG["Monitoring.alert_email"]
if email:
to_send.append(email)
for email_address in to_send:
if status.nanny_status:
nanny_msg = "Nanny status: %s" % status.nanny_status
client = aff4.FACTORY.Open(client_id, token=self.token)
hostname = client.Get(client.Schema.HOSTNAME)
url = urllib.urlencode(
(("c", client_id), ("main", "HostInformation")))
renderer = rendering.FindRendererForObject(flow_obj.state)
email_alerts.EMAIL_ALERTER.SendEmail(
email_address,
"GRR server",
"Client %s reported a crash." % client_id,
self.mail_template %
dict(client_id=client_id,
admin_ui=config_lib.CONFIG["AdminUI.url"],
hostname=hostname,
state=renderer.RawHTML(),
urn=url,
nanny_msg=nanny_msg,
signature=config_lib.CONFIG["Email.signature"]),
is_html=True)
if nanny_msg:
msg = "Client crashed, " + nanny_msg
else:
msg = "Client crashed."
# Now terminate the flow.
flow.GRRFlow.TerminateFlow(message.session_id,
reason=msg,
token=self.token,
force=True)
class NannyMessageHandler(ClientCrashEventListener):
"""A listener for nanny messages."""
EVENTS = ["NannyMessage"]
well_known_session_id = rdfvalue.SessionID(flow_name="NannyMessage")
mail_template = jinja2.Template(
"""
<html><body><h1>GRR nanny message received.</h1>
The nanny for client {{ client_id }} ({{ hostname }}) just sent a message:<br>
<br>
{{ message }}
<br>
Click <a href='{{ admin_ui }}/#{{ url }}'> here </a> to access this machine.
<p>{{ signature }}</p>
</body></html>""",
autoescape=True)
subject = "GRR nanny message received from %s."
logline = "Nanny for client %s sent: %s"
@flow.EventHandler(allow_client_access=True)
def ProcessMessage(self, message=None, event=None):
"""Processes this event."""
_ = event
client_id = message.source
message = message.payload.string
logging.info(self.logline, client_id, message)
# Write crash data to AFF4.
client = aff4.FACTORY.Open(client_id, token=self.token)
client_info = client.Get(client.Schema.CLIENT_INFO)
crash_details = rdf_client.ClientCrash(
client_id=client_id,
client_info=client_info,
crash_message=message,
timestamp=long(time.time() * 1e6),
crash_type=self.well_known_session_id)
self.WriteAllCrashDetails(client_id, crash_details)
# Also send email.
if config.CONFIG["Monitoring.alert_email"]:
client = aff4.FACTORY.Open(client_id, token=self.token)
hostname = client.Get(client.Schema.HOSTNAME)
url = "/clients/%s" % client_id.Basename()
body = self.__class__.mail_template.render(
client_id=client_id,
admin_ui=config.CONFIG["AdminUI.url"],
hostname=utils.SmartUnicode(hostname),
signature=config.CONFIG["Email.signature"],
url=url,
message=utils.SmartUnicode(message))
email_alerts.EMAIL_ALERTER.SendEmail(
config.CONFIG["Monitoring.alert_email"],
"GRR server",
self.subject % client_id,
utils.SmartStr(body),
is_html=True)
class ClientCrashHandler(ClientCrashEventListener):
"""A listener for client crashes."""
EVENTS = ["ClientCrash"]
well_known_session_id = rdfvalue.SessionID(flow_name="CrashHandler")
mail_template = jinja2.Template(
"""
<html><body><h1>GRR client crash report.</h1>
Client {{ client_id }} ({{ hostname }}) just crashed while executing an action.
Click <a href='{{ admin_ui }}/#{{ url }}'> here </a> to access this machine.
<p>Thanks,</p>
<p>{{ signature }}</p>
<p>
P.S. The state of the failing flow was:
<pre>{{ context }}</pre>
<pre>{{ state }}</pre>
<pre>{{ args }}</pre>
<pre>{{ runner_args }}</pre>
{{ nanny_msg }}
</body></html>""",
autoescape=True)
@flow.EventHandler(allow_client_access=True)
def ProcessMessage(self, message=None, event=None):
"""Processes this event."""
_ = event
nanny_msg = ""
crash_details = message.payload
client_id = crash_details.client_id
# The session id of the flow that crashed.
session_id = crash_details.session_id
flow_obj = aff4.FACTORY.Open(session_id, token=self.token)
# Log.
logging.info("Client crash reported, client %s.", client_id)
# Only kill the flow if it does not handle its own crashes. Some flows
# restart the client and therefore expect to get a crash notification.
if flow_obj.handles_crashes:
return
# Export.
stats.STATS.IncrementCounter("grr_client_crashes")
# Write crash data to AFF4.
client = aff4.FACTORY.Open(client_id, token=self.token)
client_info = client.Get(client.Schema.CLIENT_INFO)
crash_details.client_info = client_info
crash_details.crash_type = self.well_known_session_id
self.WriteAllCrashDetails(
client_id, crash_details, flow_session_id=session_id)
# Also send email.
to_send = []
try:
hunt_session_id = self._ExtractHuntId(session_id)
if hunt_session_id and hunt_session_id != session_id:
hunt_obj = aff4.FACTORY.Open(
hunt_session_id, aff4_type=implementation.GRRHunt, token=self.token)
email = hunt_obj.runner_args.crash_alert_email
if email:
to_send.append(email)
except aff4.InstantiationError:
logging.error("Failed to open hunt %s.", hunt_session_id)
email = config.CONFIG["Monitoring.alert_email"]
if email:
to_send.append(email)
for email_address in to_send:
if crash_details.nanny_status:
nanny_msg = "Nanny status: %s" % crash_details.nanny_status
client = aff4.FACTORY.Open(client_id, token=self.token)
hostname = client.Get(client.Schema.HOSTNAME)
url = "/clients/%s" % client_id.Basename()
body = self.__class__.mail_template.render(
client_id=client_id,
admin_ui=config.CONFIG["AdminUI.url"],
hostname=utils.SmartUnicode(hostname),
context=utils.SmartUnicode(flow_obj.context),
state=utils.SmartUnicode(flow_obj.state),
args=utils.SmartUnicode(flow_obj.args),
runner_args=utils.SmartUnicode(flow_obj.runner_args),
urn=url,
nanny_msg=utils.SmartUnicode(nanny_msg),
signature=config.CONFIG["Email.signature"]),
email_alerts.EMAIL_ALERTER.SendEmail(
email_address,
"GRR server",
"Client %s reported a crash." % client_id,
utils.SmartStr(body),
is_html=True)
if nanny_msg:
msg = "Client crashed, " + nanny_msg
else:
msg = "Client crashed."
# Now terminate the flow.
flow.GRRFlow.TerminateFlow(
session_id, reason=msg, token=self.token, force=True)
def testNannyMessage(self):
nanny_message = "Oh no!"
self.email_message = {}
def SendEmail(address, sender, title, message, **_):
self.email_message.update(
dict(address=address,
sender=sender,
title=title,
message=message))
with utils.Stubber(email_alerts.EMAIL_ALERTER, "SendEmail", SendEmail):
msg = rdf_flows.GrrMessage(
session_id=rdfvalue.SessionID(flow_name="NannyMessage"),
payload=rdf_protodict.DataBlob(string=nanny_message),
source=self.client_id,
auth_state=rdf_flows.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/" + queues.FLOWS.Basename() + ":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)
def RegisterWellKnownFlow(self, wkf):
session_id = rdfvalue.SessionID(flow_name=wkf.FLOW_NAME)
self.wkfs[str(session_id)] = wkf
class ClientStartupHandler(flow.EventListener):
well_known_session_id = rdfvalue.SessionID(flow_name="Startup")
@flow.EventHandler(allow_client_access=True, auth_required=False)
def ProcessMessage(self, message=None, event=None):
"""Handle a startup event."""
_ = event
# We accept unauthenticated messages so there are no errors but we don't
# store the results.
if (message.auth_state !=
rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED):
return
client_id = message.source
new_si = message.payload
drift = rdfvalue.Duration("5m")
if data_store.RelationalDBReadEnabled():
current_si = data_store.REL_DB.ReadClientStartupInfo(
client_id.Basename())
# We write the updated record if the client_info has any changes
# or the boot time is more than 5 minutes different.
if (not current_si or current_si.client_info != new_si.client_info
or abs(current_si.boot_time - new_si.boot_time) > drift):
data_store.REL_DB.WriteClientStartupInfo(
client_id.Basename(), new_si)
else:
changes = False
with aff4.FACTORY.Create(client_id,
aff4_grr.VFSGRRClient,
mode="rw",
token=self.token) as client:
old_info = client.Get(client.Schema.CLIENT_INFO)
old_boot = client.Get(client.Schema.LAST_BOOT_TIME, 0)
info = new_si.client_info
# Only write to the datastore if we have new information.
if info != old_info:
client.Set(client.Schema.CLIENT_INFO(info))
changes = True
client.AddLabels(info.labels, owner="GRR")
# Allow for some drift in the boot times (5 minutes).
if not old_boot or abs(old_boot - new_si.boot_time) > drift:
client.Set(client.Schema.LAST_BOOT_TIME(new_si.boot_time))
changes = True
if data_store.RelationalDBWriteEnabled() and changes:
try:
data_store.REL_DB.WriteClientStartupInfo(
client_id.Basename(), new_si)
except db.UnknownClientError:
pass
events.Events.PublishEventInline("ClientStartup",
message,
token=self.token)
class ClientCrashHandler(ClientCrashEventListener):
"""A listener for client crashes."""
EVENTS = ["ClientCrash"]
well_known_session_id = rdfvalue.SessionID(flow_name="CrashHandler")
mail_template = """
<html><body><h1>GRR client crash report.</h1>
Client %(client_id)s (%(hostname)s) just crashed while executing an action.
Click <a href='%(admin_ui)s/#%(urn)s'> here </a> to access this machine.
<p>Thanks,</p>
<p>%(signature)s</p>
<p>
P.S. The state of the failing flow was:
%(context)s
%(state)s
%(args)s
%(runner_args)s
%(nanny_msg)s
</body></html>"""
def ConvertRDFValueToHTML(self, rdfvalue_obj):
raw_str = utils.SmartUnicode(rdfvalue_obj)
return "<pre>%s</pre>" % cgi.escape(raw_str)
@flow.EventHandler(allow_client_access=True)
def ProcessMessage(self, message=None, event=None):
"""Processes this event."""
_ = event
nanny_msg = ""
crash_details = message.payload
client_id = crash_details.client_id
# The session id of the flow that crashed.
session_id = crash_details.session_id
flow_obj = aff4.FACTORY.Open(session_id, token=self.token)
# Log.
logging.info("Client crash reported, client %s.", client_id)
# Only kill the flow if it does not handle its own crashes. Some flows
# restart the client and therefore expect to get a crash notification.
if flow_obj.handles_crashes:
return
# Export.
stats.STATS.IncrementCounter("grr_client_crashes")
# Write crash data to AFF4.
client = aff4.FACTORY.Open(client_id, token=self.token)
client_info = client.Get(client.Schema.CLIENT_INFO)
crash_details.client_info = client_info
crash_details.crash_type = self.well_known_session_id
self.WriteAllCrashDetails(client_id,
crash_details,
flow_session_id=session_id)
# Also send email.
to_send = []
try:
hunt_session_id = self._ExtractHuntId(session_id)
if hunt_session_id and hunt_session_id != session_id:
hunt_obj = aff4.FACTORY.Open(hunt_session_id,
aff4_type=implementation.GRRHunt,
token=self.token)
email = hunt_obj.runner_args.crash_alert_email
if email:
to_send.append(email)
except aff4.InstantiationError:
logging.error("Failed to open hunt %s.", hunt_session_id)
email = config.CONFIG["Monitoring.alert_email"]
if email:
to_send.append(email)
for email_address in to_send:
if crash_details.nanny_status:
nanny_msg = "Nanny status: %s" % crash_details.nanny_status
client = aff4.FACTORY.Open(client_id, token=self.token)
hostname = client.Get(client.Schema.HOSTNAME)
url = urllib.urlencode(
(("c", client_id), ("main", "HostInformation")))
context_html = self.ConvertRDFValueToHTML(flow_obj.context)
state_html = self.ConvertRDFValueToHTML(flow_obj.state)
args_html = self.ConvertRDFValueToHTML(flow_obj.args)
runner_args_html = self.ConvertRDFValueToHTML(flow_obj.runner_args)
email_alerts.EMAIL_ALERTER.SendEmail(
email_address,
"GRR server",
"Client %s reported a crash." % client_id,
self.mail_template %
dict(client_id=client_id,
admin_ui=config.CONFIG["AdminUI.url"],
hostname=hostname,
context=context_html,
state=state_html,
args=args_html,
runner_args=runner_args_html,
urn=url,
nanny_msg=nanny_msg,
signature=config.CONFIG["Email.signature"]),
is_html=True)
if nanny_msg:
msg = "Client crashed, " + nanny_msg
else:
msg = "Client crashed."
# Now terminate the flow.
flow.GRRFlow.TerminateFlow(session_id,
reason=msg,
token=self.token,
force=True)
def testEventNotification(self):
"""Test that events are sent to listeners."""
NoClientListener.received_events = []
worker = worker_test_lib.MockWorker(token=self.token)
event = rdf_flows.GrrMessage(
session_id=rdfvalue.SessionID(flow_name="SomeFlow"),
name="test message",
payload=rdf_paths.PathSpec(path="foobar", pathtype="TSK"),
source="aff4:/C.0000000000000001",
auth_state="AUTHENTICATED")
# Not allowed to publish a message from a client..
events.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 = rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED
# Publish the event.
events.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 = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED
# Publish the event.
events.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
events.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 Handle(self, args, token=None):
"""Render robot "get files" operation status.
This handler relies on URN validation and flow type checking to check the
input parameters to avoid allowing arbitrary reads into the client aff4
space. This handler filters out only the attributes that are appropriate to
release without authorization (authentication is still required).
Args:
args: ApiGetRobotGetFilesOperationStateArgs object.
token: access token.
Returns:
ApiGetRobotGetFilesOperationStateResult object.
Raises:
RobotGetFilesOperationNotFoundError: if operation is not found (i.e.
if the flow is not found or is not a FileFinder flow).
ValueError: if operation id is incorrect. It should match the
aff4:/<client id>/flows/<flow session id> pattern exactly.
"""
# We deconstruct the operation id and rebuild it as a URN to ensure
# that it points to the flow on a client.
urn = rdfvalue.RDFURN(args.operation_id)
urn_components = urn.Split()
if len(urn_components) != 3 or urn_components[1] != "flows":
raise ValueError("Invalid operation id.")
client_id = rdf_client.ClientURN(urn_components[0])
rdfvalue.SessionID.ValidateID(urn_components[2])
flow_id = rdfvalue.SessionID(urn_components[2])
# flow_id looks like aff4:/F:ABCDEF12, convert it into a flow urn for
# the target client.
flow_urn = client_id.Add("flows").Add(flow_id.Basename())
try:
flow_obj = aff4.FACTORY.Open(flow_urn,
aff4_type=file_finder.FileFinder,
token=token)
except aff4.InstantiationError:
raise RobotGetFilesOperationNotFoundError()
try:
# TODO(user): Remove support for RDFValueCollection.
result_collection = aff4.FACTORY.Open(
flow_obj.state.context.output_urn,
aff4_type=aff4_collects.RDFValueCollection,
token=token)
result_count = len(result_collection)
except (aff4.InstantiationError, AttributeError):
try:
result_collection = aff4.FACTORY.Open(
flow_obj.GetRunner().output_urn,
aff4_type=sequential_collection.GeneralIndexedCollection,
token=token)
result_count = len(result_collection)
except aff4.InstantiationError:
result_count = 0
api_flow_obj = ApiFlow().InitFromAff4Object(flow_obj,
flow_id=flow_id.Basename())
return ApiGetRobotGetFilesOperationStateResult(
state=api_flow_obj.state, result_count=result_count)
def testFlowNameGetterReturnsCorrectValues(self):
s = rdfvalue.SessionID("A:12345678")
self.assertEqual(s.FlowName(), "12345678")
s = rdfvalue.SessionID("DEBUG-user1:12345678:hunt")
self.assertEqual(s.FlowName(), "12345678:hunt")
def testQueueGetterReturnsCorrectValues(self):
s = rdfvalue.SessionID("A:12345678")
self.assertEqual(s.Queue(), "A")
s = rdfvalue.SessionID("DEBUG-user1:12345678:hunt")
self.assertEqual(s.Queue(), "DEBUG-user1")
def testSessionIDValidation(self):
rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/A:12345678"))
rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/A:TransferStore"))
rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/DEBUG-user1:12345678"))
rdfvalue.SessionID(
rdfvalue.RDFURN("aff4:/flows/DEBUG-user1:12345678:hunt"))
class WellKnownSessionTest2(WellKnownSessionTest):
"""Another testing well known flow."""
well_known_session_id = rdfvalue.SessionID(queue=rdfvalue.RDFURN("test"),
flow_name="TestSessionId2")
def testQueueing(self):
"""Tests that queueing and fetching of requests and responses work."""
session_id = rdfvalue.SessionID("aff4:/flows/W: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 GenerateSample(self, number=0):
id_str = "%08X" % (number % 2**32)
return rdfvalue.SessionID(flow_name=id_str)
from grr.lib.rdfvalues import flows as rdf_flows
from grr.lib.rdfvalues import objects as rdf_objects
from grr.server.grr_response_server import data_store
from grr.server.grr_response_server import fleetspeak_utils
class Error(Exception):
"""Base class for errors in this module."""
class MoreDataException(Error):
"""Raised when there is more data available."""
session_id_map = {
rdfvalue.SessionID(queue=queues.ENROLLMENT, flow_name="Enrol"): "Enrol",
rdfvalue.SessionID(queue=queues.STATS, flow_name="Stats"): "StatsHandler",
}
def _GetClientIdFromQueue(q):
"""Returns q's client id, if q is a client task queue, otherwise None.
Args:
q: rdfvalue.RDFURN
Returns:
string or None
"""
split = q.Split()
if not split or len(split) < 2:
