diff --git a/cmd/main.go b/cmd/main.go index 438f2f97..c357a1e9 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -11,7 +11,6 @@ import ( "github.com/checkmarx/2ms/lib" "sync" - "time" "github.com/checkmarx/2ms/plugins" "github.com/checkmarx/2ms/reporting" @@ -129,7 +128,7 @@ func Execute() { rootCmd.AddGroup(&cobra.Group{Title: group, ID: group}) for _, plugin := range allPlugins { - subCommand, err := plugin.DefineCommand(channels) + subCommand, err := plugin.DefineCommand(channels.Items, channels.Errors) if err != nil { log.Fatal().Msg(fmt.Sprintf("error while defining command for plugin %s: %s", plugin.GetName(), err.Error())) } @@ -170,22 +169,32 @@ func preRun(cmd *cobra.Command, args []string) { log.Fatal().Msg(err.Error()) } + channels.WaitGroup.Add(1) go func() { - for { - select { - case item := <-channels.Items: - report.TotalItemsScanned++ - channels.WaitGroup.Add(1) - go secrets.Detect(item, secretsChan, channels.WaitGroup, ignoreVar) - case secret := <-secretsChan: - report.TotalSecretsFound++ - report.Results[secret.ID] = append(report.Results[secret.ID], secret) - case err, ok := <-channels.Errors: - if !ok { - return - } - log.Fatal().Msg(err.Error()) - } + defer channels.WaitGroup.Done() + + wgItems := &sync.WaitGroup{} + for item := range channels.Items { + report.TotalItemsScanned++ + wgItems.Add(1) + go secrets.Detect(item, secretsChan, wgItems, ignoreVar) + } + wgItems.Wait() + close(secretsChan) + }() + + channels.WaitGroup.Add(1) + go func() { + defer channels.WaitGroup.Done() + for secret := range secretsChan { + report.TotalSecretsFound++ + report.Results[secret.ID] = append(report.Results[secret.ID], secret) + } + }() + + go func() { + for err := range channels.Errors { + log.Fatal().Msg(err.Error()) } }() } @@ -195,9 +204,6 @@ func postRun(cmd *cobra.Command, args []string) { cfg := config.LoadConfig("2ms", Version) - // Wait for last secret to be added to report - time.Sleep(time.Millisecond * timeSleepInterval) - // ------------------------------------- // Show Report if report.TotalItemsScanned > 0 { diff --git a/plugins/confluence.go b/plugins/confluence.go index 05f00fec..de579e51 100644 --- a/plugins/confluence.go +++ b/plugins/confluence.go @@ -46,7 +46,7 @@ func (p *ConfluencePlugin) GetAuthorizationHeader() string { return lib.CreateBasicAuthCredentials(p) } -func (p *ConfluencePlugin) DefineCommand(channels Channels) (*cobra.Command, error) { +func (p *ConfluencePlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { var confluenceCmd = &cobra.Command{ Use: fmt.Sprintf("%s --%s URL", p.GetName(), argUrl), Short: "Scan Confluence server", @@ -67,11 +67,14 @@ func (p *ConfluencePlugin) DefineCommand(channels Channels) (*cobra.Command, err confluenceCmd.Run = func(cmd *cobra.Command, args []string) { err := p.initialize(cmd) if err != nil { - channels.Errors <- fmt.Errorf("error while initializing confluence plugin: %w", err) + errors <- fmt.Errorf("error while initializing confluence plugin: %w", err) return } - p.getItems(channels.Items, channels.Errors, channels.WaitGroup) + wg := &sync.WaitGroup{} + p.getItems(items, errors, wg) + wg.Wait() + close(items) } return confluenceCmd, nil @@ -90,18 +93,14 @@ func (p *ConfluencePlugin) initialize(cmd *cobra.Command) error { } func (p *ConfluencePlugin) getItems(items chan Item, errs chan error, wg *sync.WaitGroup) { - p.getSpacesItems(items, errs, wg) -} - -func (p *ConfluencePlugin) getSpacesItems(items chan Item, errs chan error, wg *sync.WaitGroup) { spaces, err := p.getSpaces() if err != nil { errs <- err } for _, space := range spaces { - go p.getSpaceItems(items, errs, wg, space) wg.Add(1) + go p.getSpaceItems(items, errs, wg, space) } } diff --git a/plugins/discord.go b/plugins/discord.go index 6f6e5e12..ed49f591 100644 --- a/plugins/discord.go +++ b/plugins/discord.go @@ -38,7 +38,7 @@ func (p *DiscordPlugin) GetName() string { return "discord" } -func (p *DiscordPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { +func (p *DiscordPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { var discordCmd = &cobra.Command{ Use: fmt.Sprintf("%s --%s TOKEN --%s SERVER", p.GetName(), tokenFlag, serversFlag), Short: "Scan Discord server", @@ -63,11 +63,14 @@ func (p *DiscordPlugin) DefineCommand(channels Channels) (*cobra.Command, error) discordCmd.Run = func(cmd *cobra.Command, args []string) { err := p.initialize(cmd) if err != nil { - channels.Errors <- fmt.Errorf("discord plugin initialization failed: %w", err) + errors <- fmt.Errorf("discord plugin initialization failed: %w", err) return } - p.getItems(channels.Items, channels.Errors, channels.WaitGroup) + wg := &sync.WaitGroup{} + p.getItems(items, errors, wg) + wg.Wait() + close(items) } return discordCmd, nil diff --git a/plugins/filesystem.go b/plugins/filesystem.go index 520d69c1..c90e8744 100644 --- a/plugins/filesystem.go +++ b/plugins/filesystem.go @@ -29,14 +29,18 @@ func (p *FileSystemPlugin) GetName() string { return "filesystem" } -func (p *FileSystemPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { +func (p *FileSystemPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { var cmd = &cobra.Command{ Use: fmt.Sprintf("%s --%s PATH", p.GetName(), flagFolder), Short: "Scan local folder", Long: "Scan local folder for sensitive information", Run: func(cmd *cobra.Command, args []string) { log.Info().Msg("Folder plugin started") - p.getFiles(channels.Items, channels.Errors, channels.WaitGroup) + + wg := &sync.WaitGroup{} + p.getFiles(items, errors, wg) + wg.Wait() + close(items) }, } diff --git a/plugins/git.go b/plugins/git.go index 554f373a..a8f20d2a 100644 --- a/plugins/git.go +++ b/plugins/git.go @@ -4,6 +4,7 @@ import ( "fmt" "os" "strings" + "sync" "github.com/gitleaks/go-gitdiff/gitdiff" "github.com/rs/zerolog/log" @@ -29,8 +30,12 @@ func (p *GitPlugin) GetName() string { return "git" } -func (p *GitPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { - p.Channels = channels +func (p *GitPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { + p.Channels = Channels{ + Items: items, + Errors: errors, + WaitGroup: &sync.WaitGroup{}, + } command := &cobra.Command{ Use: fmt.Sprintf("%s ", p.GetName()), @@ -39,7 +44,9 @@ func (p *GitPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { Args: cobra.MatchAll(cobra.ExactArgs(1), validGitRepoArgs), Run: func(cmd *cobra.Command, args []string) { log.Info().Msg("Git plugin started") - p.scanGit(args[0], p.buildScanOptions(), channels.Items, channels.Errors) + p.scanGit(args[0], p.buildScanOptions(), p.Channels.Items, p.Channels.Errors) + p.WaitGroup.Wait() + close(items) }, } flags := command.Flags() diff --git a/plugins/paligo.go b/plugins/paligo.go index 3f2a28d1..bbeaafa6 100644 --- a/plugins/paligo.go +++ b/plugins/paligo.go @@ -7,6 +7,7 @@ import ( "net/http" "strconv" "strings" + "sync" "time" "github.com/checkmarx/2ms/lib" @@ -54,8 +55,12 @@ func (p *PaligoPlugin) GetName() string { return "paligo" } -func (p *PaligoPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { - p.Channels = channels +func (p *PaligoPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { + p.Channels = Channels{ + Items: items, + Errors: errors, + WaitGroup: &sync.WaitGroup{}, + } command := &cobra.Command{ Use: fmt.Sprintf("%s --%s %s --%s %s --%s %s", @@ -73,6 +78,8 @@ func (p *PaligoPlugin) DefineCommand(channels Channels) (*cobra.Command, error) } log.Info().Msg("Paligo plugin started") p.getItems() + p.WaitGroup.Wait() + close(items) }, } diff --git a/plugins/plugins.go b/plugins/plugins.go index a22e2f67..f8236614 100644 --- a/plugins/plugins.go +++ b/plugins/plugins.go @@ -27,5 +27,5 @@ type Channels struct { type IPlugin interface { GetName() string - DefineCommand(channels Channels) (*cobra.Command, error) + DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) } diff --git a/plugins/slack.go b/plugins/slack.go index d3862bd5..b88c9505 100644 --- a/plugins/slack.go +++ b/plugins/slack.go @@ -3,6 +3,7 @@ package plugins import ( "fmt" "strconv" + "sync" "time" "github.com/rs/zerolog/log" @@ -38,8 +39,12 @@ var ( messagesCountArg int ) -func (p *SlackPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { - p.Channels = channels +func (p *SlackPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) { + p.Channels = Channels{ + Items: items, + Errors: errors, + WaitGroup: &sync.WaitGroup{}, + } command := &cobra.Command{ Use: fmt.Sprintf("%s --%s TOKEN --%s TEAM", p.GetName(), slackTokenFlag, slackTeamFlag), @@ -47,6 +52,8 @@ func (p *SlackPlugin) DefineCommand(channels Channels) (*cobra.Command, error) { Long: "Scan Slack team for sensitive information.", Run: func(cmd *cobra.Command, args []string) { p.getItems() + p.Channels.WaitGroup.Wait() + close(items) }, }