diff --git a/pkg/providers/providers.go b/pkg/providers/providers.go index 169579a..7c09b1f 100644 --- a/pkg/providers/providers.go +++ b/pkg/providers/providers.go @@ -12,6 +12,7 @@ import ( "github.com/projectdiscovery/notify/pkg/providers/gotify" "github.com/projectdiscovery/notify/pkg/providers/notion" "github.com/projectdiscovery/notify/pkg/providers/pushover" + "github.com/projectdiscovery/notify/pkg/providers/rocketchat" "github.com/projectdiscovery/notify/pkg/providers/slack" "github.com/projectdiscovery/notify/pkg/providers/smtp" "github.com/projectdiscovery/notify/pkg/providers/teams" @@ -32,6 +33,7 @@ type ProviderOptions struct { Custom []*custom.Options `yaml:"custom,omitempty"` Gotify []*gotify.Options `yaml:"gotify,omitempty"` Notion []*notion.Options `yaml:"notion,omitempty"` + Rocketchat []*rocketchat.Options `yaml:"rocketchat,omitempty"` } // Provider is an interface implemented by providers @@ -134,6 +136,15 @@ func New(providerOptions *ProviderOptions, options *types.Options) (*Client, err client.providers = append(client.providers, provider) } + if providerOptions.Rocketchat != nil && (len(options.Providers) == 0 || sliceutil.Contains(options.Providers, "rocketchat")) { + + provider, err := rocketchat.New(providerOptions.Rocketchat, options.IDs) + if err != nil { + return nil, errors.Wrap(err, "could not create rocketchat provider client") + } + client.providers = append(client.providers, provider) + } + return client, nil } diff --git a/pkg/providers/rocketchat/rocketchat.go b/pkg/providers/rocketchat/rocketchat.go new file mode 100644 index 0000000..f8dcc46 --- /dev/null +++ b/pkg/providers/rocketchat/rocketchat.go @@ -0,0 +1,179 @@ +package rocketchat + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "unicode/utf8" + + "github.com/pkg/errors" + "go.uber.org/multierr" + + "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/notify/pkg/utils" + "github.com/projectdiscovery/notify/pkg/utils/httpreq" + sliceutil "github.com/projectdiscovery/utils/slice" +) + +// defaultMaxMessageLength matches Rocket.Chat's default Message_MaxAllowedSize setting. +const defaultMaxMessageLength = 5000 + +type Provider struct { + Rocketchat []*Options `yaml:"rocketchat,omitempty"` + counter int +} + +type Options struct { + ID string `yaml:"id,omitempty"` + RocketchatWebHookURL string `yaml:"rocketchat_webhook_url,omitempty"` + RocketchatChannel string `yaml:"rocketchat_channel,omitempty"` + RocketchatUsername string `yaml:"rocketchat_username,omitempty"` + RocketchatAvatar string `yaml:"rocketchat_avatar,omitempty"` + RocketchatEmoji string `yaml:"rocketchat_emoji,omitempty"` + RocketchatFormat string `yaml:"rocketchat_format,omitempty"` + // RocketchatAttachment sends the message as an attachment instead of plain text. + RocketchatAttachment bool `yaml:"rocketchat_attachment,omitempty"` +} + +func New(options []*Options, ids []string) (*Provider, error) { + provider := &Provider{} + + for _, o := range options { + if len(ids) == 0 || sliceutil.Contains(ids, o.ID) { + provider.Rocketchat = append(provider.Rocketchat, o) + } + } + + provider.counter = 0 + + return provider, nil +} + +func (p *Provider) Send(message, CliFormat string) error { + var rocketchatErr error + p.counter++ + + for _, pr := range p.Rocketchat { + msg := utils.FormatMessage(message, utils.SelectFormat(CliFormat, pr.RocketchatFormat), p.counter) + + if !strings.HasPrefix(pr.RocketchatWebHookURL, "http://") && !strings.HasPrefix(pr.RocketchatWebHookURL, "https://") { + err := errors.Wrap(fmt.Errorf("invalid rocketchat webhook URL"), + fmt.Sprintf("failed to send rocketchat notification for id: %s ", pr.ID)) + rocketchatErr = multierr.Append(rocketchatErr, err) + continue + } + + providerSucceeded := true + for _, chunk := range splitMessage(msg, defaultMaxMessageLength) { + payload := WebhookPayload{ + Channel: pr.RocketchatChannel, + Alias: pr.RocketchatUsername, + Avatar: pr.RocketchatAvatar, + Emoji: pr.RocketchatEmoji, + } + if pr.RocketchatAttachment { + payload.Attachments = []Attachment{{Text: chunk}} + } else { + payload.Text = chunk + } + + if err := send(pr.RocketchatWebHookURL, payload); err != nil { + providerSucceeded = false + err = errors.Wrap(err, fmt.Sprintf("failed to send rocketchat notification for id: %s", pr.ID)) + rocketchatErr = multierr.Append(rocketchatErr, err) + continue + } + } + if providerSucceeded { + gologger.Verbose().Msgf("rocketchat notification sent for id: %s", pr.ID) + } + } + return rocketchatErr +} + +func send(webhookURL string, payload WebhookPayload) error { + return sendWithClient(context.Background(), httpreq.NewClient(), webhookURL, payload) +} + +type httpDoer interface { + Do(*http.Request) (*http.Response, error) +} + +func sendWithClient(ctx context.Context, client httpDoer, webhookURL string, payload WebhookPayload) error { + jsonPayload, err := json.Marshal(payload) + if err != nil { + return errors.Wrap(err, "failed to marshal rocketchat payload") + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, webhookURL, bytes.NewBuffer(jsonPayload)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + if err != nil { + return err + } + + closeErr := resp.Body.Close() + if resp.StatusCode >= 400 { + statusErr := fmt.Errorf("received non-success status: %s", resp.Status) + if closeErr != nil { + return fmt.Errorf("%w; failed to close response body: %v", statusErr, closeErr) + } + return statusErr + } + if closeErr != nil { + return errors.Wrap(closeErr, "failed to close response body") + } + return nil +} + +// splitMessage splits msg into chunks no longer than limit, breaking on newlines +// where possible so multi-line messages aren't cut mid-line. +func splitMessage(msg string, limit int) []string { + if utf8.RuneCountInString(msg) <= limit { + return []string{msg} + } + + var chunks []string + var current strings.Builder + currentLen := 0 + + lines := strings.Split(msg, "\n") + for _, line := range lines { + runes := []rune(line) + for len(runes) > limit { + if currentLen > 0 { + chunks = append(chunks, current.String()) + current.Reset() + currentLen = 0 + } + chunks = append(chunks, string(runes[:limit])) + runes = runes[limit:] + } + + if currentLen > 0 && currentLen+1+len(runes) > limit { + chunks = append(chunks, current.String()) + current.Reset() + currentLen = 0 + } + + if currentLen > 0 { + current.WriteByte('\n') + currentLen++ + } + current.WriteString(string(runes)) + currentLen += len(runes) + } + + if current.Len() > 0 { + chunks = append(chunks, current.String()) + } + + return chunks +} diff --git a/pkg/providers/rocketchat/rocketchat_test.go b/pkg/providers/rocketchat/rocketchat_test.go new file mode 100644 index 0000000..a9f64cb --- /dev/null +++ b/pkg/providers/rocketchat/rocketchat_test.go @@ -0,0 +1,288 @@ +package rocketchat + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + "unicode/utf8" +) + +func TestNewFiltersByID(t *testing.T) { + options := []*Options{ + {ID: "soc", RocketchatWebHookURL: "https://chat.example.com/hooks/soc"}, + {ID: "vulns", RocketchatWebHookURL: "https://chat.example.com/hooks/vulns"}, + } + + provider, err := New(options, []string{"soc"}) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + if len(provider.Rocketchat) != 1 || provider.Rocketchat[0].ID != "soc" { + t.Fatalf("expected only id 'soc', got %+v", provider.Rocketchat) + } + + provider, err = New(options, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + if len(provider.Rocketchat) != 2 { + t.Fatalf("expected all options when ids is empty, got %d", len(provider.Rocketchat)) + } +} + +func TestSendSuccess(t *testing.T) { + var received WebhookPayload + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&received); err != nil { + t.Fatalf("failed to decode payload: %s", err) + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + provider, err := New([]*Options{ + { + ID: "soc", + RocketchatWebHookURL: server.URL, + RocketchatChannel: "#alertas-seguranca", + RocketchatUsername: "ProjectDiscovery", + RocketchatAvatar: "https://example.com/avatar.png", + RocketchatEmoji: ":ghost:", + }, + }, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + if err := provider.Send("host.example.com is vulnerable", ""); err != nil { + t.Fatalf("unexpected send error: %s", err) + } + + if received.Text != "host.example.com is vulnerable" { + t.Errorf("unexpected text: %q", received.Text) + } + if received.Channel != "#alertas-seguranca" { + t.Errorf("unexpected channel: %q", received.Channel) + } + if received.Alias != "ProjectDiscovery" { + t.Errorf("unexpected alias: %q", received.Alias) + } + if received.Avatar != "https://example.com/avatar.png" { + t.Errorf("unexpected avatar: %q", received.Avatar) + } + if received.Emoji != ":ghost:" { + t.Errorf("unexpected emoji: %q", received.Emoji) + } +} + +func TestSendAsAttachment(t *testing.T) { + var received WebhookPayload + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&received); err != nil { + t.Fatalf("failed to decode payload: %s", err) + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + provider, err := New([]*Options{ + { + ID: "soc", + RocketchatWebHookURL: server.URL, + RocketchatAttachment: true, + }, + }, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + if err := provider.Send("finding details", ""); err != nil { + t.Fatalf("unexpected send error: %s", err) + } + + if received.Text != "" { + t.Errorf("expected empty text when using attachments, got %q", received.Text) + } + if len(received.Attachments) != 1 || received.Attachments[0].Text != "finding details" { + t.Fatalf("unexpected attachments: %+v", received.Attachments) + } +} + +func TestSendInvalidWebhookURL(t *testing.T) { + provider, err := New([]*Options{ + {ID: "soc", RocketchatWebHookURL: "not-a-url"}, + }, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + err = provider.Send("message", "") + if err == nil || !strings.Contains(err.Error(), "invalid rocketchat webhook URL") { + t.Fatalf("expected invalid webhook URL error, got: %v", err) + } +} + +func TestSendNonSuccessStatus(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + + provider, err := New([]*Options{ + {ID: "soc", RocketchatWebHookURL: server.URL}, + }, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + if err := provider.Send("message", ""); err == nil { + t.Fatal("expected error for non-success status code") + } +} + +func TestSendWithClientHonorsContext(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + defer server.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + + err := sendWithClient(ctx, &http.Client{}, server.URL, WebhookPayload{Text: "message"}) + if err == nil { + t.Fatal("expected context cancellation error") + } +} + +func TestSendWithClientReportsBodyCloseError(t *testing.T) { + client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Body: closeErrorBody{}, + }, nil + })} + + err := sendWithClient(context.Background(), client, "http://example.com", WebhookPayload{Text: "message"}) + if err == nil || !strings.Contains(err.Error(), "close") { + t.Fatalf("expected response body close error, got %v", err) + } +} + +func TestSendSplitsLargeMessages(t *testing.T) { + var payloads []WebhookPayload + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var p WebhookPayload + if err := json.NewDecoder(r.Body).Decode(&p); err != nil { + t.Fatalf("failed to decode payload: %s", err) + } + payloads = append(payloads, p) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + provider, err := New([]*Options{ + {ID: "soc", RocketchatWebHookURL: server.URL}, + }, nil) + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + line := strings.Repeat("a", 100) + lines := make([]string, 200) + for i := range lines { + lines[i] = line + } + bigMessage := strings.Join(lines, "\n") + + if err := provider.Send(bigMessage, ""); err != nil { + t.Fatalf("unexpected send error: %s", err) + } + if len(payloads) <= 1 { + t.Fatalf("expected message to be split into multiple requests, got %d", len(payloads)) + } + + var rebuilt strings.Builder + for i, p := range payloads { + if len(p.Text) > defaultMaxMessageLength { + t.Errorf("chunk exceeds max length: %d", len(p.Text)) + } + if i > 0 { + rebuilt.WriteByte('\n') + } + rebuilt.WriteString(p.Text) + } + if rebuilt.String() != bigMessage { + t.Fatal("rebuilt message from chunks does not match original") + } +} + +func TestSplitMessage(t *testing.T) { + t.Run("under limit returns single chunk", func(t *testing.T) { + chunks := splitMessage("short message", 100) + if len(chunks) != 1 || chunks[0] != "short message" { + t.Fatalf("unexpected chunks: %+v", chunks) + } + }) + + t.Run("splits on newlines without breaking lines", func(t *testing.T) { + msg := strings.Repeat("a", 40) + "\n" + strings.Repeat("b", 40) + "\n" + strings.Repeat("c", 40) + chunks := splitMessage(msg, 50) + if len(chunks) != 3 { + t.Fatalf("expected 3 chunks, got %d: %+v", len(chunks), chunks) + } + for _, c := range chunks { + if len(c) > 50 { + t.Errorf("chunk exceeds limit: %q", c) + } + } + }) + + t.Run("hard splits a single line longer than limit", func(t *testing.T) { + msg := strings.Repeat("x", 250) + chunks := splitMessage(msg, 100) + if len(chunks) != 3 { + t.Fatalf("expected 3 chunks, got %d", len(chunks)) + } + if strings.Join(chunks, "") != msg { + t.Fatal("rejoined chunks do not match original message") + } + }) + + t.Run("does not split UTF-8 characters", func(t *testing.T) { + msg := strings.Repeat("é", 60) + chunks := splitMessage(msg, 50) + if len(chunks) != 2 { + t.Fatalf("expected 2 chunks, got %d", len(chunks)) + } + for _, chunk := range chunks { + if !utf8.ValidString(chunk) { + t.Fatalf("chunk is not valid UTF-8: %q", chunk) + } + if len([]rune(chunk)) > 50 { + t.Errorf("chunk exceeds rune limit: %d", len([]rune(chunk))) + } + } + if strings.Join(chunks, "") != msg { + t.Fatal("rejoined chunks do not match original message") + } + }) +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +type closeErrorBody struct{} + +func (closeErrorBody) Read([]byte) (int, error) { return 0, io.EOF } + +func (closeErrorBody) Close() error { return io.ErrUnexpectedEOF } diff --git a/pkg/providers/rocketchat/rocketchat_types.go b/pkg/providers/rocketchat/rocketchat_types.go new file mode 100644 index 0000000..77be4ab --- /dev/null +++ b/pkg/providers/rocketchat/rocketchat_types.go @@ -0,0 +1,18 @@ +package rocketchat + +// WebhookPayload is the body sent to a Rocket.Chat Incoming Webhook. +// Reference: https://docs.rocket.chat/docs/integrations#incoming-webhook-script +type WebhookPayload struct { + Text string `json:"text,omitempty"` + Channel string `json:"channel,omitempty"` + Alias string `json:"alias,omitempty"` + Avatar string `json:"avatar,omitempty"` + Emoji string `json:"emoji,omitempty"` + Attachments []Attachment `json:"attachments,omitempty"` +} + +// Attachment represents a Rocket.Chat message attachment. +type Attachment struct { + Text string `json:"text,omitempty"` + Color string `json:"color,omitempty"` +}