def testCorrectlyExportsSingleValue(self): with self.pool: self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFString("foo"), source=self.client_id), mutation_pool=self.pool) chunks = self.ProcessPlugin() self.assertListEqual(chunks, [ "Start: aff4:/foo/bar", "Values of type: RDFString", "First pass: foo (source=%s)" % self.client_id, "Second pass: foo (source=%s)" % self.client_id, "Finish: aff4:/foo/bar" ]) # pyformat: disable
def testSendGrrMessage_allRetriesFailed(self, mock_conn): client_id = "C.0123456789abcdef" grr_message = rdf_flows.GrrMessage() mock_conn.outgoing.InsertMessage.side_effect = grpc.RpcError("Foo") with test_lib.ConfigOverrider({ "Server.fleetspeak_send_retry_attempts": 5, "Server.fleetspeak_send_retry_sleep_time_secs": 0, }): with self.assertRaisesRegex(grpc.RpcError, "Foo"): fleetspeak_utils.SendGrrMessageThroughFleetspeak( client_id, grr_message) self.assertEqual(mock_conn.outgoing.InsertMessage.call_count, 5)
def SendReply(value, session_id=None, message_type=rdf_flows.GrrMessage.Type.MESSAGE): if message_type != rdf_flows.GrrMessage.Type.MESSAGE: return if str(session_id) in self.wkfs: message = rdf_flows.GrrMessage(name=action_cls.__name__, payload=value, auth_state="AUTHENTICATED", session_id=session_id) self.wkfs[str(session_id)].ProcessMessage(message) else: responses.append(value)
def SendOKStatus(self, response_id, session_id): """Send a message to the flow.""" message = rdf_flows.GrrMessage( request_id=1, response_id=response_id, session_id=session_id, type=rdf_flows.GrrMessage.Type.STATUS, auth_state=rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED) status = rdf_flows.GrrStatus( status=rdf_flows.GrrStatus.ReturnedStatus.OK) message.payload = status self.SendMessage(message)
def testCPUAccounting(self): with contextlib.ExitStack() as stack: server_cpu_time = 1.0 server_sys_time = 1.1 stack.enter_context( mock.patch.object(communication, "TotalServerCpuTime", lambda: server_cpu_time)) stack.enter_context( mock.patch.object(communication, "TotalServerSysTime", lambda: server_sys_time)) process_cpu_time = 1.2 process_sys_time = 1.3 class FakeProcess(object): def __init__(self, pid=None): pass def cpu_times(self): # pylint: disable=invalid-name return collections.namedtuple( "pcputimes", ["user", "system"])(process_cpu_time, process_sys_time) stack.enter_context( mock.patch.object(psutil, "Process", FakeProcess)) class _ProgressAction(ProgressAction): def Run(self, *args): super().Run(*args) nonlocal server_cpu_time, server_sys_time server_cpu_time = 42.0 server_sys_time = 43.0 nonlocal process_cpu_time, process_sys_time process_cpu_time = 10.0 process_sys_time = 11.0 message = rdf_flows.GrrMessage(name="ProgressAction", runtime_limit_us=0) worker = mock.MagicMock() action = _ProgressAction(worker) action.SendReply = mock.MagicMock() action.Execute(message) self.assertEqual(action.SendReply.call_count, 1) self.assertAlmostEqual( action.SendReply.call_args[0][0].cpu_time_used.user_cpu_time, 42.0 - 1.0 + 10.0 - 1.2) self.assertAlmostEqual( action.SendReply.call_args[0][0].cpu_time_used.system_cpu_time, 43.0 - 1.1 + 11.0 - 1.3)
def testReceiveMessages(self): fs_server = fleetspeak_frontend_server.GRRFSServer() client_id = "C.1234567890123456" flow_id = "12345678" data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_enabled=True) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow_id, create_time=rdfvalue.RDFDatetime.Now()) data_store.REL_DB.WriteFlowObject(rdf_flow) flow_request = rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=1) data_store.REL_DB.WriteFlowRequests([flow_request]) session_id = "%s/%s" % (client_id, flow_id) fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id) fs_messages = [] for i in range(1, 10): grr_message = rdf_flows.GrrMessage(request_id=1, response_id=i + 1, session_id=session_id, payload=rdfvalue.RDFInteger(i)) fs_message = fs_common_pb2.Message( message_type="GrrMessage", source=fs_common_pb2.Address(client_id=fs_client_id, service_name=FS_SERVICE_NAME)) fs_message.data.Pack(grr_message.AsPrimitiveProto()) fs_messages.append(fs_message) with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)): for fs_message in fs_messages: fs_server.Process(fs_message, None) # Ensure the last-ping timestamp gets updated. client_data = data_store.REL_DB.MultiReadClientMetadata([client_id]) self.assertEqual(client_data[client_id].ping, rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)) flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses( client_id, flow_id) self.assertLen(flow_data, 1) stored_flow_request, flow_responses = flow_data[0] self.assertEqual(stored_flow_request, flow_request) self.assertLen(flow_responses, 9)
def testEventNotification(self): """Test that events are sent to listeners.""" TestListener.received_events = [] 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") events.Events.PublishEvent("TestEvent", event, token=self.token) # Make sure the source is correctly propagated. self.assertEqual(TestListener.received_events[0], event)
def testUsesDefaultClientURNIfGrrMessageHasNoSource(self): with self.pool: self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFString("foo"), source=None), mutation_pool=self.pool) chunks = self.ProcessPlugin( source_urn=rdf_client.ClientURN("C.1111222233334444")) self.assertListEqual(chunks, [ "Start: aff4:/foo/bar", "Values of type: RDFString", "First pass: foo (source=aff4:/C.1111222233334444)", "Second pass: foo (source=aff4:/C.1111222233334444)", "Finish: aff4:/foo/bar" ]) # pyformat: disable
def GRRMessageFromClientActionRequest(request): stub = action_registry.ACTION_STUB_BY_ID[request.action_identifier] name = compatibility.GetName(stub) return rdf_flows.GrrMessage( session_id="%s/%s" % (request.client_id, request.flow_id), name=name, request_id=request.request_id, queue=rdf_client.ClientURN(request.client_id), payload=request.action_args, cpu_limit=request.cpu_limit_ms / 1000.0, network_bytes_limit=request.network_bytes_limit, # Legacy clients will fail if the task id is not set. # TODO(amoser): Remove task ids after April 2021. generate_task_id=True)
def testSizeLimitedQueueOverflow(self): msg_a = rdf_flows.GrrMessage(name="A") msg_b = rdf_flows.GrrMessage(name="B") msg_c = rdf_flows.GrrMessage(name="C") msg_d = rdf_flows.GrrMessage(name="D") queue = comms.SizeLimitedQueue(maxsize=3 * len(msg_a.SerializeToString()), heart_beat_cb=lambda: None) queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False) queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False) queue.Put(msg_c, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False) with self.assertRaises(Queue.Full): queue.Put(msg_d, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False)
def testSizeLimitedQueue(self): limited_queue = comms.SizeLimitedQueue( maxsize=10000000, heart_beat_cb=lambda: None) msg_a = rdf_flows.GrrMessage(name="A") msg_b = rdf_flows.GrrMessage(name="B") msg_c = rdf_flows.GrrMessage(name="C") for _ in range(10): limited_queue.Put(msg_a) limited_queue.Put(msg_b) limited_queue.Put(msg_c) result = limited_queue.GetMessages() self.assertCountEqual(list(result.job), [msg_c] * 10 + [msg_a, msg_b] * 10) # Tests a partial Get(). for _ in range(7): limited_queue.Put(msg_a) limited_queue.Put(msg_b) limited_queue.Put(msg_c) result = limited_queue.GetMessages( soft_size_limit=len(msg_a.SerializeToString()) * 5 - 1) self.assertLen(list(result.job), 5) for _ in range(3): limited_queue.Put(msg_a) limited_queue.Put(msg_b) limited_queue.Put(msg_c) # Append the remaining messages to the same result. result.job.Extend(limited_queue.GetMessages().job) self.assertCountEqual(list(result.job), [msg_c] * 10 + [msg_a, msg_b] * 10)
def testOldStyleHuntIDsDontError(self): """Tests receiving messages with old style hunt ids.""" client_id = "C.1234567890123456" hunt_id = "aff4:/hunts/H:3479C8EA/C.917bf16e123bf731/H:5E69190A/H:A21A1AA2" messages = [ rdf_flows.GrrMessage(request_id=1, response_id=1, session_id=hunt_id, auth_state="AUTHENTICATED", payload=rdfvalue.RDFInteger(1)) ] # This must not raise even though we don't generate session ids like the one # above anymore. ReceiveMessages(client_id, messages)
def testProcessesMultipleYaraProcessDumpInformationCorrectly(self): yara_dump = rdf_memory.YaraProcessDumpResponse(dumped_processes=[ rdf_memory.YaraProcessDumpInformation(dump_files=[ rdf_paths.PathSpec(path="my_proc_123_f0_fa.tmp", pathtype=TMPFILE), rdf_paths.PathSpec(path="my_proc_123_fa_104.tmp", pathtype=TMPFILE) ]), rdf_memory.YaraProcessDumpInformation(dump_files=[ rdf_paths.PathSpec(path="foobar_456_f0_fa.tmp", pathtype=TMPFILE), ]) ]) stat_entry_1 = rdf_client_fs.StatEntry( pathspec=rdf_paths.PathSpec( path="my_proc_123_f0_fa.tmp", pathtype=TMPFILE)) stat_entry_2 = rdf_client_fs.StatEntry( pathspec=rdf_paths.PathSpec( path="my_proc_123_fa_104.tmp", pathtype=TMPFILE)) stat_entry_3 = rdf_client_fs.StatEntry( pathspec=rdf_paths.PathSpec( path="foobar_456_f0_fa.tmp", pathtype=TMPFILE)) m0 = rdf_flows.GrrMessage(source=self.client_id, payload=yara_dump) m1 = rdf_flows.GrrMessage(source=self.client_id, payload=stat_entry_1) m2 = rdf_flows.GrrMessage(source=self.client_id, payload=stat_entry_2) m3 = rdf_flows.GrrMessage(source=self.client_id, payload=stat_entry_3) self._ProcessValuesWithPlugin([m0, m1, m2, m3]) self.assertEqual(self.plugin.OutputMemoryDump.call_count, 2) self.plugin.OutputMemoryDump.assert_any_call( rdf_memory.YaraProcessDumpInformation(dump_files=[ rdf_paths.PathSpec(path="my_proc_123_f0_fa.tmp", pathtype=TMPFILE), rdf_paths.PathSpec(path="my_proc_123_fa_104.tmp", pathtype=TMPFILE) ]), self.client_id) self.plugin.OutputMemoryDump.assert_any_call( rdf_memory.YaraProcessDumpInformation(dump_files=[ rdf_paths.PathSpec(path="foobar_456_f0_fa.tmp", pathtype=TMPFILE), ]), self.client_id)
def SendMessages(self, response_ids, session_id, authenticated=True): """Send messages to the flow.""" for response_id in response_ids: message = rdf_flows.GrrMessage( request_id=1, response_id=response_id, session_id=session_id) blob = rdf_protodict.DataBlob() blob.SetValue(response_id) message.payload = blob if authenticated: auth_state = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED message.auth_state = auth_state self.SendMessage(message)
def __iter__(self): """An iterator which returns all the responses in order.""" old_response_id = None action_registry = server_stubs.ClientActionStub.classes expected_response_classes = [] is_client_request = False # This is the client request so this response packet was sent by a client. if self.request.HasField("request"): is_client_request = True client_action_name = self.request.request.name if client_action_name not in action_registry: raise RuntimeError("Got unknown client action: %s." % client_action_name) expected_response_classes = action_registry[ client_action_name].out_rdfvalues for message in self._responses: self.message = rdf_flows.GrrMessage(message) # Handle retransmissions if self.message.response_id == old_response_id: continue else: old_response_id = self.message.response_id if self.message.type == self.message.Type.MESSAGE: if is_client_request: # Let's do some verification for requests that came from clients. if not expected_response_classes: raise RuntimeError( "Client action %s does not specify out_rdfvalue." % client_action_name) else: args_rdf_name = self.message.args_rdf_name if not args_rdf_name: raise RuntimeError( "Deprecated message format received: " "args_rdf_name is None.") elif args_rdf_name not in [ x.__name__ for x in expected_response_classes ]: raise RuntimeError( "Response type was %s but expected %s for %s." % (args_rdf_name, expected_response_classes, client_action_name)) yield self.message.payload
def testDestroyFlowStates(self): """Check that we can efficiently destroy the flow's request queues.""" session_id = rdfvalue.SessionID(flow_name="test2") request = rdf_flow_runner.RequestState( id=1, client_id=test_lib.TEST_CLIENT_ID, next_state="TestState", session_id=session_id) with queue_manager.QueueManager(token=self.token) as manager: manager.QueueRequest(request) manager.QueueResponse( rdf_flows.GrrMessage(request_id=1, response_id=1, session_id=session_id)) # 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) # Read the response directly. responses = data_store.DB.ReadResponsesForRequestId(session_id, 1) self.assertEqual(len(responses), 1) response = responses[0] self.assertEqual(response.request_id, 1) self.assertEqual(response.response_id, 1) self.assertEqual(response.session_id, session_id) with queue_manager.QueueManager(token=self.token) as manager: manager.DestroyFlowStates(session_id) all_requests = list(manager.FetchRequestsAndResponses(session_id)) self.assertEqual(len(all_requests), 0) # Check that the response is gone. responses = data_store.DB.ReadResponsesForRequestId(session_id, 1) self.assertEqual(len(responses), 0) # Ensure the rows are gone from the data store. Some data stores # don't store the queues in that way but there is no harm in # checking. self.assertEqual( data_store.DB.ResolveRow(session_id.Add("state/request:00000001")), []) self.assertEqual(data_store.DB.ResolveRow(session_id.Add("state")), [])
def testClientSummaryModalIsShownWhenClientInfoButtonClicked(self): client_id = self.SetupClient(0) h = self.CreateSampleHunt() with data_store.DB.GetMutationPool() as pool: h.ResultCollection().Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFString("foo-result"), source=client_id), mutation_pool=pool) self.Open("/#/hunts/%s/results" % h.urn.Basename()) self.Click("css=td:contains('%s') button:has(.glyphicon-info-sign)" % client_id.Basename()) self.WaitUntil( self.IsElementPresent, "css=.modal-dialog:contains('Client %s')" % client_id.Basename())
def testFastPoll(self): """Test fast poll settings propagated to status results.""" for i in range(10): message = rdf_flows.GrrMessage( name="MockAction", session_id=self.session_id.Basename() + str(i), auth_state=rdf_flows.GrrMessage.AuthorizationState. UNAUTHENTICATED, request_id=1, require_fastpoll=i % 2, generate_task_id=True) self.context.HandleMessage(message) message_list = self.context.Drain(max_size=1000000).job self.assertEqual(len(message_list), 10) self.assertItemsEqual([m.require_fastpoll for m in message_list], [0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
def _RunFileFinder( args: rdf_file_finder.FileFinderArgs ) -> List[rdf_file_finder.FileFinderResult]: results = [] def SendReply(rdf_value, *args, **kwargs): del args, kwargs # Unused. results.append(rdf_value) ff = vfs_file_finder.VfsFileFinder() ff.grr_worker = mock.MagicMock() ff.SendReply = SendReply ff.message = rdf_flows.GrrMessage(payload=args) ff.Run(args) return results
def _FillInStubResults(self): results = implementation.GRRHunt.ResultCollectionForHID( self.hunt.urn, token=self.token) result = results[0] with data_store.DB.GetMutationPool() as pool: for i in range(self.handler.MAX_RECORDS_TO_CHECK): wrong_result = rdf_flows.GrrMessage( payload=rdfvalue.RDFString("foo/bar"), age=(result.age - (self.handler.MAX_RECORDS_TO_CHECK - i + 1) * rdfvalue.Duration("1s")), source=self.client_id) results.Add( wrong_result, timestamp=wrong_result.age, mutation_pool=pool) return result
def ExecuteAction(self, action_cls, arg=None, grr_worker=None, session_id=None): message = rdf_flows.GrrMessage(name=action_cls.__name__, payload=arg, auth_state="AUTHENTICATED", session_id=session_id) self.results = [] action = self._GetActionInstance(action_cls, grr_worker=grr_worker) action.Execute(message) return self.results
def GenerateStatusMessage(self, message, response_id=1): cpu_time_used = rdf_client_stats.CpuSeconds( user_cpu_time=self.user_cpu_usage.next(), system_cpu_time=self.system_cpu_usage.next()) network_bytes_sent = self.network_usage.next() return rdf_flows.GrrMessage( session_id=message.session_id, name=message.name, response_id=response_id, request_id=message.request_id, payload=rdf_flows.GrrStatus( status=rdf_flows.GrrStatus.ReturnedStatus.OK, cpu_time_used=cpu_time_used, network_bytes_sent=network_bytes_sent), type=rdf_flows.GrrMessage.Type.STATUS)
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 the following tag. NOTE: supported in REL_DB implementation only. Raises: ValueError: If responses is not of the correct type. """ del tag if not isinstance(response, rdfvalue.RDFValue): raise ValueError("SendReply can only send a Semantic Value") # Only send the reply if we have a parent, indicated by knowing our parent's # request state. if self.runner_args.request_state.session_id: request_state = self.runner_args.request_state request_state.response_count += 1 # Make a response message msg = rdf_flows.GrrMessage( session_id=request_state.session_id, request_id=request_state.id, response_id=request_state.response_count, auth_state=rdf_flows.GrrMessage.AuthorizationState. AUTHENTICATED, type=rdf_flows.GrrMessage.Type.MESSAGE, payload=response, args_rdf_name=response.__class__.__name__, args_age=int(response.age)) # Queue the response now self.queue_manager.QueueResponse(msg) if self.runner_args.write_intermediate_results: self.QueueReplyForResultCollection(response) else: # Only write the reply to the collection if we are the parent flow. self.QueueReplyForResultCollection(response)
def testResponsesAnyRequestTriggerClientMessageDeletion(self): # Write a flow that is waiting for request #2. client_id, flow_id = self._SetupClientAndFlow( next_request_to_process=2) msg = rdf_flows.GrrMessage(queue=client_id, generate_task_id=True) self.db.WriteClientMessages([msg]) self.assertTrue(self.db.ReadClientMessages(client_id)) self._WriteRequestAndCompleteResponses(client_id, flow_id, request_id=1, num_responses=3, task_id=msg.task_id) self.assertFalse(self.db.ReadClientMessages(client_id))
def testTransferNetworkByteLimitError(self): message = rdf_flows.GrrMessage( name="TransferBuffer", payload=self.buffer_ref, network_bytes_limit=300, generate_task_id=True) # We just get a client alert and a status message back. responses = self.transfer_buf.HandleMessage(message) client_alert = responses[0].payload self.assertIn("Network limit exceeded", str(client_alert)) status = responses[1].payload self.assertIn("Action exceeded network send limit", str(status.backtrace)) self.assertEqual(status.status, rdf_flows.GrrStatus.ReturnedStatus.NETWORK_LIMIT_EXCEEDED)
def SendMessages(self, response_ids, session_id, authenticated=True, args_rdf_name="DataBlob"): """Send messages to the flow.""" for response_id in response_ids: message = rdf_flows.GrrMessage(request_id=1, response_id=response_id, session_id=session_id, args_rdf_name=args_rdf_name) if authenticated: auth_state = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED message.auth_state = auth_state self.SendMessage(message)
def _HandleMockAction(self, message): """Handles the action in case it's a mock.""" responses = getattr(self, message.name)(message.payload) ret = [] for i, r in enumerate(responses): ret.append( rdf_flows.GrrMessage(session_id=message.session_id, request_id=message.request_id, task_id=message.task_id, name=message.name, response_id=i + 1, payload=r, type=rdf_flows.GrrMessage.Type.MESSAGE)) ret.append( self.GenerateStatusMessage(message, response_id=len(ret) + 1)) return ret
def AsLegacyGrrMessage(self): payload = rdf_flows.GrrStatus(status=inv_status_map[self.status]) if self.error_message: payload.error_message = self.error_message if self.backtrace: payload.backtrace = self.backtrace if self.cpu_time_used: payload.cpu_time_used = self.cpu_time_used if self.network_bytes_sent: payload.network_bytes_sent = self.network_bytes_sent return rdf_flows.GrrMessage(session_id="%s/flows/%s" % (self.client_id, self.flow_id), request_id=self.request_id, response_id=self.response_id, type="STATUS", timestamp=self.timestamp, payload=payload)
def testSendGrrMessage_retrySuccess(self, mock_conn, mock_sleep): client_id = "C.0123456789abcdef" grr_message = rdf_flows.GrrMessage() def MockInsertMessage(*args): nonlocal mock_conn mock_conn.outgoing.InsertMessage.side_effect = None raise grpc.RpcError("Foo") mock_conn.outgoing.InsertMessage.side_effect = MockInsertMessage with test_lib.ConfigOverrider({ "Server.fleetspeak_send_retry_attempts": 2, "Server.fleetspeak_send_retry_sleep_time_secs": 3, }): fleetspeak_utils.SendGrrMessageThroughFleetspeak(client_id, grr_message) self.assertEqual(mock_conn.outgoing.InsertMessage.call_count, 2) self.assertEqual(mock_sleep.call_count, 1) mock_sleep.assert_called_with(3)
def __init__(self, grr_worker=None): """Initializes the action plugin. Args: grr_worker: The grr client worker object which may be used to e.g. send new actions on. """ self.grr_worker = grr_worker self.response_id = INITIAL_RESPONSE_ID self.cpu_used = None self.nanny_controller = None self.status = rdf_flows.GrrStatus( status=rdf_flows.GrrStatus.ReturnedStatus.OK) self._last_gc_run = rdfvalue.RDFDatetime.Now() self._gc_frequency = config.CONFIG["Client.gc_frequency"] self.proc = psutil.Process() self.cpu_start = self.proc.cpu_times() self.cpu_limit = rdf_flows.GrrMessage().cpu_limit