diff --git a/docs/configuration/sinks/mattermost.rst b/docs/configuration/sinks/mattermost.rst index 5bd3866e6..52fbeb78c 100644 --- a/docs/configuration/sinks/mattermost.rst +++ b/docs/configuration/sinks/mattermost.rst @@ -75,6 +75,7 @@ Now we're ready to configure the Mattermost sink. token_id: (the token id visible in bot panel) channel: (the channel name you want to send messages to - either display name or channel name divided by hyphen (e.g. channel-name)) team_id: (OPTIONAL - this is only needed if your mattermost bot is not an admin) + user_agent: (OPTIONAL - custom User-Agent header for Mattermost API requests) Save the file and run diff --git a/src/robusta/core/sinks/mattermost/mattermost_sink.py b/src/robusta/core/sinks/mattermost/mattermost_sink.py index 4ec02e2cc..5ef6d27a3 100644 --- a/src/robusta/core/sinks/mattermost/mattermost_sink.py +++ b/src/robusta/core/sinks/mattermost/mattermost_sink.py @@ -16,10 +16,13 @@ def __init__(self, sink_config: MattermostSinkConfigWrapper, registry): token_id=sink_config.mattermost_sink.token_id, team=sink_config.mattermost_sink.team, team_id=sink_config.mattermost_sink.team_id, + user_agent=sink_config.mattermost_sink.user_agent, ) self.sender = MattermostSender( - cluster_name=self.cluster_name, account_id=self.account_id, client=client, - sink_params=sink_config.mattermost_sink + cluster_name=self.cluster_name, + account_id=self.account_id, + client=client, + sink_params=sink_config.mattermost_sink, ) def write_finding(self, finding: Finding, platform_enabled: bool): diff --git a/src/robusta/core/sinks/mattermost/mattermost_sink_params.py b/src/robusta/core/sinks/mattermost/mattermost_sink_params.py index 7cc4387a1..dfd9c0dda 100644 --- a/src/robusta/core/sinks/mattermost/mattermost_sink_params.py +++ b/src/robusta/core/sinks/mattermost/mattermost_sink_params.py @@ -14,6 +14,7 @@ class MattermostSinkParams(SinkBaseParams): channel: str team: Optional[str] team_id: Optional[str] + user_agent: Optional[str] @classmethod def _get_sink_type(cls): diff --git a/src/robusta/integrations/mattermost/client.py b/src/robusta/integrations/mattermost/client.py index 81fe103b6..ca11e91f2 100644 --- a/src/robusta/integrations/mattermost/client.py +++ b/src/robusta/integrations/mattermost/client.py @@ -6,12 +6,23 @@ _API_PREFIX = "api/v4" + class MattermostClient: channel_id: str bot_id: str team_id: Optional[str] - - def __init__(self, url: str, token: str, token_id: str, channel_name: str, team: Optional[str], team_id: Optional[str]): + user_agent: Optional[str] + + def __init__( + self, + url: str, + token: str, + token_id: str, + channel_name: str, + team: Optional[str], + team_id: Optional[str], + user_agent: Optional[str], + ): """ Set the Mattermost webhook url. """ @@ -19,12 +30,15 @@ def __init__(self, url: str, token: str, token_id: str, channel_name: str, team: self.token = token self.token_id = token_id self.team_id = team_id + self.user_agent = user_agent self.is_admin = self.is_admin_bot() self._init_setup(channel_name, team) def _send_mattermost_request(self, url: str, method: HttpMethod, **kwargs): headers = kwargs.pop("headers", {}) headers["Authorization"] = f"Bearer {self.token}" + if self.user_agent: + headers["User-Agent"] = self.user_agent return process_request(url, method, headers=headers, **kwargs) def _get_full_mattermost_url(self, endpoint: str) -> str: