def NotifyAboutEnd(self):
super(GetFile, self).NotifyAboutEnd()
if not self.runner.ShouldSendNotifications():
return
stat_entry = self.state.stat_entry
if not stat_entry:
stat_entry = rdf_client.StatEntry(pathspec=self.args.pathspec)
urn = stat_entry.AFF4Path(self.client_id)
components = urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
if not self.state.get("success"):
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_FILE_COLLECTION_FAILED, "File transfer failed.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
else:
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File transferred successfully.",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def setUp(self):
super(ApiDeletePendingUserNotificationHandlerTest, self).setUp()
self.handler = user_plugin.ApiDeletePendingUserNotificationHandler()
self.client_id = self.SetupClient(0)
with test_lib.FakeTime(self.TIME_0):
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"<some message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"<some message with identical time>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
with test_lib.FakeTime(self.TIME_1):
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_APPROVAL_GRANTED,
"<some other message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
def GenerateNotifications(cls, client_id, token):
"""Generates fake notifications of different notification types."""
session_id = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=discovery.Interrogate.__name__,
token=token)
with aff4.FACTORY.Open(session_id, mode="rw", token=token) as flow_obj:
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"Fake discovery message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=client_id.Basename())))
# ViewObject: VirtualFileSystem
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
"File fetch completed",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["proc", "10", "exe"])))
gui_test_lib.CreateFileVersion(
client_id,
"fs/os/proc/10/exe",
"",
timestamp=gui_test_lib.TIME_0,
token=token)
# ViewObject: Flow
notification.Notify(
token.username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_COMPLETED,
"Fake view flow message",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=rdf_objects.FlowReference(
client_id=client_id.Basename(),
flow_id=flow_obj.urn.Basename())))
# FlowError
flow_obj.GetRunner().Error("Fake flow error")
return session_id
def NotifyAboutEnd(self):
if not self.runner.ShouldSendNotifications():
return
status_text = "Recursive Directory Listing complete %d nodes, %d dirs"
urn = self.state.first_directory
if not urn:
try:
urn = self.args.pathspec.AFF4Path(self.client_id)
except ValueError:
pass
if urn:
components = urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
notification.Notify(
self.token.username, rdf_objects.UserNotification.Type.
TYPE_VFS_RECURSIVE_LIST_DIRECTORY_COMPLETED,
status_text % (self.state.file_count, self.state.dir_count),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
def NotifyAboutEnd(self):
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED, "",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
def testHuntNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_HUNT_STOPPED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(hunt_id="H:123456")))
self.assertEqual(n.reference.type, "HUNT")
self.assertEqual(n.reference.hunt.hunt_urn, "aff4:/hunts/H:123456")
def testCronNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_CRON_JOB_APPROVAL_GRANTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CRON_JOB,
cron_job=rdf_objects.CronJobReference(cron_job_id="FooBar")))
self.assertEqual(n.reference.type, "CRON")
self.assertEqual(n.reference.cron.cron_job_urn, "aff4:/cron/FooBar")
def testClientApprovalGrantedNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_CLIENT_APPROVAL_GRANTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
self.assertEqual(n.reference.type, "DISCOVERY")
self.assertEqual(n.reference.discovery.client_id, self.client_id)
def testDiscoveryNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=self.client_id.Basename())))
self.assertEqual(n.reference.type, "CLIENT")
self.assertEqual(n.reference.client.client_id, self.client_id)
def testFlowSuccessNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_COMPLETED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=rdf_objects.FlowReference(
client_id=self.client_id.Basename(), flow_id="F:123456")))
self.assertEqual(n.reference.type, "FLOW")
self.assertEqual(n.reference.flow.client_id, self.client_id)
self.assertEqual(n.reference.flow.flow_id, "F:123456")
def _SendNotifications(username, client_id):
with test_lib.FakeTime(42):
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED,
"<some message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CLIENT,
client=rdf_objects.ClientReference(
client_id=client_id.Basename())))
with test_lib.FakeTime(44):
notification.Notify(
username,
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTION_FAILED,
"<some other message>",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["foo"])))
def testVfsNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_VFS_FILE_COLLECTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=rdf_objects.VfsFileReference(
client_id=self.client_id.Basename(),
path_type=rdf_objects.PathInfo.PathType.OS,
path_components=["foo", "bar"])))
self.assertEqual(n.reference.type, "VFS")
self.assertEqual(n.reference.vfs.client_id, self.client_id)
self.assertEqual(n.reference.vfs.vfs_path, "fs/os/foo/bar")
def testHuntApprovalNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_HUNT_APPROVAL_REQUESTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.APPROVAL_REQUEST,
approval_request=rdf_objects.ApprovalRequestReference(
approval_type=rdf_objects.ApprovalRequest.ApprovalType.
APPROVAL_TYPE_HUNT,
approval_id="foo-bar",
subject_id="H:123456",
requestor_username=self.token.username)))
self.assertEqual(n.reference.type, "HUNT_APPROVAL")
self.assertEqual(n.reference.hunt_approval.hunt_id, "H:123456")
self.assertEqual(n.reference.hunt_approval.username, self.token.username)
self.assertEqual(n.reference.hunt_approval.approval_id, "foo-bar")
def testCronJobNotificationIsShownAndClickable(self):
notification.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_CRON_JOB_APPROVAL_GRANTED,
"Test CronJob notification",
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CRON_JOB,
cron_job=rdf_objects.CronJobReference(cron_job_id="OSBreakDown")))
self.Open("/")
self.Click("css=button[id=notification_button]")
self.Click("css=a:contains('Test CronJob notification')")
self.WaitUntil(self.IsElementPresent,
"css=tr.row-selected td:contains('OSBreakDown')")
self.WaitUntil(self.IsElementPresent, "css=dd:contains('OSBreakDown')")
def testCronJobApprovalNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_CRON_JOB_APPROVAL_REQUESTED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.APPROVAL_REQUEST,
approval_request=rdf_objects.ApprovalRequestReference(
approval_type=rdf_objects.ApprovalRequest.ApprovalType.
APPROVAL_TYPE_CRON_JOB,
approval_id="foo-bar",
subject_id="FooBar",
requestor_username=self.token.username)))
self.assertEqual(n.reference.type, "CRON_JOB_APPROVAL")
self.assertEqual(n.reference.cron_job_approval.cron_job_id, "FooBar")
self.assertEqual(n.reference.cron_job_approval.username,
self.token.username)
self.assertEqual(n.reference.cron_job_approval.approval_id, "foo-bar")
def Error(self, backtrace, client_id=None, status_code=None):
"""Terminates this flow with an error."""
try:
self.queue_manager.DestroyFlowStates(self.session_id)
except queue_manager.MoreDataException:
pass
if not self.IsRunning():
return
# Set an error status
reply = rdf_flows.GrrStatus()
if status_code is None:
reply.status = rdf_flows.GrrStatus.ReturnedStatus.GENERIC_ERROR
else:
reply.status = status_code
client_id = client_id or self.runner_args.client_id
if backtrace:
reply.error_message = backtrace
logging.error("Error in flow %s (%s). Trace: %s", self.session_id,
client_id, backtrace)
self.context.backtrace = backtrace
else:
logging.error("Error in flow %s (%s).", self.session_id, client_id)
self._SendTerminationMessage(reply)
self.context.state = rdf_flows.FlowContext.State.ERROR
if self.ShouldSendNotifications():
flow_ref = None
if client_id:
flow_ref = rdf_objects.FlowReference(
client_id=client_id.Basename(), flow_id=self.session_id.Basename())
notification_lib.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_FAILED,
"Flow (%s) terminated due to error" % self.session_id,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=flow_ref))
self.flow_obj.Flush()
def NotifyAboutEnd(self):
if not self.runner.ShouldSendNotifications():
return
if self.state.urn:
components = self.state.urn.Split()
file_ref = None
if len(components) > 3:
file_ref = rdf_objects.VfsFileReference(
client_id=components[0],
path_type=components[2].upper(),
path_components=components[3:])
notification.Notify(
self.token.username,
notification.UserNotification.Type.TYPE_VFS_LIST_DIRECTORY_COMPLETED,
"List of {0} completed.".format(utils.SmartStr(self.args.pathspec)),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.VFS_FILE,
vfs_file=file_ref))
else:
super(IteratedListDirectory, self).NotifyAboutEnd()
def NotifyAboutEnd(self):
"""Send out a final notification about the end of this flow."""
if not self.runner.ShouldSendNotifications():
return
flow_ref = None
if self.runner_args.client_id:
flow_ref = rdf_objects.FlowReference(
client_id=self.runner_args.client_id.Basename(),
flow_id=self.urn.Basename())
num_results = len(self.ResultCollection())
notification_lib.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_COMPLETED,
"Flow %s completed with %d %s" %
(self.__class__.__name__, num_results,
num_results == 1 and "result" or "results"),
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=flow_ref))
def ObjectReference(self):
return objects.ObjectReference(
reference_type=objects.ObjectReference.Type.CLIENT,
client=objects.ClientReference(
client_id=utils.SmartStr(self.client_id)))
def ObjectReference(self):
return rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(
hunt_id=utils.SmartStr(self.hunt_id)))
def ObjectReference(self):
return rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.CRON_JOB,
cron_job=rdf_objects.CronJobReference(
cron_job_id=utils.SmartStr(self.cron_job_id)))