def testFlowNotification(self):
FlowDoneListener.received_events = []
# Run the flow in the simulated way
client_mock = action_mocks.ActionMock()
for _ in test_lib.TestFlowHelper("DummyLogFlow",
client_mock,
client_id=self.client_id,
notification_urn=rdfvalue.SessionID(
queue=rdfvalue.RDFURN("EV"),
flow_name="FlowDone"),
token=self.token):
pass
# The event goes to an external queue so we need another worker.
worker = test_lib.MockWorker(queues=[rdfvalue.RDFURN("EV")],
token=self.token)
worker.Simulate()
self.assertEqual(len(FlowDoneListener.received_events), 1)
flow_event = FlowDoneListener.received_events[0].payload
self.assertEqual(flow_event.flow_name, "DummyLogFlow")
self.assertEqual(flow_event.client_id, "aff4:/C.1000000000000000")
self.assertEqual(flow_event.status,
rdf_flows.FlowNotification.Status.OK)
def testAuditEntryIsCreatedForEveryClient(self):
self.handler.Handle(client_plugin.ApiAddClientsLabelsArgs(
client_ids=self.client_ids, labels=["drei", "ein", "zwei"]),
token=self.token)
# We need to run .Simulate() so that the appropriate event is fired,
# collected, and finally written to the logs that we inspect.
mock_worker = test_lib.MockWorker(token=self.token)
mock_worker.Simulate()
parentdir = aff4.FACTORY.Open("aff4:/audit/logs", token=self.token)
log = list(parentdir.ListChildren())[0]
fd = aff4.FACTORY.Open(log, token=self.token)
for client_id in self.client_ids:
found_event = None
for event in fd:
if (event.action == events.AuditEvent.Action.CLIENT_ADD_LABEL
and event.client == rdf_client.ClientURN(client_id)):
found_event = event
break
self.assertFalse(found_event is None)
self.assertEqual(found_event.user, self.token.username)
self.assertEqual(
found_event.description, "%s.drei,%s.ein,%s.zwei" %
(self.token.username, self.token.username,
self.token.username))
def testUserModificationAudit(self):
audit.AuditEventListener.created_logs.clear()
worker = test_lib.MockWorker(token=self.token)
token = self.GenerateToken(username="******", reason="reason")
maintenance_utils.AddUser(
"testuser", password="******", labels=["admin"], token=token)
worker.Simulate()
maintenance_utils.UpdateUser(
"testuser", "xxx", delete_labels=["admin"], token=token)
worker.Simulate()
maintenance_utils.DeleteUser("testuser", token=token)
worker.Simulate()
log_entries = []
for log in aff4.FACTORY.Open(
"aff4:/audit/logs", token=self.token).OpenChildren():
log_entries.extend(log)
self.assertEqual(len(log_entries), 3)
self.assertEqual(log_entries[0].action, "USER_ADD")
self.assertEqual(log_entries[0].urn, "aff4:/users/testuser")
self.assertEqual(log_entries[0].user, "usermodtest")
self.assertEqual(log_entries[1].action, "USER_UPDATE")
self.assertEqual(log_entries[1].urn, "aff4:/users/testuser")
self.assertEqual(log_entries[1].user, "usermodtest")
self.assertEqual(log_entries[2].action, "USER_DELETE")
self.assertEqual(log_entries[2].urn, "aff4:/users/testuser")
self.assertEqual(log_entries[2].user, "usermodtest")
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_grr.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 = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED
events.Events.PublishEvent("FileStore.AddFileToStore",
rdf_flows.GrrMessage(payload=urn,
auth_state=auth_state),
token=token)
worker = test_lib.MockWorker(token=token)
worker.Simulate()
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 = pathspec.AFF4Path(client_id)
client_mock = action_mocks.GetFileClientMock()
for _ in test_lib.TestFlowHelper(
transfer.GetFile.__name__,
client_mock,
token=token,
client_id=client_id,
pathspec=pathspec):
pass
auth_state = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED
events.Events.PublishEvent(
"FileStore.AddFileToStore",
rdf_flows.GrrMessage(payload=urn, auth_state=auth_state),
token=token)
worker = test_lib.MockWorker(token=token)
worker.Simulate()
return urn
def _RunRateLimitedHunt(self, client_ids, start_time):
with hunts.GRRHunt.StartHunt(
hunt_name="DummyHunt",
regex_rules=[
rdf_foreman.ForemanAttributeRegex(attribute_name="GRR client",
attribute_regex="GRR"),
],
client_rate=1, token=self.token) as hunt:
hunt.Run()
# Pretend to be the foreman now and dish out hunting jobs to all the
# clients..
foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token)
for client_id in client_ids:
foreman.AssignTasksToClient(client_id)
self.assertEqual(len(DummyHunt.client_ids), 0)
# Run the hunt.
worker_mock = test_lib.MockWorker(check_flow_errors=True,
queues=queues.HUNTS,
token=self.token)
# One client is scheduled in the first minute.
with test_lib.FakeTime(start_time + 2):
worker_mock.Simulate()
self.assertEqual(len(DummyHunt.client_ids), 1)
# No further clients will be scheduled until the end of the first minute.
with test_lib.FakeTime(start_time + 59):
worker_mock.Simulate()
self.assertEqual(len(DummyHunt.client_ids), 1)
return worker_mock, hunt.urn
def testDelayedCallState(self):
"""Tests the ability to delay a CallState invocation."""
with test_lib.FakeTime(10000):
client_mock = ClientMock()
client_mock = test_lib.MockClient(self.client_id, client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
flow.GRRFlow.StartFlow(
client_id=self.client_id, flow_name="DelayedCallStateFlow",
token=self.token)
self.Work(client_mock, worker_mock)
# We should have done the first CallState so far.
self.assertEqual(DelayedCallStateFlow.flow_ran, 1)
with test_lib.FakeTime(10050):
# 50 seconds more is not enough.
self.Work(client_mock, worker_mock)
self.assertEqual(DelayedCallStateFlow.flow_ran, 1)
with test_lib.FakeTime(10100):
# But 100 is.
self.Work(client_mock, worker_mock)
self.assertEqual(DelayedCallStateFlow.flow_ran, 2)
def PopulateData(token=None):
"""Populates some usage information into the data store."""
flows = ["ListDirectory", "GetFile", "Flow1", "KillFlow", "Interrogate"]
def SimulateUserActivity(username, client_id, timestamp, num=10):
for flow_name in flows[:num]:
event = rdfvalue.AuditEvent(user=username,
action="RUN_FLOW",
flow_name=flow_name,
client=client_id,
age=timestamp)
flow.Events.PublishEvent("Audit", event, token=token)
now = int(rdfvalue.RDFDatetime().Now())
week_duration = 7 * 24 * 60 * 60 * 1e6
SimulateUserActivity("test", "C.0000000000000001", now)
SimulateUserActivity("test", "C.0000000000000002", now)
SimulateUserActivity("admin", "C.0000000000000001", now)
SimulateUserActivity("user", "C.0000000000000003", now)
SimulateUserActivity("test", "C.0000000000000001", now - week_duration)
SimulateUserActivity("admin", "C.0000000000000001", now - week_duration)
SimulateUserActivity("user", "C.0000000000000001", now - 2 * week_duration)
SimulateUserActivity("admin", "C.0000000000000001",
now - 2 * week_duration)
test_lib.MockWorker(token=token).Simulate()
def testAuditEntryIsCreatedForEveryClient(self):
client_ids = self.SetupClients(3)
flow.GRRFlow.StartFlow(flow_name="ApplyLabelsToClientsFlow",
clients=client_ids,
labels=["drei", "ein", "zwei"],
token=self.token)
mock_worker = test_lib.MockWorker(token=self.token)
mock_worker.Simulate()
fd = aff4.FACTORY.Open("aff4:/audit/log", token=self.token)
for client_id in client_ids:
found_event = None
for event in fd:
if (event.action == rdfvalue.AuditEvent.Action.CLIENT_ADD_LABEL
and event.client == rdfvalue.ClientURN(client_id)):
found_event = event
break
self.assertFalse(found_event is None)
self.assertEqual(found_event.flow_name, "ApplyLabelsToClientsFlow")
self.assertEqual(found_event.user, self.token.username)
self.assertEqual(found_event.description,
"test.drei,test.ein,test.zwei")
def testAttributesOfFileFoundInHashFileStoreAreSetCorrectly(self):
client_ids = self.SetupClients(2)
filename = os.path.join(self.base_path, "tcpip.sig")
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=filename)
urn1 = aff4_grr.VFSGRRClient.PathspecToURN(pathspec, client_ids[0])
urn2 = aff4_grr.VFSGRRClient.PathspecToURN(pathspec, client_ids[1])
for client_id in client_ids:
client_mock = action_mocks.FileFinderClientMock()
for _ in test_lib.TestFlowHelper(
file_finder.FileFinder.__name__,
client_mock,
token=self.token,
client_id=client_id,
paths=[filename],
action=rdf_file_finder.FileFinderAction(
action_type=rdf_file_finder.FileFinderAction.Action.DOWNLOAD)):
pass
# Running worker to make sure FileStore.AddFileToStore event is processed
# by the worker.
worker = test_lib.MockWorker(token=self.token)
worker.Simulate()
fd1 = aff4.FACTORY.Open(urn1, token=self.token)
self.assertTrue(isinstance(fd1, aff4_grr.VFSBlobImage))
fd2 = aff4.FACTORY.Open(urn2, token=self.token)
self.assertTrue(isinstance(fd2, filestore.FileStoreImage))
self.assertEqual(fd1.Get(fd1.Schema.STAT), fd2.Get(fd2.Schema.STAT))
self.assertEqual(fd1.Get(fd1.Schema.SIZE), fd2.Get(fd2.Schema.SIZE))
self.assertEqual(
fd1.Get(fd1.Schema.CONTENT_LAST), fd2.Get(fd2.Schema.CONTENT_LAST))
def testFlowNotification(self):
FlowDoneListener.received_events = []
# Install the mock
vfs.VFS_HANDLERS[rdfvalue.PathSpec.PathType.OS] = MockVFSHandler
path = rdfvalue.PathSpec(path="/",
pathtype=rdfvalue.PathSpec.PathType.OS)
# Run the flow in the simulated way
client_mock = action_mocks.ActionMock("IteratedListDirectory")
for _ in test_lib.TestFlowHelper(
"IteratedListDirectory", client_mock, client_id=self.client_id,
notification_urn=rdfvalue.SessionID(queue=rdfvalue.RDFURN("EV"),
flow_name="FlowDone"),
pathspec=path, token=self.token):
pass
# The event goes to an external queue so we need another worker.
worker = test_lib.MockWorker(queues=[rdfvalue.RDFURN("EV")],
token=self.token)
worker.Simulate()
self.assertEqual(len(FlowDoneListener.received_events), 1)
flow_event = FlowDoneListener.received_events[0].payload
self.assertEqual(flow_event.flow_name, "IteratedListDirectory")
self.assertEqual(flow_event.client_id, "aff4:/C.1000000000000000")
self.assertEqual(flow_event.status, rdfvalue.FlowNotification.Status.OK)
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=collects.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=collects.PackedVersionedCollection,
token=self.token))
self.assertEqual(len(global_crashes), 1)
self.assertEqual(global_crashes[0], crash)
def testExportWithDummyPlugin(self):
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(self.base_path, "winexec_img.dd"))
pathspec.Append(path="/Ext2IFS_1_10b.exe",
pathtype=rdf_paths.PathSpec.PathType.TSK)
urn = aff4.AFF4Object.VFSGRRClient.PathspecToURN(pathspec, self.client_id)
client_mock = action_mocks.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 = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED
flow.Events.PublishEvent(
"FileStore.AddFileToStore",
rdf_flows.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, aff4_filestore.FileStoreHash))
self.assertTrue(aff4_filestore.FileStoreHash(
fingerprint_type="pecoff",
hash_type="md5",
hash_value="a3a3259f7b145a21c7b512d876a5da06") in responses)
self.assertTrue(aff4_filestore.FileStoreHash(
fingerprint_type="pecoff",
hash_type="sha1",
hash_value="019bddad9cac09f37f3941a7f285c79d3c7e7801") in responses)
self.assertTrue(aff4_filestore.FileStoreHash(
fingerprint_type="generic",
hash_type="md5",
hash_value="bb0a15eefe63fd41f8dc9dee01c5cf9a") in responses)
self.assertTrue(aff4_filestore.FileStoreHash(
fingerprint_type="generic",
hash_type="sha1",
hash_value="7dd6bee591dfcb6d75eb705405302c3eab65e21a") in responses)
self.assertTrue(aff4_filestore.FileStoreHash(
fingerprint_type="generic",
hash_type="sha256",
hash_value="0e8dc93e150021bb4752029ebbff51394aa36f06"
"9cf19901578e4f06017acdb5") in responses)
def testWorkerPrioritization(self):
"""Test that flow priorities work on the worker side."""
result = []
client_mock = PriorityClientMock(result)
client_mock = test_lib.MockClient(self.client_id,
client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
# Start some flows with different priorities.
# pyformat: disable
args = [
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"),
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"),
(rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2")
]
# pyformat: enable
server_result = []
PriorityFlow.storage = server_result
for (priority, msg) in args:
flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="PriorityFlow",
msg=msg,
priority=priority,
token=self.token)
while True:
# Run all the clients first so workers have messages to choose from.
client_processed = 1
while client_processed:
client_processed = client_mock.Next()
# Now process the results, this should happen in the correct order.
flows_run = []
for flow_run in worker_mock.Next():
flows_run.append(flow_run)
if not flows_run:
break
# The flows should be run in order of priority.
self.assertEqual(server_result[0:1], [u"high priority"])
self.assertEqual(sorted(server_result[1:3]),
[u"medium priority", u"medium priority2"])
self.assertEqual(sorted(server_result[3:5]),
[u"low priority", u"low priority2"])
def testHuntClientRate(self):
"""Check that clients are scheduled slowly by the hunt."""
start_time = 10
# Set up 10 clients.
client_ids = self.SetupClients(10)
with test_lib.FakeTime(start_time):
with hunts.GRRHunt.StartHunt(hunt_name="DummyHunt",
regex_rules=[
rdfvalue.ForemanAttributeRegex(
attribute_name="GRR client",
attribute_regex="GRR"),
],
client_rate=1,
token=self.token) as hunt:
hunt.Run()
# Pretend to be the foreman now and dish out hunting jobs to all the
# clients..
foreman = aff4.FACTORY.Open("aff4:/foreman",
mode="rw",
token=self.token)
for client_id in client_ids:
foreman.AssignTasksToClient(client_id)
self.assertEqual(len(DummyHunt.client_ids), 0)
# Run the hunt.
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
time.time = lambda: start_time + 2
# One client is scheduled in the first minute.
worker_mock.Simulate()
self.assertEqual(len(DummyHunt.client_ids), 1)
# No further clients will be scheduled until the end of the first minute.
time.time = lambda: start_time + 59
worker_mock.Simulate()
self.assertEqual(len(DummyHunt.client_ids), 1)
# One client will be processed every minute.
for i in range(len(client_ids)):
time.time = lambda: start_time + 1 + 60 * i
worker_mock.Simulate()
self.assertEqual(len(DummyHunt.client_ids), i + 1)
def testAuditEntryIsCreatedForEveryClient(self):
self.handler.Handle(client_plugin.ApiAddClientsLabelsArgs(
client_ids=self.client_ids, labels=["drei", "ein", "zwei"]),
token=self.token)
# We need to run .Simulate() so that the appropriate event is fired,
# collected, and finally written to the logs that we inspect.
mock_worker = test_lib.MockWorker(token=self.token)
mock_worker.Simulate()
event = self._FindAuditEvent()
self.assertFalse(event is None)
self.assertEqual(event.user, self.token.username)
self.assertEqual(
event.description, "%s.drei,%s.ein,%s.zwei" %
(self.token.username, self.token.username, self.token.username))
def testStatEntryToExportedFileConverterWithHashedAFF4File(self):
client_ids = self.SetupClients(1)
client_id = client_ids[0]
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, client_id)
client_mock = test_lib.ActionMock("TransferBuffer", "StatFile",
"HashBuffer")
for _ in test_lib.TestFlowHelper(
"GetFile", client_mock, token=self.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=self.token)
worker = test_lib.MockWorker(token=self.token)
worker.Simulate()
fd = aff4.FACTORY.Open(urn, token=self.token)
hash_value = fd.Get(fd.Schema.HASH)
self.assertTrue(hash_value)
converter = export.StatEntryToExportedFileConverter(
options=rdfvalue.ExportOptions(export_files_hashes=True))
results = list(converter.Convert(rdfvalue.ExportedMetadata(),
rdfvalue.StatEntry(aff4path=urn,
pathspec=pathspec),
token=self.token))
self.assertEqual(results[0].hash_md5,
"bb0a15eefe63fd41f8dc9dee01c5cf9a")
self.assertEqual(results[0].hash_sha1,
"7dd6bee591dfcb6d75eb705405302c3eab65e21a")
self.assertEqual(
results[0].hash_sha256,
"0e8dc93e150021bb4752029ebbff51394aa36f069cf19901578e4f06017acdb5")
def __init__(self, method_name=None):
super(GRRFuseTest, self).__init__(method_name)
# Set up just once for the whole test suite, since we don't have any
# per-test setup to do.
super(GRRFuseTest, self).setUp()
self.client_name = str(self.client_id)[len("aff4:/"):]
with aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw") as fd:
fd.Set(fd.Schema.SYSTEM("Linux"))
kb = fd.Schema.KNOWLEDGE_BASE()
fd.Set(kb)
# Ignore cache so our tests always get client side updates.
self.grr_fuse = fuse_mount.GRRFuse(root="/", token=self.token,
ignore_cache=True)
self.action_mock = test_lib.ActionMock("TransferBuffer", "StatFile", "Find",
"HashFile", "HashBuffer",
"UpdateVFSFile",
"EnumerateInterfaces",
"EnumerateFilesystems",
"GetConfiguration",
"GetConfig", "GetClientInfo",
"GetInstallDate", "GetPlatformInfo",
"EnumerateUsers", "ListDirectory")
client_mock = test_lib.MockClient(self.client_id, self.action_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token)
# All the flows we've run so far. We'll check them for errors at the end of
# each test.
self.total_flows = set()
# We add the thread as a class variable since we'll be referring to it in
# the tearDownClass method, and we want all tests to share it.
self.__class__.fake_server_thread = threading.Thread(
target=self.RunFakeWorkerAndClient,
args=(client_mock,
worker_mock))
self.fake_server_thread.start()
def testClientPrioritization(self):
"""Test that flow priorities work on the client side."""
result = []
client_mock = PriorityClientMock(result)
client_mock = test_lib.MockClient(self.client_id,
client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
# Start some flows with different priorities.
# pyformat: disable
args = [
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"),
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"),
(rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2")
]
# pyformat: enable
for (priority, msg) in args:
flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="PriorityFlow",
msg=msg,
priority=priority,
token=self.token)
while True:
client_processed = client_mock.Next()
flows_run = []
for flow_run in worker_mock.Next():
flows_run.append(flow_run)
if client_processed == 0 and not flows_run:
break
# The flows should be run in order of priority.
self.assertEqual(result[0:1], [u"high priority"])
self.assertEqual(sorted(result[1:3]),
[u"medium priority", u"medium priority2"])
self.assertEqual(sorted(result[3:5]),
[u"low priority", u"low priority2"])
def testClientEventNotification(self):
"""Make sure that client events handled securely."""
ClientListener.received_events = []
NoClientListener.received_events = []
event = rdfvalue.GrrMessage(
source="C.1395c448a443c7d9",
auth_state=rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED)
event.payload = rdfvalue.PathSpec(path="foobar")
flow.Events.PublishEvent("TestEvent", event, token=self.token)
test_lib.MockWorker(token=self.token).Simulate()
# The same event should be sent to both listeners, but only the listener
# which accepts client messages should register it.
self.assertProtoEqual(ClientListener.received_events[0][0].payload,
event.payload)
self.assertEqual(NoClientListener.received_events, [])
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 testRegexChangeNotification(self):
"""Test the AFF4RegexNotificationRule rule."""
client_name = "C." + "0" * 16
# Create the notification rule.
rule_fd = aff4.FACTORY.Create("aff4:/config/aff4_rules/new_rule",
aff4_type="AFF4RegexNotificationRule",
token=self.token)
rule_fd.Set(rule_fd.Schema.CLIENT_PATH_REGEX("b.*"))
rule_fd.Set(rule_fd.Schema.EVENT_NAME("MockChangeEvent"))
rule_fd.Set(rule_fd.Schema.NOTIFY_ONLY_IF_NEW(0))
rule_fd.Close()
# Force notification rules to be reloaded.
aff4.FACTORY.UpdateNotificationRules()
fd = aff4.FACTORY.Create(rdfvalue.ClientURN(client_name).Add("a"),
token=self.token,
aff4_type="AFF4Object")
fd.Close()
worker_mock = test_lib.MockWorker(token=self.token)
while worker_mock.Next():
pass
# No notifications are expected, because path doesn't match the regex
self.assertEqual(len(MockChangeEvent.CHANGED_URNS), 0)
fd = aff4.FACTORY.Create(rdfvalue.ClientURN(client_name).Add("b"),
token=self.token,
aff4_type="AFF4Object")
fd.Close()
while worker_mock.Next():
pass
# Now we get a notification, because the path matches
self.assertEqual(len(MockChangeEvent.CHANGED_URNS), 1)
self.assertEqual(MockChangeEvent.CHANGED_URNS[0],
rdfvalue.ClientURN(client_name).Add("b"))
MockChangeEvent.CHANGED_URNS = []
# Write again to the same file and check that there's notification again
fd = aff4.FACTORY.Create(rdfvalue.ClientURN(client_name).Add("b"),
token=self.token,
aff4_type="AFF4Object")
fd.Close()
while worker_mock.Next():
pass
self.assertEqual(len(MockChangeEvent.CHANGED_URNS), 1)
self.assertEqual(MockChangeEvent.CHANGED_URNS[0],
rdfvalue.ClientURN(client_name).Add("b"))
MockChangeEvent.CHANGED_URNS = []
# Change the rule to notify only if file is written for the first time
rule_fd = aff4.FACTORY.Open("aff4:/config/aff4_rules/new_rule",
mode="rw",
token=self.token)
rule_fd.Set(rule_fd.Schema.NOTIFY_ONLY_IF_NEW, rdfvalue.RDFInteger(1))
rule_fd.Close()
# Force update of the rules in the factory
aff4.FACTORY.UpdateNotificationRules()
# Check that we don't get a notification for overwriting existing file
fd = aff4.FACTORY.Create(rdfvalue.ClientURN(client_name).Add("b"),
token=self.token,
aff4_type="AFF4Object")
fd.Close()
while worker_mock.Next():
pass
self.assertEqual(len(MockChangeEvent.CHANGED_URNS), 0)
# Check that we do get a notification for writing a new file
fd = aff4.FACTORY.Create(rdfvalue.ClientURN(client_name).Add("b2"),
token=self.token,
aff4_type="AFF4Object")
fd.Close()
while worker_mock.Next():
pass
self.assertEqual(len(MockChangeEvent.CHANGED_URNS), 1)
self.assertEqual(MockChangeEvent.CHANGED_URNS[0],
rdfvalue.ClientURN(client_name).Add("b2"))