Make detectors configurable (#1084)

* Make detectors configurable.

* remove redundant check.

* add number of detectors.

* update comment.

* remove reflect.

* inline key.

* replace name w/ type.

* remove temp var.

* fix test name.

* fix engine start.

* add filter unverified to engine.

* reorder engine args.

* Address comments.

* Add include and exclude.

* update comments.

* add comment.

* add comment.

main.go

@@ -50,6 +50,8 @@ var (
 	archiveMaxSize       = cli.Flag("archive-max-size", "Maximum size of archive to scan.").Bytes()
 	archiveMaxDepth      = cli.Flag("archive-max-depth", "Maximum depth of archive to scan.").Int()
 	archiveTimeout       = cli.Flag("archive-timeout", "Maximum time to spend extracting an archive.").Duration()
+	includeDetectors     = cli.Flag("include-detectors", "Comma separated list of detector types to include. Types must exist in detectors.proto").String()
+	excludeDetectors     = cli.Flag("exclude-detectors", "Comma separated list of detector types to exclude. Types must exist in detectors.proto").String()
 
 	gitScan             = cli.Command("git", "Find credentials in git repositories.")
 	gitScanURI          = gitScan.Arg("uri", "Git repository URL. https://, file://, or ssh:// schema expected.").Required().String()
@@ -211,12 +213,31 @@ func run(state overseer.State) {
 	}
 
 	ctx := context.TODO()
+	var detectorsOption engine.EngineOption
+
+	detectorCfg, err := engine.NewDetectorsConfig(*includeDetectors, *excludeDetectors)
+	if err != nil {
+		logger.Error(err, "failed to create detectors config")
+		os.Exit(1)
+	}
+	dts, err := engine.Detectors(ctx, detectorCfg)
+	if err != nil {
+		logger.Error(err, "failed to create detectors")
+		os.Exit(1)
+	}
+
+	if len(dts) > 0 {
+		detectorsOption = engine.WithDetectors(!*noVerification, dts...)
+	} else {
+		detectorsOption = engine.WithDetectors(!*noVerification, engine.DefaultDetectors()...)
+	}
+
 	e := engine.Start(ctx,
 		engine.WithConcurrency(*concurrency),
 		engine.WithDecoders(decoders.DefaultDecoders()...),
-		engine.WithDetectors(!*noVerification, engine.DefaultDetectors()...),
 		engine.WithDetectors(!*noVerification, conf.Detectors...),
 		engine.WithFilterUnverified(*filterUnverified),
+		detectorsOption,
 	)
 
 	var repoPath string

pkg/engine/detectors.go

@@ -0,0 +1,103 @@
+package engine
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
+)
+
+// DetectorFilter allows for include/exclude detectors.
+type DetectorFilter interface {
+	detectors() []string
+	filter(map[string]struct{}) ([]detectors.Detector, error)
+}
+
+// IncludeDetectorFilter is a detector filter that only includes the specified detectors.
+type IncludeDetectorFilter struct {
+	includeDetectors []string
+}
+
+func (i *IncludeDetectorFilter) detectors() []string {
+	return i.includeDetectors
+}
+
+func (i *IncludeDetectorFilter) filter(include map[string]struct{}) ([]detectors.Detector, error) {
+	ds := make([]detectors.Detector, 0, len(i.includeDetectors))
+	for _, d := range DefaultDetectors() {
+		if _, ok := include[strings.ToLower(d.Type().String())]; ok {
+			ds = append(ds, d)
+		}
+	}
+
+	if len(ds) != len(include) {
+		return nil, fmt.Errorf("1 or more detectors are invalid, please check your detector types")
+	}
+
+	return ds, nil
+}
+
+// ExcludeDetectorFilter is a detector filter that excludes the specified detectors.
+type ExcludeDetectorFilter struct {
+	excludeDetectors []string
+}
+
+func (e *ExcludeDetectorFilter) detectors() []string {
+	return e.excludeDetectors
+}
+
+func (e *ExcludeDetectorFilter) filter(exclude map[string]struct{}) ([]detectors.Detector, error) {
+	defaultDetectors := DefaultDetectors()
+	l := len(defaultDetectors) - len(exclude)
+	result := make([]detectors.Detector, 0, l)
+	for _, detector := range defaultDetectors {
+		if _, ok := exclude[strings.ToLower(detector.Type().String())]; !ok {
+			result = append(result, detector)
+		}
+	}
+
+	if len(result) != l {
+		return nil, fmt.Errorf("1 or more detectors are invalid, please check your detector types")
+	}
+	return result, nil
+}
+
+// NewDetectorsConfig creates a new detector filter based on the include/exclude arguments.
+func NewDetectorsConfig(include, exclude string) (DetectorFilter, error) {
+	if len(include) == 0 && len(exclude) == 0 {
+		return nil, fmt.Errorf("no detectors specified")
+	}
+
+	if len(include) > 0 && len(exclude) > 0 {
+		return nil, fmt.Errorf("cannot specify both include and exclude detectors")
+	}
+
+	// Determine the detector filter type.
+	if len(include) > len(exclude) {
+		return &IncludeDetectorFilter{includeDetectors: strings.Split(include, ",")}, nil
+	}
+	return &ExcludeDetectorFilter{excludeDetectors: strings.Split(exclude, ",")}, nil
+}
+
+// Detectors only returns a specific set of detectors if they are specified in the
+// detector list and are valid. Otherwise, an error is returned.
+func Detectors(ctx context.Context, dt DetectorFilter) ([]detectors.Detector, error) {
+	configured := setDetectors(ctx, dt.detectors())
+
+	if len(configured) == 0 {
+		return nil, fmt.Errorf("no detectors specified")
+	}
+
+	return dt.filter(configured)
+}
+
+func setDetectors(ctx context.Context, dts []string) map[string]struct{} {
+	valid := make(map[string]struct{}, len(dts))
+	for _, d := range dts {
+		ctx.Logger().Info("setting detector", "detector-name", d)
+		valid[strings.ToLower(d)] = struct{}{}
+	}
+
+	return valid
+}

pkg/engine/detectors_test.go

@@ -0,0 +1,126 @@
+package engine
+
+import (
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/aws"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/azure"
+)
+
+func TestNewDectorsConfig(t *testing.T) {
+	tests := []struct {
+		name     string
+		inc      string
+		exc      string
+		expected DetectorFilter
+		wantErr  bool
+	}{
+		{
+			name:    "No dts specified, returns error",
+			inc:     "",
+			exc:     "",
+			wantErr: true,
+		},
+		{
+			name:     "Include dts specified, returns config with include",
+			inc:      "AWS",
+			expected: &IncludeDetectorFilter{includeDetectors: []string{"AWS"}},
+		},
+		{
+			name:     "Exclude dts specified, returns config with exclude",
+			exc:      "AWS",
+			expected: &ExcludeDetectorFilter{excludeDetectors: []string{"AWS"}},
+		},
+		{
+			name:    "Include and exclude dts specified, returns error",
+			inc:     "AWS",
+			exc:     "Azure",
+			wantErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		result, err := NewDetectorsConfig(test.inc, test.exc)
+		if err != nil && !test.wantErr {
+			t.Errorf("NewDetectorsConfig(%v, %v) got %v, want %v", test.inc, test.exc, err, test.expected)
+		}
+
+		if result != nil {
+			if !reflect.DeepEqual(result, test.expected) {
+				t.Errorf("NewDetectorsConfig(%v, %v) got %v, want %v", test.inc, test.exc, result, test.expected)
+			}
+		}
+	}
+}
+
+func TestDetectors(t *testing.T) {
+	ctx := context.Background()
+	tests := []struct {
+		name     string
+		dt       DetectorFilter
+		expected []detectors.Detector
+		wantErr  bool
+	}{
+		{
+			name:    "No dts specified, returns default set",
+			dt:      &IncludeDetectorFilter{},
+			wantErr: true,
+		},
+		{
+			name:     "Valid include detectors (1) specified, returns valid set",
+			dt:       &IncludeDetectorFilter{includeDetectors: []string{"AWS"}},
+			expected: []detectors.Detector{aws.New()},
+		},
+		{
+			name:     "Valid include detectors (2) specified, returns valid set",
+			dt:       &IncludeDetectorFilter{includeDetectors: []string{"AWS", "Azure"}},
+			expected: []detectors.Detector{aws.New(), &azure.Scanner{}},
+		},
+		{
+			name:    "Invalid include detectors specified, returns error",
+			dt:      &IncludeDetectorFilter{includeDetectors: []string{"AWS", "Azure", "Invalid"}},
+			wantErr: true,
+		},
+		{
+			name:     "Valid exlude detectors (1) specified, returns valid set",
+			dt:       &ExcludeDetectorFilter{excludeDetectors: []string{"AWS"}},
+			expected: excludeDetectors(map[string]struct{}{"aws": {}}),
+		},
+		{
+			name:     "Valid exclude detectors (2) specified, returns valid set",
+			dt:       &ExcludeDetectorFilter{excludeDetectors: []string{"AWS", "Azure"}},
+			expected: excludeDetectors(map[string]struct{}{"aws": {}, "azure": {}}),
+		},
+		{
+			name:    "Invalid exclude detector specified, returns error",
+			dt:      &ExcludeDetectorFilter{excludeDetectors: []string{"AWS", "Azure", "Invalid"}},
+			wantErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		result, err := Detectors(ctx, test.dt)
+		if err != nil && !test.wantErr {
+			t.Errorf("Detectors(%v) got %v, want %v", test.dt, err, test.expected)
+		}
+
+		if !reflect.DeepEqual(result, test.expected) {
+			t.Errorf("Detectors(%v) got %v, want %v", test.dt, result, test.expected)
+		}
+	}
+}
+
+func excludeDetectors(exclude map[string]struct{}) []detectors.Detector {
+	defaultDetectors := DefaultDetectors()
+	result := make([]detectors.Detector, 0, len(defaultDetectors)-len(exclude))
+	for _, detector := range defaultDetectors {
+		if _, ok := exclude[strings.ToLower(detector.Type().String())]; !ok {
+			result = append(result, detector)
+		}
+	}
+	return result
+}