Skip to content

Commit

Permalink
feat: discord plugin (#53)
Browse files Browse the repository at this point in the history
Resolves #31

### Features

- We expect the user to give us a *Personal Access Token*. This token
can be retrieved from the browser *Dev Tool*, or by authenticating a
*Discord App*. See
#31 (comment).
- The user must give at least one `--discord-server`, we will not loop
over all the user's servers.
- The *Server* in Discord called *Guild* in the API.
- If the user doesn't give `--discord-channel`, we will scan all the
channels in a server.
- We will scan all messages until `--discord-duration` or
`--discord-messages-count` (the closest one).
- **Threads**: Only *bots* can get all the **threads** from a
**channel**. As we currently use *Personal Access Token*, we can get the
**thread** from the **message** started it.
So we will scan the **threads** that *started* in the time/limit
arguments (as they are messages that scanned), and for each **thread**
we will scan the messages in the time/limit requirements.

### Questions

- It is a little confusing if you give multiple **servers** and multiple
**channels**, because each **channel** relates to its **server**, and
also if you will give **channels** only for one **server**, the other
**server** will not be scanned at all.
Compared to *Confluence*, which has **servers** and **spaces**, we are
not scanning multiple **servers** in one scan.

Waits for #52
  • Loading branch information
Baruch Odem (Rothkoff) authored May 11, 2023
1 parent ebb3e7d commit 0a85890
Show file tree
Hide file tree
Showing 5 changed files with 312 additions and 10 deletions.
9 changes: 6 additions & 3 deletions cmd/main.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,16 @@
package cmd

import (
"github.com/checkmarx/2ms/plugins"
"github.com/checkmarx/2ms/reporting"
"github.com/checkmarx/2ms/secrets"
"os"
"strings"

"sync"
"time"

"github.com/checkmarx/2ms/plugins"
"github.com/checkmarx/2ms/reporting"
"github.com/checkmarx/2ms/secrets"

"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/spf13/cobra"
Expand All @@ -27,6 +29,7 @@ var Version = ""

var allPlugins = []plugins.IPlugin{
&plugins.ConfluencePlugin{},
&plugins.DiscordPlugin{},
&plugins.RepositoryPlugin{},
}

Expand Down
3 changes: 3 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module github.com/checkmarx/2ms
go 1.20

require (
github.com/bwmarrin/discordgo v0.27.1
github.com/rs/zerolog v1.29.0
github.com/spf13/cobra v1.6.1
github.com/stretchr/testify v1.8.1
Expand All @@ -16,6 +17,7 @@ require (
github.com/fatih/semgroup v1.2.0 // indirect
github.com/fsnotify/fsnotify v1.6.0 // indirect
github.com/gitleaks/go-gitdiff v0.8.0 // indirect
github.com/gorilla/websocket v1.4.2 // indirect
github.com/h2non/filetype v1.1.3 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
Expand All @@ -38,6 +40,7 @@ require (
github.com/spf13/pflag v1.0.5 // indirect
github.com/spf13/viper v1.15.0 // indirect
github.com/subosito/gotenv v1.4.2 // indirect
golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa // indirect
golang.org/x/sync v0.1.0 // indirect
golang.org/x/sys v0.6.0 // indirect
golang.org/x/text v0.8.0 // indirect
Expand Down
5 changes: 5 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03
github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
github.com/bwmarrin/discordgo v0.27.1 h1:ib9AIc/dom1E/fSIulrBwnez0CToJE113ZGt4HoliGY=
github.com/bwmarrin/discordgo v0.27.1/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
github.com/charmbracelet/lipgloss v0.7.1 h1:17WMwi7N1b1rVWOjMT+rCh7sQkvDU75B2hbZpc5Kc1E=
github.com/charmbracelet/lipgloss v0.7.1/go.mod h1:yG0k3giv8Qj8edTCbbg6AlQ5e8KNWpFujkNawKNhE2c=
Expand Down Expand Up @@ -127,6 +129,8 @@ github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/h2non/filetype v1.1.3 h1:FKkx9QbD7HR/zjK1Ia5XiBsq9zdLi5Kf3zGyFTAFkGg=
github.com/h2non/filetype v1.1.3/go.mod h1:319b3zT68BvV+WRj7cwy856M2ehB3HqNOt6sy1HndBY=
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
Expand Down Expand Up @@ -231,6 +235,7 @@ golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8U
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa h1:zuSxTR4o9y82ebqCUJYNGJbGPo6sKVl54f/TVDObg1c=
golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
Expand Down
16 changes: 9 additions & 7 deletions plugins/confluence.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,15 @@ import (
"github.com/spf13/cobra"
)

const argConfluence = "confluence"
const argConfluenceSpaces = "confluence-spaces"
const argConfluenceUsername = "confluence-username"
const argConfluenceToken = "confluence-token"
const argConfluenceHistory = "history"
const confluenceDefaultWindow = 25
const confluenceMaxRequests = 500
const (
argConfluence = "confluence"
argConfluenceSpaces = "confluence-spaces"
argConfluenceUsername = "confluence-username"
argConfluenceToken = "confluence-token"
argConfluenceHistory = "history"
confluenceDefaultWindow = 25
confluenceMaxRequests = 500
)

type ConfluencePlugin struct {
Plugin
Expand Down
289 changes: 289 additions & 0 deletions plugins/discord.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,289 @@
package plugins

import (
"fmt"
"sync"
"time"

"github.com/bwmarrin/discordgo"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/spf13/cobra"
)

const (
discordTokenFlag = "discord-token"
discordServersFlag = "discord-server"
discordChannelsFlag = "discord-channel"
discordFromDateFlag = "discord-duration"
discordMessagesCountFlag = "discord-messages-count"
)

const defaultDateFrom = time.Hour * 24 * 14

type DiscordPlugin struct {
Enabled bool
Token string
Guilds []string
Channels []string
Count int
BackwardDuration time.Duration
Session *discordgo.Session

errChan chan error
itemChan chan Item
waitGroup *sync.WaitGroup
}

func (p *DiscordPlugin) DefineCommandLineArgs(cmd *cobra.Command) error {
flags := cmd.Flags()

flags.String(discordTokenFlag, "", "discord token")
flags.StringArray(discordServersFlag, []string{}, "discord servers")
flags.StringArray(discordChannelsFlag, []string{}, "discord channels")
flags.Duration(discordFromDateFlag, defaultDateFrom, "discord from date")
flags.Int(discordMessagesCountFlag, 0, "discord messages count")

cmd.MarkFlagsRequiredTogether(discordTokenFlag, discordServersFlag)

return nil
}

func (p *DiscordPlugin) Initialize(cmd *cobra.Command) error {
flags := cmd.Flags()
token, _ := flags.GetString(discordTokenFlag)
if token == "" {
return fmt.Errorf("discord token arg is missing. Plugin initialization failed")
}

guilds, _ := flags.GetStringArray(discordServersFlag)
if len(guilds) == 0 {
return fmt.Errorf("discord servers arg is missing. Plugin initialization failed")
}

channels, _ := flags.GetStringArray(discordChannelsFlag)
if len(channels) == 0 {
log.Warn().Msg("discord channels not provided. Will scan all channels")
}

fromDate, _ := flags.GetDuration(discordFromDateFlag)
count, _ := flags.GetInt(discordMessagesCountFlag)
if count == 0 && fromDate == 0 {
return fmt.Errorf("discord messages count or from date arg is missing. Plugin initialization failed")
}

p.Token = token
p.Guilds = guilds
p.Channels = channels
p.Count = count
p.BackwardDuration = fromDate
p.Enabled = true

return nil
}

func (p *DiscordPlugin) IsEnabled() bool {
return p.Enabled
}

func (p *DiscordPlugin) GetItems(itemsChan chan Item, errChan chan error, wg *sync.WaitGroup) {
defer wg.Done()

p.errChan = errChan
p.itemChan = itemsChan
p.waitGroup = wg

err := p.getDiscordReady()
if err != nil {
errChan <- err
return
}

guilds := p.getGuildsByNameOrIDs()
log.Info().Msgf("Found %d guilds", len(guilds))

wg.Add(len(guilds))
for _, guild := range guilds {
go p.readGuildMessages(guild)
}
}

func (p *DiscordPlugin) getDiscordReady() (err error) {
p.Session, err = discordgo.New(p.Token)
if err != nil {
return err
}

p.Session.StateEnabled = true
ready := make(chan error)
p.Session.AddHandlerOnce(func(s *discordgo.Session, r *discordgo.Ready) {
ready <- nil
})
go func() {
err := p.Session.Open()
if err != nil {
ready <- err
}
}()
time.AfterFunc(time.Second*10, func() {
ready <- fmt.Errorf("discord session timeout")
})

err = <-ready
if err != nil {
return err
}

return nil
}

func (p *DiscordPlugin) getGuildsByNameOrIDs() []*discordgo.Guild {
var result []*discordgo.Guild

for _, guild := range p.Guilds {
for _, g := range p.Session.State.Guilds {
if g.Name == guild || g.ID == guild {
result = append(result, g)
}
}
}

return result
}

func (p *DiscordPlugin) readGuildMessages(guild *discordgo.Guild) {
defer p.waitGroup.Done()

guildLogger := log.With().Str("guild", guild.Name).Logger()
guildLogger.Debug().Send()

selectedChannels := p.getChannelsByNameOrIDs(guild)
guildLogger.Info().Msgf("Found %d channels", len(selectedChannels))

p.waitGroup.Add(len(selectedChannels))
for _, channel := range selectedChannels {
go p.readChannelMessages(channel)
}
}

func (p *DiscordPlugin) getChannelsByNameOrIDs(guild *discordgo.Guild) []*discordgo.Channel {
var result []*discordgo.Channel
if len(p.Channels) == 0 {
return guild.Channels
}

for _, channel := range p.Channels {
for _, c := range guild.Channels {
if c.Name == channel || c.ID == channel {
result = append(result, c)
}
}
}

return result
}

func (p *DiscordPlugin) readChannelMessages(channel *discordgo.Channel) {
defer p.waitGroup.Done()

channelLogger := log.With().Str("guildID", channel.GuildID).Str("channel", channel.Name).Logger()
channelLogger.Debug().Send()

permission, err := p.Session.UserChannelPermissions(p.Session.State.User.ID, channel.ID)
if err != nil {
if err, ok := err.(*discordgo.RESTError); ok {
if err.Message.Code == 50001 {
channelLogger.Debug().Msg("No read permissions")
return
}
}

channelLogger.Error().Err(err).Msg("Failed to get permissions")
p.errChan <- err
return
}
if permission&discordgo.PermissionViewChannel == 0 {
channelLogger.Debug().Msg("No read permissions")
return
}
if channel.Type != discordgo.ChannelTypeGuildText {
channelLogger.Debug().Msg("Not a text channel")
return
}

messages, err := p.getMessages(channel.ID, channelLogger)
if err != nil {
channelLogger.Error().Err(err).Msg("Failed to get messages")
p.errChan <- err
return
}
channelLogger.Info().Msgf("Found %d messages", len(messages))

items := convertMessagesToItems(channel.GuildID, &messages)
for _, item := range *items {
p.itemChan <- item
}
}

func (p *DiscordPlugin) getMessages(channelID string, logger zerolog.Logger) ([]*discordgo.Message, error) {
var messages []*discordgo.Message
threadMessages := []*discordgo.Message{}

var beforeID string

m, err := p.Session.ChannelMessages(channelID, 100, beforeID, "", "")
if err != nil {
return nil, err
}

lastMessage := false
for len(m) > 0 && !lastMessage {

for _, message := range m {

timeSince := time.Since(message.Timestamp)
if p.BackwardDuration > 0 && timeSince > p.BackwardDuration {
logger.Debug().Msgf("Reached time limit (%s). Last message is %s old", p.BackwardDuration.String(), timeSince.Round(time.Hour).String())
lastMessage = true
break
}

if p.Count > 0 && len(messages) == p.Count {
logger.Debug().Msgf("Reached message count (%d)", p.Count)
lastMessage = true
break
}

if message.Thread != nil {
logger.Info().Msgf("Found thread %s", message.Thread.Name)
tMgs, err := p.getMessages(message.Thread.ID, logger.With().Str("thread", message.Thread.Name).Logger())
if err != nil {
return nil, err
}
threadMessages = append(threadMessages, tMgs...)
}

messages = append(messages, message)
beforeID = message.ID
}

m, err = p.Session.ChannelMessages(channelID, 100, beforeID, "", "")
if err != nil {
return nil, err
}
}

return append(messages, threadMessages...), nil
}

func convertMessagesToItems(guildId string, messages *[]*discordgo.Message) *[]Item {
items := []Item{}
for _, message := range *messages {
items = append(items, Item{
Content: message.Content,
Source: fmt.Sprintf("https://discord.com/channels/%s/%s/%s", guildId, message.ChannelID, message.ID),
ID: message.ID,
})
}
return &items
}

0 comments on commit 0a85890

Please sign in to comment.