Merge branch 'master' into rob-1126_headers_for_centralized_prometheus

Author: nherment

.github/workflows/test_robusta.yaml

@@ -19,7 +19,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: 3.11
 
       # setup a KIND cluster for tests which need a kubernetes image
       - name: Create k8s Kind Cluster

src/robusta/integrations/slack/sender.py

@@ -4,21 +4,22 @@
 import tempfile
 from datetime import datetime, timedelta
 from itertools import chain
-from typing import Any, Dict, List, Set
+from typing import Any, Dict, List, Optional, Set
 
 import certifi
 import humanize
 from dateutil import tz
 from slack_sdk import WebClient
-from slack_sdk.http_retry import all_builtin_retry_handlers
 from slack_sdk.errors import SlackApiError
+from slack_sdk.http_retry import all_builtin_retry_handlers
 
 from robusta.core.model.base_params import AIInvestigateParams, ResourceInfo
 from robusta.core.model.env_vars import (
     ADDITIONAL_CERTIFICATE,
+    HOLMES_ASK_SLACK_BUTTON_ENABLED,
     HOLMES_ENABLED,
     SLACK_REQUEST_TIMEOUT,
-    SLACK_TABLE_COLUMNS_LIMIT, HOLMES_ASK_SLACK_BUTTON_ENABLED,
+    SLACK_TABLE_COLUMNS_LIMIT,
 )
 from robusta.core.playbooks.internal.ai_integration import ask_holmes
 from robusta.core.reporting.base import Emojis, EnrichmentType, Finding, FindingStatus, LinkType
@@ -69,7 +70,12 @@ def __init__(self, slack_token: str, account_id: str, cluster_name: str, signing
             except Exception as e:
                 logging.exception(f"Failed to use custom certificate. {e}")
 
-        self.slack_client = WebClient(token=slack_token, ssl=ssl_context, timeout=SLACK_REQUEST_TIMEOUT, retry_handlers=all_builtin_retry_handlers())
+        self.slack_client = WebClient(
+            token=slack_token,
+            ssl=ssl_context,
+            timeout=SLACK_REQUEST_TIMEOUT,
+            retry_handlers=all_builtin_retry_handlers(),
+        )
         self.signing_key = signing_key
         self.account_id = account_id
         self.cluster_name = cluster_name
@@ -213,37 +219,74 @@ def __to_slack(self, block: BaseBlock, sink_name: str) -> List[SlackBlock]:
             logging.warning(f"cannot convert block of type {type(block)} to slack format block: {block}")
             return []  # no reason to crash the entire report
 
-    def __upload_file_to_slack(self, block: FileBlock, max_log_file_limit_kb: int) -> str:
+    def _upload_temp_file(self, f, file_reference, truncated_content: bytes, filename: str) -> Optional[str]:
+        """Helper to upload a file-like or file path to Slack."""
+        f.write(truncated_content)
+        f.flush()
+        f.seek(0)
+
+        result = self.slack_client.files_upload_v2(
+            title=filename,
+            file_uploads=[{"file": file_reference, "filename": filename, "title": filename}],
+        )
+        return result["file"]["permalink"]
+
+    def __upload_file_to_slack(self, block: FileBlock, max_log_file_limit_kb: int) -> Optional[str]:
+        """Upload a file to Slack and return a permalink to it."""
         truncated_content = block.truncate_content(max_file_size_bytes=max_log_file_limit_kb * 1000)
+        filename = block.filename
 
-        """Upload a file to slack and return a link to it"""
-        with tempfile.NamedTemporaryFile() as f:
-            f.write(truncated_content)
-            f.flush()
-            result = self.slack_client.files_upload_v2(title=block.filename, file=f.name, filename=block.filename)
-            return result["file"]["permalink"]
+        try:
+            with tempfile.NamedTemporaryFile() as f:
+                logging.debug("Trying NamedTemporaryFile for Slack upload")
+                return self._upload_temp_file(f, f.name, truncated_content, filename)
+        except Exception as e:
+            logging.debug(f"NamedTemporaryFile failed: {e}")
+        try:
+            SPOOLED_FILE_SIZE = 10 * 1000 * 1000  # 10MB to protect against OOM
+            with tempfile.SpooledTemporaryFile(max_size=SPOOLED_FILE_SIZE) as f:
+                logging.debug("Trying SpooledTemporaryFile for Slack upload")
+                return self._upload_temp_file(f, f, truncated_content, filename)
+        except Exception as e2:
+            logging.exception(f"SpooledTemporaryFile also failed: {e2}")
+            return None
 
     def prepare_slack_text(self, message: str, max_log_file_limit_kb: int, files: List[FileBlock] = []):
+        error_files = []
+
         if files:
             # it's a little annoying but it seems like files need to be referenced in `title` and not just `blocks`
             # in order to be actually shared. well, I'm actually not sure about that, but when I tried adding the files
             # to a separate block and not including them in `title` or the first block then the link was present but
             # the file wasn't actually shared and the link was broken
             uploaded_files = []
+
             for file_block in files:
                 # slack throws an error if you write empty files, so skip it
                 if len(file_block.contents) == 0:
                     continue
                 permalink = self.__upload_file_to_slack(file_block, max_log_file_limit_kb=max_log_file_limit_kb)
-                uploaded_files.append(f"* <{permalink} | {file_block.filename}>")
+                if permalink:
+                    uploaded_files.append(f"* <{permalink} | {file_block.filename}>")
+                else:
+                    error_files.append(file_block.filename)
 
             file_references = "\n".join(uploaded_files)
             message = f"{message}\n{file_references}"
 
         if len(message) == 0:
             return "empty-message"  # blank messages aren't allowed
+        message = Transformer.apply_length_limit(message, MAX_BLOCK_CHARS)
 
-        return Transformer.apply_length_limit(message, MAX_BLOCK_CHARS)
+        if error_files:
+            error_msg = (
+                "_Failed to send file(s) "
+                + ", ".join(error_files)
+                + " to slack._\n _See robusta-runner logs for details._"
+            )
+            return message, error_msg
+
+        return message, None
 
     def __send_blocks_to_slack(
         self,
@@ -266,9 +309,12 @@ def __send_blocks_to_slack(
         file_blocks.extend(Transformer.tableblock_to_fileblocks(other_blocks, SLACK_TABLE_COLUMNS_LIMIT))
         file_blocks.extend(Transformer.tableblock_to_fileblocks(report_attachment_blocks, SLACK_TABLE_COLUMNS_LIMIT))
 
-        message = self.prepare_slack_text(
+        message, error_msg = self.prepare_slack_text(
             title, max_log_file_limit_kb=sink_params.max_log_file_limit_kb, files=file_blocks
         )
+        if error_msg:
+            other_blocks.append(MarkdownBlock(error_msg))
+
         output_blocks = []
         for block in other_blocks:
             output_blocks.extend(self.__to_slack(block, sink_params.name))
@@ -310,15 +356,14 @@ def __send_blocks_to_slack(
                 f"error sending message to slack\ne={e}\ntext={message}\nchannel={channel}\nblocks={*output_blocks,}\nattachment_blocks={*attachment_blocks,}"
             )
 
-
     def __limit_labels_size(self, labels: dict, max_size: int = 1000) -> dict:
         # slack can only send 2k tokens in a callback so the labels are limited in size
 
         low_priority_labels = ["job", "prometheus", "severity", "service"]
         current_length = len(str(labels))
         if current_length <= max_size:
             return labels
-        
+
         limited_labels = copy.deepcopy(labels)
 
         # first remove the low priority labels if needed
@@ -348,7 +393,7 @@ def __create_holmes_callback(self, finding: Finding) -> CallbackBlock:
             "robusta_issue_id": str(finding.id),
             "issue_type": finding.aggregation_key,
             "source": finding.source.name,
-            "labels": self.__limit_labels_size(labels=finding.subject.labels)
+            "labels": self.__limit_labels_size(labels=finding.subject.labels),
         }
 
         return CallbackBlock(
@@ -714,12 +759,7 @@ def update_slack_message(self, channel: str, ts: str, blocks: list, text: str =
                 return
 
             # Call Slack's chat_update method
-            resp = self.slack_client.chat_update(
-                channel=channel,
-                ts=ts,
-                text=text,
-                blocks=blocks
-            )
+            resp = self.slack_client.chat_update(channel=channel, ts=ts, text=text, blocks=blocks)
             logging.debug(f"Message updated successfully: {resp['ts']}")
             return resp["ts"]
 

tests/test_slack.py

@@ -1,4 +1,7 @@
-from robusta.api import Finding, MarkdownBlock, SlackSender, TableBlock
+import logging
+from unittest.mock import patch
+
+from robusta.api import FileBlock, Finding, MarkdownBlock, SlackSender, TableBlock
 from robusta.core.sinks.slack.slack_sink_params import SlackSinkParams
 from tests.config import CONFIG
 from tests.utils.slack_utils import SlackChannel
@@ -7,6 +10,10 @@
 TEST_CLUSTER = "test cluster"
 TEST_KEY = "test key"
 
+TEST_FILE_NAME = "test.txt"
+TEST_FILE_CONTENT = "test file content"
+TEST_FINDING_TITLE = "Test Text File Upload"
+
 
 def test_send_to_slack(slack_channel: SlackChannel):
     slack_sender = SlackSender(
@@ -48,3 +55,84 @@ def test_long_table_columns(slack_channel: SlackChannel):
     )
     slack_params = SlackSinkParams(name="test_slack", slack_channel=slack_channel.channel_name, api_key="")
     slack_sender.send_finding_to_slack(finding, slack_params, False)
+
+
+def test_send_file_spooled_tempfile_fails(slack_channel: SlackChannel):
+    slack_sender = SlackSender(
+        CONFIG.PYTEST_IN_CLUSTER_SLACK_TOKEN, TEST_ACCOUNT, TEST_CLUSTER, TEST_KEY, slack_channel.channel_name
+    )
+
+    # Test with a text file
+    finding = Finding(title=TEST_FINDING_TITLE, aggregation_key="TestTextFileUpload")
+    finding.add_enrichment([FileBlock(TEST_FILE_NAME, TEST_FILE_CONTENT)])
+
+    slack_params = SlackSinkParams(name="test_slack", slack_channel=slack_channel.channel_name, api_key="")
+
+    # verify NamedTemporaryFile sending works
+    # verify SpooledTemporaryFile sending works
+    with patch("tempfile.NamedTemporaryFile", side_effect=FileNotFoundError("No usable temporary directory found")):
+        slack_sender.send_finding_to_slack(finding, slack_params, False)
+
+    # Verify that the message contains the finding title but not the file content
+    latest_message = slack_channel.get_latest_message()
+    logging.warning(latest_message)
+    assert TEST_FINDING_TITLE in latest_message
+    assert TEST_FILE_NAME in latest_message
+
+
+def test_send_file_named_tempfile_fails(slack_channel: SlackChannel):
+    slack_sender = SlackSender(
+        CONFIG.PYTEST_IN_CLUSTER_SLACK_TOKEN, TEST_ACCOUNT, TEST_CLUSTER, TEST_KEY, slack_channel.channel_name
+    )
+
+    finding = Finding(title=TEST_FINDING_TITLE, aggregation_key="TestTextFileUpload")
+    finding.add_enrichment([FileBlock(TEST_FILE_NAME, TEST_FILE_CONTENT)])
+
+    slack_params = SlackSinkParams(name="test_slack", slack_channel=slack_channel.channel_name, api_key="")
+
+    # Simulate tempfile failure
+    with patch("tempfile.SpooledTemporaryFile", side_effect=FileNotFoundError("No usable temporary directory found")):
+        slack_sender.send_finding_to_slack(finding, slack_params, False)
+
+    # Check the Slack message
+    latest_message = slack_channel.get_latest_message()
+    assert TEST_FINDING_TITLE in latest_message
+    assert TEST_FILE_NAME in latest_message
+
+
+def test_temporary_file_creation_failure(slack_channel: SlackChannel):
+    slack_sender = SlackSender(
+        CONFIG.PYTEST_IN_CLUSTER_SLACK_TOKEN, TEST_ACCOUNT, TEST_CLUSTER, TEST_KEY, slack_channel.channel_name
+    )
+
+    # Test with a text file
+    finding = Finding(title=TEST_FINDING_TITLE, aggregation_key="TestTextFileUpload")
+    finding.add_enrichment(
+        [FileBlock(TEST_FILE_NAME, TEST_FILE_CONTENT.encode()), FileBlock("file2.txt", TEST_FILE_CONTENT.encode())]
+    )
+
+    slack_params = SlackSinkParams(name="test_slack", slack_channel=slack_channel.channel_name, api_key="")
+
+    # Mock NamedTemporaryFile to raise an exception
+    with patch("tempfile.NamedTemporaryFile", side_effect=FileNotFoundError("No usable temporary directory found")):
+        with patch("tempfile.SpooledTemporaryFile", side_effect=FileNotFoundError("Cant create spooled file")):
+            slack_sender.send_finding_to_slack(finding, slack_params, False)
+
+        # Verify that the message contains the finding title but not the file content
+        latest_message = slack_channel.get_latest_message()
+        assert TEST_FINDING_TITLE in latest_message
+        assert TEST_FILE_NAME not in latest_message  # File should not be included
+
+    # Test with a binary file (PNG)
+    png_content = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82"  # 1x1 transparent PNG
+    finding = Finding(title="Test PNG File Upload", aggregation_key="TestPNGFileUpload")
+    finding.add_enrichment([FileBlock("test.png", png_content)])
+
+    with patch("tempfile.NamedTemporaryFile", side_effect=FileNotFoundError("No usable temporary directory found")):
+        with patch("tempfile.SpooledTemporaryFile", side_effect=FileNotFoundError("Cant create spooled file")):
+            slack_sender.send_finding_to_slack(finding, slack_params, False)
+
+        # Verify that the message contains the finding title but not the file content
+        latest_message = slack_channel.get_latest_message()
+        assert "Test PNG File Upload" in latest_message
+        assert "test.png" not in latest_message  # File should not be included

tests/utils/slack_utils.py

@@ -1,3 +1,5 @@
+import logging
+
 from slack_sdk import WebClient
 
 
@@ -25,7 +27,12 @@ def _create_or_join_channel(client: WebClient, channel_name: str) -> str:
         """Creates or joins the specified slack channel and returns its channel_id"""
         for channel in client.conversations_list()["channels"]:
             if channel["name"] == channel_name:
-                # TODO: join the channel if necessary
+                try:
+                    # Attempt to join the channel if not already joined
+                    client.conversations_join(channel=channel["id"])
+                except Exception as e:
+                    # It's ok if already in channel or can't join (e.g., private and no permission)
+                    logging.warning(f"Could not join channel {channel_name}: {e}")
                 return channel["id"]
 
         # TODO: make this a private channel not a public channel. it shouldn't be visible to anyone who joins the public