def testCorrectlyExportsFourValuesOfTwoDifferentTypes(self):
with self.pool:
self.collection.Add(rdf_flows.GrrMessage(
payload=rdfvalue.RDFString("foo"), source=self.client_id),
mutation_pool=self.pool)
self.collection.Add(rdf_flows.GrrMessage(
payload=rdfvalue.RDFInteger(42), source=self.client_id),
mutation_pool=self.pool)
self.collection.Add(rdf_flows.GrrMessage(
payload=rdfvalue.RDFString("bar"), source=self.client_id),
mutation_pool=self.pool)
self.collection.Add(rdf_flows.GrrMessage(
payload=rdfvalue.RDFInteger(43), source=self.client_id),
mutation_pool=self.pool)
chunks = self.ProcessPlugin()
self.assertListEqual(chunks, [
"Start: aff4:/foo/bar", "Values of type: RDFInteger",
"First pass: 42 (source=%s)" % self.client_id,
"First pass: 43 (source=%s)" % self.client_id,
"Second pass: 42 (source=%s)" % self.client_id,
"Second pass: 43 (source=%s)" % self.client_id,
"Values of type: RDFString",
"First pass: foo (source=%s)" % self.client_id,
"First pass: bar (source=%s)" % self.client_id,
"Second pass: foo (source=%s)" % self.client_id,
"Second pass: bar (source=%s)" % self.client_id,
"Finish: aff4:/foo/bar"
]) # pyformat: disable
def testResultsViewGetsAutoRefreshed(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("/")
# Ensure auto-refresh updates happen every second.
self.GetJavaScriptValue(
"grrUi.core.resultsCollectionDirective.setAutoRefreshInterval(1000);"
)
self.Click("css=a[grrtarget=hunts]")
self.Click("css=td:contains('GenericHunt')")
self.Click("css=li[heading=Results]")
self.WaitUntil(self.IsElementPresent,
"css=grr-results-collection td:contains('foo-result')")
self.WaitUntilNot(
self.IsElementPresent,
"css=grr-results-collection td:contains('bar-result')")
with data_store.DB.GetMutationPool() as pool:
h.ResultCollection().Add(rdf_flows.GrrMessage(
payload=rdfvalue.RDFString("bar-result"), source=client_id),
mutation_pool=pool)
self.WaitUntil(self.IsElementPresent,
"css=grr-results-collection td:contains('bar-result')")
def Run(self):
hunt_urn = rdfvalue.RDFURN("aff4:/hunts/H:123456")
results = implementation.GRRHunt.ResultCollectionForHID(hunt_urn)
with data_store.DB.GetMutationPool() as pool:
result = rdf_flows.GrrMessage(
payload=rdfvalue.RDFString("blah1"),
age=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1))
results.Add(result,
timestamp=result.age + rdfvalue.Duration("1s"),
mutation_pool=pool)
result = rdf_flows.GrrMessage(
payload=rdfvalue.RDFString("blah2-foo"),
age=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42))
results.Add(result,
timestamp=result.age + rdfvalue.Duration("1s"),
mutation_pool=pool)
self.Check("ListHuntResults",
args=hunt_plugin.ApiListHuntResultsArgs(hunt_id="H:123456"))
self.Check("ListHuntResults",
args=hunt_plugin.ApiListHuntResultsArgs(hunt_id="H:123456",
count=1))
self.Check("ListHuntResults",
args=hunt_plugin.ApiListHuntResultsArgs(hunt_id="H:123456",
offset=1,
count=1))
self.Check("ListHuntResults",
args=hunt_plugin.ApiListHuntResultsArgs(hunt_id="H:123456",
filter="foo"))
def testLengthOfCollectionIsCorrectWhenMultipleTypesAreUsed(self):
with self.pool:
for i in range(100):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)),
mutation_pool=self.pool)
self.assertEqual(200, len(self.collection))
def testEqualTimestampNotifications(self):
frontend_server = frontend_lib.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
message_expiry_time=100,
threadpool_prefix="notification-test")
# This schedules 10 requests.
session_id = flow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
# We pretend that the client processed all the 10 requests at once and
# sends the replies in a single http poll.
messages = [
rdf_flows.GrrMessage(
request_id=i,
response_id=1,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="test%s" % i),
auth_state="AUTHENTICATED",
generate_task_id=True) for i in range(1, 11)
]
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
statuses = [
rdf_flows.GrrMessage(request_id=i,
response_id=2,
session_id=session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS,
auth_state="AUTHENTICATED",
generate_task_id=True) for i in range(1, 11)
]
frontend_server.ReceiveMessages(self.client_id, messages + statuses)
with queue_manager.QueueManager(token=self.token) as q:
all_notifications = q.GetNotificationsByPriorityForAllShards(
rdfvalue.RDFURN("aff4:/F"))
medium_priority = rdf_flows.GrrNotification.Priority.MEDIUM_PRIORITY
medium_notifications = all_notifications[medium_priority]
my_notifications = [
n for n in medium_notifications if n.session_id == session_id
]
# There must not be more than one notification.
self.assertEqual(len(my_notifications), 1)
notification = my_notifications[0]
self.assertEqual(notification.first_queued, notification.timestamp)
self.assertEqual(notification.last_status, 10)
def testStoresEmptyGrrMessage(self):
with self.pool:
self.collection.Add(rdf_flows.GrrMessage(),
mutation_pool=self.pool)
self.assertListEqual([rdf_flows.GrrMessage.__name__],
list(self.collection.ListStoredTypes()))
def testLengthIsReportedCorrectlyForEveryType(self):
with self.pool:
for i in range(99):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)),
mutation_pool=self.pool)
for i in range(101):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)),
mutation_pool=self.pool)
self.assertEqual(
99, self.collection.LengthByType(rdfvalue.RDFInteger.__name__))
self.assertEqual(
101, self.collection.LengthByType(rdfvalue.RDFString.__name__))
def ProcessResponses(self,
plugin_args=None,
responses=None,
process_responses_separately=False):
plugin = bigquery_plugin.BigQueryOutputPlugin(
source_urn=self.results_urn,
output_base_urn=self.base_urn,
args=plugin_args,
token=self.token)
plugin.InitializeState()
messages = []
for response in responses:
messages.append(
rdf_flows.GrrMessage(source=self.client_id, payload=response))
with test_lib.FakeTime(1445995873):
with mock.patch.object(bigquery,
"GetBigQueryClient") as mock_bigquery:
if process_responses_separately:
for message in messages:
plugin.ProcessResponses([message])
else:
plugin.ProcessResponses(messages)
plugin.Flush()
return [
x[0] for x in mock_bigquery.return_value.InsertData.call_args_list
]
def _SendTerminationMessage(self, status=None):
"""This notifies the parent flow of our termination."""
if not self.runner_args.request_state.session_id:
# No parent flow, nothing to do here.
return
if status is None:
status = rdf_flows.GrrStatus()
client_resources = self.context.client_resources
user_cpu = client_resources.cpu_usage.user_cpu_time
sys_cpu = client_resources.cpu_usage.system_cpu_time
status.cpu_time_used.user_cpu_time = user_cpu
status.cpu_time_used.system_cpu_time = sys_cpu
status.network_bytes_sent = self.context.network_bytes_sent
status.child_session_id = self.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.STATUS,
payload=status)
# Queue the response now
self.queue_manager.QueueResponse(msg)
self.QueueNotification(session_id=request_state.session_id)
def Execute(self, action_cls, args):
responses = list()
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)
message = rdf_flows.GrrMessage(name=action_cls.__name__,
payload=args,
auth_state="AUTHENTICATED",
session_id=rdfvalue.SessionID())
action = action_cls(grr_worker=worker_mocks.FakeClientWorker())
action.SendReply = SendReply
action.Execute(message)
return responses
def testWellKnownFlows(self):
"""Test the well known flows."""
test_flow = self.FlowSetup(flow_test_lib.WellKnownSessionTest.__name__)
# Make sure the session ID is well known
self.assertEqual(
test_flow.session_id,
flow_test_lib.WellKnownSessionTest.well_known_session_id)
# Messages to Well Known flows can be unauthenticated
messages = [
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i))
for i in range(10)
]
for message in messages:
test_flow.ProcessMessage(message)
# The messages might be processed in arbitrary order
test_flow.messages.sort()
# Make sure that messages were processed even without a status
# message to complete the transaction (Well known flows do not
# have transactions or states - all messages always get to the
# ProcessMessage method):
self.assertEqual(test_flow.messages, range(10))
def testUnauthenticated(self):
"""What happens if an unauthenticated message is sent to the client?
RuntimeError needs to be issued, and the client needs to send a
GrrStatus message with the traceback in it.
"""
# Push a request on it
message = rdf_flows.GrrMessage(
name="MockAction",
session_id=self.session_id,
auth_state=rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED,
request_id=1,
generate_task_id=True)
self.context.HandleMessage(message)
# We expect to receive an GrrStatus to indicate an exception was
# raised:
# Check the response - one data and one status
message_list = self.context.Drain().job
self.assertEqual(len(message_list), 1)
self.assertEqual(message_list[0].session_id, self.session_id)
self.assertEqual(message_list[0].response_id, 1)
status = rdf_flows.GrrStatus(message_list[0].payload)
self.assertIn("not Authenticated", status.error_message)
self.assertIn("RuntimeError", status.error_message)
self.assertNotEqual(status.status,
rdf_flows.GrrStatus.ReturnedStatus.OK)
def testDelete(self):
"""Test that we can delete tasks."""
test_queue = rdfvalue.RDFURN("fooDelete")
task = rdf_flows.GrrMessage(
queue=test_queue, session_id="aff4:/Test", generate_task_id=True)
with data_store.DB.GetMutationPool() as pool:
manager = queue_manager.QueueManager(token=self.token)
manager.Schedule([task], pool)
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].session_id, "aff4:/Test")
# Now delete the task
with data_store.DB.GetMutationPool() as pool:
manager.Delete(test_queue, tasks, mutation_pool=pool)
# Task is now deleted.
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 0)
# If we try to get another lease on it we should fail - even after
# expiry time.
self._current_mock_time += 1000
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 0)
def testCountsActualNumberOfCompletedResponsesWhenApplyingTheLimit(self):
session_id = rdfvalue.SessionID(flow_name="test")
# Now queue more requests and responses:
with queue_manager.QueueManager(token=self.token) as manager:
# Start with request 1 - leave request 1 un-responded to.
for request_id in range(5):
request = rdf_flow_runner.RequestState(
id=request_id,
client_id=test_lib.TEST_CLIENT_ID,
next_state="TestState",
session_id=session_id)
manager.QueueRequest(request)
# Don't queue any actual responses, just a status message with a
# fake response_id.
manager.QueueResponse(
rdf_flows.GrrMessage(
session_id=session_id,
request_id=request_id,
response_id=1000,
type=rdf_flows.GrrMessage.Type.STATUS))
# Check that even though status message for every request indicates 1000
# responses, only the actual response count is used to apply the limit
# when FetchCompletedResponses is called.
completed_response = list(
manager.FetchCompletedResponses(session_id, limit=5))
self.assertEqual(len(completed_response), 5)
for i, (request, responses) in enumerate(completed_response):
self.assertEqual(request.id, i)
# Responses contain just the status message.
self.assertEqual(len(responses), 1)
def testSchedule(self):
"""Test the ability to schedule a task."""
test_queue = rdfvalue.RDFURN("fooSchedule")
task = rdf_flows.GrrMessage(
queue=test_queue,
task_ttl=5,
session_id="aff4:/Test",
generate_task_id=True)
manager = queue_manager.QueueManager(token=self.token)
with data_store.DB.GetMutationPool() as pool:
manager.Schedule([task], pool)
self.assertGreater(task.task_id, 0)
self.assertGreater(task.task_id & 0xffffffff, 0)
self.assertEqual((long(self._current_mock_time * 1000) & 0xffffffff) << 32,
task.task_id
& 0x1fffffff00000000)
self.assertEqual(task.task_ttl, 5)
stored_tasks = data_store.DB.QueueQueryTasks(test_queue, limit=100000)
self.assertEqual(len(stored_tasks), 1)
stored_task = stored_tasks[0]
self.assertGreater(stored_task.leased_until, 0)
stored_task.leased_until = None
self.assertRDFValuesEqual(stored_task, task)
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 4)
self.assertEqual(tasks[0].session_id, "aff4:/Test")
# If we try to get another lease on it we should fail
self._current_mock_time += 10
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 0)
# However after 100 seconds this should work again
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 3)
# Check now that after a few retransmits we drop the message
for i in range(2, 0, -1):
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, i)
# The task is now gone
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100)
self.assertEqual(len(tasks), 0)
def testTaskRetransmissionsAreCorrectlyAccounted(self):
test_queue = rdfvalue.RDFURN("fooSchedule")
task = rdf_flows.GrrMessage(
queue=test_queue,
task_ttl=5,
session_id="aff4:/Test",
generate_task_id=True)
manager = queue_manager.QueueManager(token=self.token)
with data_store.DB.GetMutationPool() as pool:
manager.Schedule([task], pool)
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 4)
self.assertEqual(
stats.STATS.GetMetricValue("grr_task_retransmission_count"),
self.retransmission_metric_value)
# Get a lease on the task 100 seconds later
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 3)
self.assertEqual(
stats.STATS.GetMetricValue("grr_task_retransmission_count"),
self.retransmission_metric_value + 1)
def testNotificationClaimsTimeout(self):
collection_urn = rdfvalue.RDFURN(
"aff4:/testNotificationClaimsTimeout/collection")
with data_store.DB.GetMutationPool() as pool:
for i in range(5):
hunts_results.HuntResultCollection.StaticAdd(
collection_urn,
rdf_flows.GrrMessage(request_id=i),
mutation_pool=pool)
results_1 = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(5, len(results_1[1]))
# Check that we have a claim - that another read returns nothing.
results_2 = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(0, len(results_2[1]))
# Push time forward past the default claim timeout, then we should be able
# to re-read (and re-claim).
with test_lib.FakeTime(rdfvalue.RDFDatetime.Now() +
rdfvalue.Duration("45m")):
results_3 = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(results_3, results_1)
def ProcessResponses(self,
plugin_args=None,
responses=None,
process_responses_separately=False):
plugin = email_plugin.EmailOutputPlugin(source_urn=self.results_urn,
args=plugin_args,
token=self.token)
plugin.InitializeState()
messages = []
for response in responses:
messages.append(
rdf_flows.GrrMessage(source=self.client_id, payload=response))
def SendEmail(address, sender, title, message, **_):
self.email_messages.append(
dict(address=address,
sender=sender,
title=title,
message=message))
with utils.Stubber(email_alerts.EMAIL_ALERTER, "SendEmail", SendEmail):
if process_responses_separately:
for message in messages:
plugin.ProcessResponses([message])
else:
plugin.ProcessResponses(messages)
plugin.Flush()
def ProcessValues(self, values_by_cls):
chunks = []
chunks.extend(list(self.plugin.Start()))
for value_cls in sorted(values_by_cls, key=lambda cls: cls.__name__):
values = values_by_cls[value_cls]
messages = []
for value in values:
messages.append(
rdf_flows.GrrMessage(source=self.client_id, payload=value))
# pylint: disable=cell-var-from-loop
chunks.extend(
list(self.plugin.ProcessValues(value_cls, lambda: messages)))
# pylint: enable=cell-var-from-loop
chunks.extend(list(self.plugin.Finish()))
fd_path = os.path.join(self.temp_dir, self.plugin.output_file_name)
with open(fd_path, "wb") as fd:
for chunk in chunks:
fd.write(chunk)
return fd_path
def CreateGenericHuntWithCollection(self, values=None):
self.client_ids = self.SetupClients(10)
CreateFileVersion(self.client_ids[0],
"fs/os/c/bin/bash",
token=self.token)
if values is None:
values = [
rdfvalue.RDFURN("aff4:/sample/1"),
rdfvalue.RDFURN("aff4:/%s/fs/os/c/bin/bash" %
self.client_ids[0].Basename()),
rdfvalue.RDFURN("aff4:/sample/3")
]
with implementation.GRRHunt.StartHunt(
hunt_name=standard.GenericHunt.__name__,
client_rule_set=self._CreateForemanClientRuleSet(),
output_plugins=[],
token=self.token) as hunt:
runner = hunt.GetRunner()
runner.Start()
collection = hunt.ResultCollection()
with data_store.DB.GetMutationPool() as pool:
for value in values:
collection.Add(rdf_flows.GrrMessage(
payload=value, source=self.client_ids[0]),
mutation_pool=pool)
return hunt.urn
def testDeleteRequest(self):
"""Check that we can efficiently destroy a single flow request."""
session_id = rdfvalue.SessionID(flow_name="test3")
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(
session_id=session_id, 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.DeleteRequest(request)
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 0)
def StaticAdd(cls,
collection_urn,
rdf_value,
timestamp=None,
suffix=None,
mutation_pool=None):
"""Adds an rdf value to a collection.
Adds an rdf value to a collection. Does not require that the collection be
open. NOTE: The caller is responsible for ensuring that the collection
exists and is of the correct type.
Args:
collection_urn: The urn of the collection to add to.
rdf_value: The rdf value to add to the collection. If this value is not
GrrMessage, it will be wrapped into GrrMessage (later when
collection is iterated, this value will still be returned wrapped
in GrrMessage).
timestamp: The timestamp (in microseconds) to store the rdf value
at. Defaults to the current time.
suffix: A 'fractional timestamp' suffix to reduce the chance of
collisions. Defaults to a random number.
mutation_pool: A MutationPool object to write to.
Returns:
The pair (timestamp, suffix) which identifies the value within the
collection.
Raises:
ValueError: rdf_value has unexpected type.
"""
if rdf_value is None:
raise ValueError("Can't add None to MultiTypeCollection")
if mutation_pool is None:
raise ValueError("Mutation pool can't be none.")
if not isinstance(rdf_value, rdf_flows.GrrMessage):
rdf_value = rdf_flows.GrrMessage(payload=rdf_value)
value_type = rdf_value.args_rdf_name or rdf_flows.GrrMessage.__name__
# In order to make this fast, we never actually generate the
# subcollections, we just use them. This means that we cannot use
# ListChildren to get all the items stored in this
# MultiTypeCollection.
subpath = collection_urn.Add(value_type)
sequential_collection.GrrMessageCollection.StaticAdd(
subpath,
rdf_value,
timestamp=timestamp,
suffix=suffix,
mutation_pool=mutation_pool)
mutation_pool.CollectionAddStoredTypeIndex(collection_urn, value_type)
def testExtractsTypesFromGrrMessage(self):
with self.pool:
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(0)),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString("foo")),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFURN("aff4:/foo/bar")),
mutation_pool=self.pool)
self.assertEqual(
set([
rdfvalue.RDFInteger.__name__, rdfvalue.RDFString.__name__,
rdfvalue.RDFURN.__name__
]), set(self.collection.ListStoredTypes()))
def testValuesOfMultipleTypesCanBeIteratedPerType(self):
with self.pool:
for i in range(100):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)),
mutation_pool=self.pool)
for index, (_, v) in enumerate(
self.collection.ScanByType(rdfvalue.RDFInteger.__name__)):
self.assertEqual(index, v.payload)
for index, (_, v) in enumerate(
self.collection.ScanByType(rdfvalue.RDFString.__name__)):
self.assertEqual(str(index), v.payload)
def testValuesOfMultipleTypesCanBeIteratedTogether(self):
with self.pool:
original_values = set()
for i in range(100):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)),
mutation_pool=self.pool)
original_values.add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)))
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)),
mutation_pool=self.pool)
original_values.add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)))
self.assertEqual(sorted([v.payload for v in original_values]),
sorted([v.payload for v in self.collection]))
def testValuesOfSingleTypeAreAddedAndIterated(self):
with self.pool:
for i in range(100):
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)),
mutation_pool=self.pool)
for index, v in enumerate(self.collection):
self.assertEqual(index, v.payload)
def SendReply(self,
rdf_value,
message_type=rdf_flows.GrrMessage.Type.MESSAGE,
**kw):
message = rdf_flows.GrrMessage(type=message_type,
payload=rdf_value,
**kw)
self.responses.append(message)
def testDeletingCollectionDeletesAllSubcollections(self):
if not isinstance(data_store.DB, fake_data_store.FakeDataStore):
self.skipTest("Only supported on FakeDataStore.")
with self.pool:
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(0)),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFString("foo")),
mutation_pool=self.pool)
self.collection.Add(
rdf_flows.GrrMessage(payload=rdfvalue.RDFURN("aff4:/foo/bar")),
mutation_pool=self.pool)
self.collection.Delete()
for urn in data_store.DB.subjects.keys():
self.assertFalse(
utils.SmartStr(self.collection.collection_id) in urn)
def ExecuteAction(self, action_cls, arg=None, grr_worker=None):
message = rdf_flows.GrrMessage(
name=action_cls.__name__, payload=arg, auth_state="AUTHENTICATED")
self.results = []
action = self._GetActionInstance(action_cls, grr_worker=grr_worker)
action.Execute(message)
return self.results
def testLinuxTransactionLog(self):
"""Tests the linux transaction log."""
with tempfile.NamedTemporaryFile() as fd:
log = client_utils_linux.TransactionLog(logfile=fd.name)
grr_message = rdf_flows.GrrMessage(session_id="W:test")
log.Write(grr_message)
self.assertRDFValuesEqual(grr_message, log.Get())
log.Clear()
self.assertIsNone(log.Get())
