From b151331651fa26f61730380858832a6afec264a3 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 4 Jun 2025 13:17:51 +0000 Subject: [PATCH] feat: Add support for Telegram message threads This commit introduces the ability to send notifications to specific message threads (topics) within Telegram groups. The following changes were made: - Added a new optional `telegram_thread_id` field to the Telegram provider configuration in `pkg/providers/telegram/telegram.go`. - Modified the `Send` function in the Telegram provider to append the `telegram_thread_id` to the `telegram_chat_id` in the Shoutrrr URL if provided. The format used is `chat_id:thread_id`. - Updated `README.md` to include documentation for the new `telegram_thread_id` field and provided a link to the Shoutrrr documentation for instructions on obtaining the `message_thread_id`. - Added unit tests in `pkg/providers/telegram/telegram_test.go` to verify the correct construction of the Telegram URL with and without the `telegram_thread_id`. This involved refactoring `pkg/providers/telegram/telegram.go` to allow mocking of the `shoutrrr.Send` function for testing purposes. This feature allows you to direct notifications to specific topics in your Telegram groups, providing more granular control over where notifications appear. --- README.md | 1 + go.mod | 3 + pkg/providers/telegram/telegram.go | 13 ++- pkg/providers/telegram/telegram_test.go | 104 ++++++++++++++++++++++++ 4 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 pkg/providers/telegram/telegram_test.go diff --git a/README.md b/README.md index 9260952..e483769 100644 --- a/README.md +++ b/README.md @@ -113,6 +113,7 @@ telegram: - id: "tel" telegram_api_key: "XXXXXXXXXXXX" telegram_chat_id: "XXXXXXXX" + telegram_thread_id: "YYYYYYYYYY" # Optional: The message thread ID to send messages to (https://www.shoutrrr.com/services/telegram/#message-threads) telegram_format: "{{data}}" telegram_parsemode: "Markdown" # None/Markdown/MarkdownV2/HTML (https://core.telegram.org/bots/api#formatting-options) diff --git a/go.mod b/go.mod index 746bf0d..2cf7431 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/projectdiscovery/goflags v0.1.64 github.com/projectdiscovery/gologger v1.1.29 github.com/projectdiscovery/utils v0.2.16 + github.com/stretchr/testify v1.9.0 go.uber.org/multierr v1.11.0 go.uber.org/ratelimit v0.3.0 gopkg.in/yaml.v3 v3.0.1 @@ -35,6 +36,7 @@ require ( github.com/charmbracelet/x/ansi v0.3.2 // indirect github.com/cheggaaa/pb/v3 v3.1.4 // indirect github.com/cnf/structhash v0.0.0-20201127153200-e1b16c1ebc08 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/dlclark/regexp2 v1.11.4 // indirect github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5 // indirect github.com/fatih/color v1.15.0 // indirect @@ -67,6 +69,7 @@ require ( github.com/muesli/termenv v0.15.3-0.20240618155329-98d742f6907a // indirect github.com/nwaples/rardecode v1.1.3 // indirect github.com/pierrec/lz4/v4 v4.1.2 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect github.com/projectdiscovery/blackrock v0.0.1 // indirect github.com/projectdiscovery/machineid v0.0.0-20240226150047-2e2c51e35983 // indirect diff --git a/pkg/providers/telegram/telegram.go b/pkg/providers/telegram/telegram.go index 54edba9..fa60b9b 100644 --- a/pkg/providers/telegram/telegram.go +++ b/pkg/providers/telegram/telegram.go @@ -6,12 +6,14 @@ import ( "github.com/containrrr/shoutrrr" "github.com/pkg/errors" "go.uber.org/multierr" - "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/notify/pkg/utils" sliceutil "github.com/projectdiscovery/utils/slice" ) +// shoutrrrSendFunc is a package-level variable to allow mocking in tests. +var shoutrrrSendFunc = shoutrrr.Send + type Provider struct { Telegram []*Options `yaml:"telegram,omitempty"` counter int @@ -21,6 +23,7 @@ type Options struct { ID string `yaml:"id,omitempty"` TelegramAPIKey string `yaml:"telegram_api_key,omitempty"` TelegramChatID string `yaml:"telegram_chat_id,omitempty"` + TelegramThreadID string `yaml:"telegram_thread_id,omitempty"` TelegramFormat string `yaml:"telegram_format,omitempty"` TelegramParseMode string `yaml:"telegram_parsemode,omitempty"` } @@ -47,8 +50,12 @@ func (p *Provider) Send(message, CliFormat string) error { if pr.TelegramParseMode == "" { pr.TelegramParseMode = "None" } - url := fmt.Sprintf("telegram://%s@telegram?channels=%s&parsemode=%s", pr.TelegramAPIKey, pr.TelegramChatID, pr.TelegramParseMode) - err := shoutrrr.Send(url, msg) + telegramChatID := pr.TelegramChatID + if pr.TelegramThreadID != "" { + telegramChatID = fmt.Sprintf("%s:%s", pr.TelegramChatID, pr.TelegramThreadID) + } + url := fmt.Sprintf("telegram://%s@telegram?channels=%s&parsemode=%s", pr.TelegramAPIKey, telegramChatID, pr.TelegramParseMode) + err := shoutrrrSendFunc(url, msg) if err != nil { err = errors.Wrap(err, fmt.Sprintf("failed to send telegram notification for id: %s ", pr.ID)) TelegramErr = multierr.Append(TelegramErr, err) diff --git a/pkg/providers/telegram/telegram_test.go b/pkg/providers/telegram/telegram_test.go new file mode 100644 index 0000000..db31c92 --- /dev/null +++ b/pkg/providers/telegram/telegram_test.go @@ -0,0 +1,104 @@ +package telegram + +import ( + "net/url" + "testing" + + "github.com/stretchr/testify/require" +) + +// capturedURL will store the URL passed to the mocked shoutrrrSendFunc +var capturedURL string + +// TestTelegramSendURLWithThreadID tests the Send method of the Telegram provider, +// focusing on how the URL is constructed with and without a thread ID. +func TestTelegramSendURLWithThreadID(t *testing.T) { + // Store the original shoutrrrSendFunc and defer its restoration + originalSendFunc := shoutrrrSendFunc + defer func() { + shoutrrrSendFunc = originalSendFunc + }() + + // Mock shoutrrrSendFunc to capture the URL and avoid actual sending + shoutrrrSendFunc = func(serviceURL string, message string) error { + capturedURL = serviceURL + return nil // Simulate success + } + + tests := []struct { + name string + options Options + expectedChatIDInURL string + }{ + { + name: "with thread_id", + options: Options{ + ID: "test-with-thread", + TelegramAPIKey: "testAPIKey", + TelegramChatID: "testChatID", + TelegramThreadID: "testThreadID", + }, + expectedChatIDInURL: "testChatID:testThreadID", + }, + { + name: "without thread_id", + options: Options{ + ID: "test-without-thread", + TelegramAPIKey: "testAPIKey2", + TelegramChatID: "testChatID2", + }, + expectedChatIDInURL: "testChatID2", + }, + { + name: "with thread_id but empty", + options: Options{ + ID: "test-with-empty-thread", + TelegramAPIKey: "testAPIKey3", + TelegramChatID: "testChatID3", + TelegramThreadID: "", // Explicitly empty + }, + expectedChatIDInURL: "testChatID3", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capturedURL = "" // Reset captured URL for each test run + provider, err := New([]*Options{&tt.options}, nil) + require.NoError(t, err, "New() should not return an error") + require.NotNil(t, provider, "New() should return a provider") + require.Len(t, provider.Telegram, 1, "Provider should have one Telegram option") + + err = provider.Send("test message", "") + require.NoError(t, err, "Send() should not return an error") + + parsedURL, err := url.Parse(capturedURL) + require.NoError(t, err, "Captured URL should be parseable") + + channels := parsedURL.Query().Get("channels") + require.Equal(t, tt.expectedChatIDInURL, channels, "Chat ID in URL does not match expected") + + // Verify other parts of the URL + expectedScheme := "telegram" + require.Equal(t, expectedScheme, parsedURL.Scheme, "URL scheme does not match") + + // Check API Key (Username part of Userinfo) + expectedAPIKey := tt.options.TelegramAPIKey + require.NotNil(t, parsedURL.User, "URL Userinfo should not be nil") + actualAPIKey := parsedURL.User.Username() + require.Equal(t, expectedAPIKey, actualAPIKey, "URL API key (username) does not match") + + // Check Host part + expectedHost := "telegram" + require.Equal(t, expectedHost, parsedURL.Host, "URL host does not match") + + parseMode := parsedURL.Query().Get("parsemode") + // Default ParseMode is "None" if not specified in options + expectedParseMode := "None" + if tt.options.TelegramParseMode != "" { + expectedParseMode = tt.options.TelegramParseMode + } + require.Equal(t, expectedParseMode, parseMode, "URL parsemode does not match") + }) + } +}