From 49415c7ef1c40e7940999302ea80ad0b640474cc Mon Sep 17 00:00:00 2001
From: Adam <adamm@iguazio.com>
Date: Wed, 22 May 2024 11:27:38 +0300
Subject: [PATCH] [Notifications] Redact Secret Params from Notification Data
 in Webhook Notification (#5610)

redact secret params

Co-authored-by: quaark <a.melnick@icloud.com>
---
 server/api/main.py                      |  3 ++-
 server/api/utils/notification_pusher.py | 26 ++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/server/api/main.py b/server/api/main.py
index 4953bab980e..2b258e2aac8 100644
--- a/server/api/main.py
+++ b/server/api/main.py
@@ -41,6 +41,7 @@
 import server.api.runtime_handlers
 import server.api.utils.clients.chief
 import server.api.utils.clients.log_collector
+import server.api.utils.notification_pusher
 from mlrun.config import config
 from mlrun.errors import err_to_str
 from mlrun.runtimes import RuntimeClassMode, RuntimeKinds
@@ -757,7 +758,7 @@ def _push_terminal_run_notifications(
     logger.debug(
         "Got terminal runs with configured notifications", runs_amount=len(runs)
     )
-    mlrun.utils.notifications.NotificationPusher(unmasked_runs).push()
+    server.api.utils.notification_pusher.RunNotificationPusher(unmasked_runs).push()
 
 
 def _generate_event_on_failed_runs(
diff --git a/server/api/utils/notification_pusher.py b/server/api/utils/notification_pusher.py
index 52ae0f4f99a..5db8ce0bd8c 100644
--- a/server/api/utils/notification_pusher.py
+++ b/server/api/utils/notification_pusher.py
@@ -21,9 +21,33 @@
 import mlrun.model
 import mlrun.utils.helpers
 import server.api.api.utils
+import server.api.constants
 from mlrun.utils import logger
 from mlrun.utils.notifications.notification import NotificationBase, NotificationTypes
-from mlrun.utils.notifications.notification_pusher import _NotificationPusherBase
+from mlrun.utils.notifications.notification_pusher import (
+    NotificationPusher,
+    _NotificationPusherBase,
+)
+
+
+class RunNotificationPusher(NotificationPusher):
+    def _prepare_notification_args(
+        self, run: mlrun.model.RunObject, notification_object: mlrun.model.Notification
+    ):
+        """
+        Prepare notification arguments for the notification pusher.
+        In the server side implementation, we need to mask the notification parameters on the task as they are
+        unmasked to extract the credentials required to send the notification.
+        """
+        message, severity, runs = super()._prepare_notification_args(
+            run, notification_object
+        )
+        for run in runs:
+            server.api.api.utils.mask_notification_params_on_task(
+                run, server.api.constants.MaskOperations.REDACT
+            )
+
+        return message, severity, runs
 
 
 class AlertNotificationPusher(_NotificationPusherBase):