def testForemanMessageHandler(self):
with mock.patch.object(foreman.Foreman,
"AssignTasksToClient") as instr:
# Send a message to the Foreman.
client_id = "C.1100110011001100"
data_store.REL_DB.WriteMessageHandlerRequests([
rdf_objects.MessageHandlerRequest(
client_id=client_id,
handler_name="ForemanHandler",
request_id=12345,
request=rdf_protodict.DataBlob())
])
done = threading.Event()
def handle(l):
worker_lib.ProcessMessageHandlerRequests(l)
done.set()
data_store.REL_DB.RegisterMessageHandler(
handle,
worker_lib.GRRWorker.message_handler_lease_time,
limit=1000)
try:
self.assertTrue(done.wait(10))
# Make sure there are no leftover requests.
self.assertEqual(
data_store.REL_DB.ReadMessageHandlerRequests(), [])
instr.assert_called_once_with(client_id)
finally:
data_store.REL_DB.UnregisterMessageHandler(timeout=60)
def testMessageHandlerRequests(self):
requests = [
rdf_objects.MessageHandlerRequest(client_id="C.1000000000000000",
handler_name="Testhandler",
request_id=i * 100,
request=rdfvalue.RDFInteger(i))
for i in range(5)
]
self.db.WriteMessageHandlerRequests(requests)
read = self.db.ReadMessageHandlerRequests()
for r in read:
self.assertTrue(r.timestamp)
r.timestamp = None
self.assertEqual(sorted(read, key=lambda req: req.request_id),
requests)
self.db.DeleteMessageHandlerRequests(requests[:2])
self.db.DeleteMessageHandlerRequests(requests[4:5])
read = self.db.ReadMessageHandlerRequests()
self.assertLen(read, 2)
read = sorted(read, key=lambda req: req.request_id)
for r in read:
r.timestamp = None
self.assertEqual(requests[2:4], read)
self.db.DeleteMessageHandlerRequests(read)
def testEnrollmentHandler(self):
self._ClearClient()
# First 406 queues an EnrolmentRequest.
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 406)
# Send it to the server.
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 406)
self.assertLen(self.messages, 1)
self.assertEqual(self.messages[0].session_id.Basename(),
"E:%s" % ca_enroller.EnrolmentHandler.handler_name)
request = rdf_objects.MessageHandlerRequest(
client_id=self.messages[0].source.Basename(),
handler_name="Enrol",
request_id=12345,
request=self.messages[0].payload)
handler = ca_enroller.EnrolmentHandler(token=self.token)
handler.ProcessMessages([request])
# The next client communication should give a 200.
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 200)
def _PushHandlerMessage(self, message):
"""Pushes a message that goes to a message handler."""
# We only accept messages of type MESSAGE.
if message.type != rdf_flows.GrrMessage.Type.MESSAGE:
raise ValueError("Unexpected message type: %s" % type(message))
if not message.session_id:
raise ValueError("Message without session_id: %s" % message)
# Assume the message is authenticated and comes from this client.
message.source = self.client_id
message.auth_state = "AUTHENTICATED"
session_id = message.session_id
handler_name = message_handlers.session_id_map.get(session_id, None)
if handler_name is None:
raise ValueError("Unknown well known session id in msg %s" %
message)
logging.info("Running message handler: %s", handler_name)
handler_cls = handler_registry.handler_name_map.get(handler_name)
handler_request = rdf_objects.MessageHandlerRequest(
client_id=self.client_id,
handler_name=handler_name,
request_id=message.response_id,
request=message.payload)
handler_cls().ProcessMessages([handler_request])
def _PushHandlerMessage(self, message):
"""Pushes a message that goes to a message handler."""
# We only accept messages of type MESSAGE.
if message.type != rdf_flows.GrrMessage.Type.MESSAGE:
raise ValueError("Unexpected message type: %s" % type(message))
if not message.session_id:
raise ValueError("Message without session_id: %s" % message)
# Assume the message is authenticated and comes from this client.
message.source = self.client_id
message.auth_state = "AUTHENTICATED"
session_id = message.session_id
if data_store.RelationalDBEnabled():
handler_name = queue_manager.session_id_map.get(session_id, None)
if handler_name is None:
raise ValueError("Unknown well known session id in msg %s" %
message)
logging.info("Running message handler: %s", handler_name)
handler_cls = handler_registry.handler_name_map.get(handler_name)
handler_request = rdf_objects.MessageHandlerRequest(
client_id=self.client_id.Basename(),
handler_name=handler_name,
request_id=message.response_id,
request=message.payload)
handler_cls(token=self.token).ProcessMessages([handler_request])
else:
logging.info("Running well known flow: %s", session_id)
self.well_known_flows[session_id.FlowName()].ProcessMessage(
message)
def _testProcessMessagesWellKnown(self):
worker_obj = self._TestWorker()
# Send a message to a WellKnownFlow - ClientStatsAuto.
session_id = administrative.GetClientStatsAuto.well_known_session_id
client_id = self.SetupClient(100)
if data_store.RelationalDBReadEnabled():
done = threading.Event()
def handle(l):
worker_obj._ProcessMessageHandlerRequests(l)
done.set()
data_store.REL_DB.RegisterMessageHandler(
handle, worker_obj.well_known_flow_lease_time, limit=1000)
data_store.REL_DB.WriteMessageHandlerRequests([
rdf_objects.MessageHandlerRequest(
client_id=client_id.Basename(),
handler_name="StatsHandler",
request_id=12345,
request=rdf_client_stats.ClientStats(RSS_size=1234))
])
self.assertTrue(done.wait(10))
else:
self.SendResponse(session_id,
data=rdf_client_stats.ClientStats(RSS_size=1234),
client_id=client_id,
well_known=True)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
if data_store.RelationalDBReadEnabled():
results = data_store.REL_DB.ReadClientStats(
client_id=client_id.Basename(),
min_timestamp=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0),
max_timestamp=rdfvalue.RDFDatetime.Now())
self.assertLen(results, 1)
stats = results[0]
else:
client = aff4.FACTORY.Open(client_id.Add("stats"),
token=self.token)
stats = client.Get(client.Schema.STATS)
self.assertEqual(stats.RSS_size, 1234)
# Make sure no notifications have been sent.
user = aff4.FACTORY.Open("aff4:/users/%s" % self.token.username,
token=self.token)
notifications = user.Get(user.Schema.PENDING_NOTIFICATIONS)
self.assertIsNone(notifications)
if data_store.RelationalDBReadEnabled():
data_store.REL_DB.UnregisterMessageHandler(timeout=60)
def testEnrollment(self):
"""Test the http response to unknown clients."""
self._ClearClient()
# Now communicate with the server.
self.SendToServer()
status = self.client_communicator.RunOnce()
# We expect to receive a 406 and all client messages will be tagged as
# UNAUTHENTICATED.
self.assertEqual(status.code, 406)
self.assertLen(self.messages, 10)
self.assertEqual(
self.messages[0].auth_state,
rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED)
# The next request should be an enrolling request.
self.client_communicator.RunOnce()
self.assertLen(self.messages, 11)
enrolment_messages = []
expected_id = "E:%s" % ca_enroller.EnrolmentHandler.handler_name
for m in self.messages:
if m.session_id.Basename() == expected_id:
enrolment_messages.append(m)
self.assertLen(enrolment_messages, 1)
# Now we manually run the enroll well known flow with the enrollment
# request. This will start a new flow for enrolling the client, sign the
# cert and add it to the data store.
handler = ca_enroller.EnrolmentHandler()
req = rdf_objects.MessageHandlerRequest(
client_id=self.client_id, request=enrolment_messages[0].payload)
handler.ProcessMessages([req])
# The next client communication should be enrolled now.
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 200)
# There should be a cert for the client right now.
md = data_store.REL_DB.ReadClientMetadata(self.client_id)
self.assertTrue(md.certificate)
# Now communicate with the server once again.
self.SendToServer()
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 200)
def PushToStateQueue(self, manager, message, **kw):
"""Push given message to the state queue."""
# Assume the client is authorized
message.auth_state = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED
# Update kw args
for k, v in kw.items():
setattr(message, k, v)
# Handle well known flows
if message.request_id == 0:
# Well known flows only accept messages of type MESSAGE.
if message.type == rdf_flows.GrrMessage.Type.MESSAGE:
# Assume the message is authenticated and comes from this client.
message.source = self.client_id
message.auth_state = "AUTHENTICATED"
session_id = message.session_id
if session_id:
handler_name = queue_manager.session_id_map.get(session_id)
if handler_name:
logging.info("Running message handler: %s",
handler_name)
handler_cls = handler_registry.handler_name_map.get(
handler_name)
handler_request = rdf_objects.MessageHandlerRequest(
client_id=self.client_id.Basename(),
handler_name=handler_name,
request_id=message.response_id,
request=message.payload)
handler_cls(
token=self.token).ProcessMessages(handler_request)
else:
logging.info("Running well known flow: %s", session_id)
self.well_known_flows[
session_id.FlowName()].ProcessMessage(message)
return
manager.QueueResponse(message)
def testMessageHandlerRequestLeasing(self):
requests = [
rdf_objects.MessageHandlerRequest(client_id="C.1000000000000000",
handler_name="Testhandler",
request_id=i * 100,
request=rdfvalue.RDFInteger(i))
for i in range(10)
]
lease_time = rdfvalue.Duration("5m")
with test_lib.FakeTime(
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)):
self.db.WriteMessageHandlerRequests(requests)
leased = queue.Queue()
self.db.RegisterMessageHandler(leased.put, lease_time, limit=5)
got = []
while len(got) < 10:
try:
l = leased.get(True, timeout=6)
except queue.Empty:
self.fail(
"Timed out waiting for messages, expected 10, got %d" %
len(got))
self.assertLessEqual(len(l), 5)
for m in l:
self.assertEqual(m.leased_by, utils.ProcessIdString())
self.assertGreater(m.leased_until, rdfvalue.RDFDatetime.Now())
self.assertLess(m.timestamp, rdfvalue.RDFDatetime.Now())
m.leased_by = None
m.leased_until = None
m.timestamp = None
got += l
self.db.DeleteMessageHandlerRequests(got)
got.sort(key=lambda req: req.request_id)
self.assertEqual(requests, got)
def testMessageHandlers(self):
client_id = self.SetupClient(100)
done = threading.Event()
def handle(l):
worker_lib.ProcessMessageHandlerRequests(l)
done.set()
data_store.REL_DB.RegisterMessageHandler(
handle,
worker_lib.GRRWorker.message_handler_lease_time,
limit=1000)
data_store.REL_DB.WriteMessageHandlerRequests([
rdf_objects.MessageHandlerRequest(
client_id=client_id,
handler_name="StatsHandler",
request_id=12345,
request=rdf_client_stats.ClientStats(RSS_size=1234))
])
self.assertTrue(done.wait(10))
results = data_store.REL_DB.ReadClientStats(
client_id=client_id,
min_timestamp=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0),
max_timestamp=rdfvalue.RDFDatetime.Now())
self.assertLen(results, 1)
stats = results[0]
self.assertEqual(stats.RSS_size, 1234)
data_store.REL_DB.UnregisterMessageHandler(timeout=60)
# Make sure there are no leftover requests.
self.assertEqual(data_store.REL_DB.ReadMessageHandlerRequests(), [])
def ReceiveMessages(self, client_id: str,
messages: Iterable[rdf_flows.GrrMessage]):
"""Receives and processes the messages.
For each message we update the request object, and place the
response in that request's queue. If the request is complete, we
send a message to the worker.
Args:
client_id: The client which sent the messages.
messages: A list of GrrMessage RDFValues.
"""
now = time.time()
unprocessed_msgs = []
worker_message_handler_requests = []
frontend_message_handler_requests = []
dropped_count = 0
msgs_by_session_id = collection.Group(messages, lambda m: m.session_id)
for session_id, msgs in msgs_by_session_id.items():
for msg in msgs:
if (msg.auth_state != msg.AuthorizationState.AUTHENTICATED
and msg.session_id != self.unauth_allowed_session_id):
dropped_count += 1
continue
session_id_str = str(session_id)
if session_id_str in message_handlers.session_id_map:
request = rdf_objects.MessageHandlerRequest(
client_id=msg.source.Basename(),
handler_name=message_handlers.
session_id_map[session_id],
request_id=msg.response_id or random.UInt32(),
request=msg.payload)
if request.handler_name in self._SHORTCUT_HANDLERS:
frontend_message_handler_requests.append(request)
else:
worker_message_handler_requests.append(request)
elif session_id_str in self.legacy_well_known_session_ids:
logging.debug(
"Dropping message for legacy well known session id %s",
session_id)
else:
unprocessed_msgs.append(msg)
if dropped_count:
logging.info("Dropped %d unauthenticated messages for %s",
dropped_count, client_id)
if unprocessed_msgs:
flow_responses = []
for message in unprocessed_msgs:
try:
flow_responses.append(
rdf_flow_objects.FlowResponseForLegacyResponse(
message))
except ValueError as e:
logging.warning(
"Failed to parse legacy FlowResponse:\n%s\n%s", e,
message)
data_store.REL_DB.WriteFlowResponses(flow_responses)
for msg in unprocessed_msgs:
if msg.type == rdf_flows.GrrMessage.Type.STATUS:
stat = rdf_flows.GrrStatus(msg.payload)
if stat.status == rdf_flows.GrrStatus.ReturnedStatus.CLIENT_KILLED:
# A client crashed while performing an action, fire an event.
crash_details = rdf_client.ClientCrash(
client_id=client_id,
session_id=msg.session_id,
backtrace=stat.backtrace,
crash_message=stat.error_message,
nanny_status=stat.nanny_status,
timestamp=rdfvalue.RDFDatetime.Now())
events.Events.PublishEvent("ClientCrash",
crash_details,
token=self.token)
if worker_message_handler_requests:
data_store.REL_DB.WriteMessageHandlerRequests(
worker_message_handler_requests)
if frontend_message_handler_requests:
worker_lib.ProcessMessageHandlerRequests(
frontend_message_handler_requests)
logging.debug("Received %s messages from %s in %s sec", len(messages),
client_id,
time.time() - now)
def ReceiveMessagesRelationalFlows(self, client_id, messages):
"""Receives and processes messages for flows stored in the relational db.
Args:
client_id: The client which sent the messages.
messages: A list of GrrMessage RDFValues.
"""
now = time.time()
unprocessed_msgs = []
message_handler_requests = []
dropped_count = 0
for session_id, msgs in iteritems(
collection.Group(messages, operator.attrgetter("session_id"))):
# Remove and handle messages to WellKnownFlows
leftover_msgs = self.HandleWellKnownFlows(msgs)
for msg in leftover_msgs:
if (msg.auth_state != msg.AuthorizationState.AUTHENTICATED
and msg.session_id != self.unauth_allowed_session_id):
dropped_count += 1
continue
if session_id in queue_manager.session_id_map:
message_handler_requests.append(
rdf_objects.MessageHandlerRequest(
client_id=msg.source.Basename(),
handler_name=queue_manager.
session_id_map[session_id],
request_id=msg.response_id,
request=msg.payload))
elif session_id in self.legacy_well_known_session_ids:
logging.debug(
"Dropping message for legacy well known session id %s",
session_id)
else:
unprocessed_msgs.append(msg)
if dropped_count:
logging.info("Dropped %d unauthenticated messages for %s",
dropped_count, client_id)
if unprocessed_msgs:
flow_responses = []
for message in unprocessed_msgs:
flow_responses.append(
rdf_flow_objects.FlowResponseForLegacyResponse(message))
data_store.REL_DB.WriteFlowResponses(flow_responses)
for msg in unprocessed_msgs:
if msg.type == rdf_flows.GrrMessage.Type.STATUS:
stat = rdf_flows.GrrStatus(msg.payload)
if stat.status == rdf_flows.GrrStatus.ReturnedStatus.CLIENT_KILLED:
# A client crashed while performing an action, fire an event.
crash_details = rdf_client.ClientCrash(
client_id=client_id,
session_id=msg.session_id,
backtrace=stat.backtrace,
crash_message=stat.error_message,
nanny_status=stat.nanny_status,
timestamp=rdfvalue.RDFDatetime.Now())
events.Events.PublishEvent("ClientCrash",
crash_details,
token=self.token)
if message_handler_requests:
data_store.REL_DB.WriteMessageHandlerRequests(
message_handler_requests)
logging.debug("Received %s messages from %s in %s sec", len(messages),
client_id,
time.time() - now)
def Flush(self):
"""Writes the changes in this object to the datastore."""
if data_store.RelationalDBReadEnabled(category="message_handlers"):
message_handler_requests = []
leftover_responses = []
for r, timestamp in self.response_queue:
if r.request_id == 0 and r.session_id in session_id_map:
message_handler_requests.append(
rdf_objects.MessageHandlerRequest(
client_id=r.source and r.source.Basename(),
handler_name=session_id_map[r.session_id],
request_id=r.response_id,
request=r.payload))
else:
leftover_responses.append((r, timestamp))
if message_handler_requests:
data_store.REL_DB.WriteMessageHandlerRequests(
message_handler_requests)
self.response_queue = leftover_responses
self.data_store.StoreRequestsAndResponses(
new_requests=self.request_queue,
new_responses=self.response_queue,
requests_to_delete=self.requests_to_delete)
# We need to make sure that notifications are written after the requests so
# we flush after writing all requests and only notify afterwards.
mutation_pool = self.data_store.GetMutationPool()
with mutation_pool:
if data_store.RelationalDBReadEnabled(category="client_messages"):
if self.client_messages_to_delete:
data_store.REL_DB.DeleteClientMessages(
list(itervalues(self.client_messages_to_delete)))
else:
messages_by_queue = utils.GroupBy(
list(itervalues(self.client_messages_to_delete)),
lambda request: request.queue)
for queue, messages in iteritems(messages_by_queue):
self.Delete(queue, messages, mutation_pool=mutation_pool)
if self.new_client_messages:
for timestamp, messages in iteritems(
utils.GroupBy(self.new_client_messages,
lambda x: x[1])):
self.Schedule([x[0] for x in messages],
timestamp=timestamp,
mutation_pool=mutation_pool)
if self.notifications:
for notification in itervalues(self.notifications):
self.NotifyQueue(notification, mutation_pool=mutation_pool)
mutation_pool.Flush()
self.request_queue = []
self.response_queue = []
self.requests_to_delete = []
self.client_messages_to_delete = {}
self.notifications = {}
self.new_client_messages = []
def testEnrollment(self):
"""Test the http response to unknown clients."""
self._ClearClient()
# Now communicate with the server.
self.SendToServer()
status = self.client_communicator.RunOnce()
# We expect to receive a 406 and all client messages will be tagged as
# UNAUTHENTICATED.
self.assertEqual(status.code, 406)
self.assertLen(self.messages, 10)
self.assertEqual(
self.messages[0].auth_state,
rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED)
# The next request should be an enrolling request.
status = self.client_communicator.RunOnce()
self.assertLen(self.messages, 11)
enrolment_messages = []
for m in self.messages:
if m.session_id == ca_enroller.Enroler.well_known_session_id:
enrolment_messages.append(m)
self.assertLen(enrolment_messages, 1)
# Now we manually run the enroll well known flow with the enrollment
# request. This will start a new flow for enrolling the client, sign the
# cert and add it to the data store.
if data_store.AFF4Enabled():
flow_obj = ca_enroller.Enroler(
ca_enroller.Enroler.well_known_session_id,
mode="rw",
token=self.token)
flow_obj.ProcessMessage(enrolment_messages[0])
else:
handler = ca_enroller.EnrolmentHandler()
req = rdf_objects.MessageHandlerRequest(
client_id=self.client_id,
request=enrolment_messages[0].payload)
handler.ProcessMessages([req])
# The next client communication should be enrolled now.
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 200)
# There should be a cert for the client right now.
if data_store.AFF4Enabled():
client = aff4.FACTORY.Create(self.client_cn,
aff4_grr.VFSGRRClient,
mode="rw",
token=self.token)
self.assertTrue(client.Get(client.Schema.CERT))
else:
md = data_store.REL_DB.ReadClientMetadata(self.client_id)
self.assertTrue(md.certificate)
# Now communicate with the server once again.
self.SendToServer()
status = self.client_communicator.RunOnce()
self.assertEqual(status.code, 200)
