Skip to content

Commit

Permalink
Fix include and exclude detector logic (#1267)
Browse files Browse the repository at this point in the history
* Fix include and exclude detector logic

* Fix test

* Add more clarifying comments
  • Loading branch information
mcastorina committed Apr 26, 2023
1 parent 622700b commit 5a86c18
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 34 deletions.
85 changes: 58 additions & 27 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/handlers"
"github.com/trufflesecurity/trufflehog/v3/pkg/log"
"github.com/trufflesecurity/trufflehog/v3/pkg/output"
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
"github.com/trufflesecurity/trufflehog/v3/pkg/sources/git"
"github.com/trufflesecurity/trufflehog/v3/pkg/updater"
Expand Down Expand Up @@ -233,8 +232,8 @@ func run(state overseer.State) {
handlers.SetArchiveMaxTimeout(*archiveTimeout)
}

// Build include and exclude detector filter sets.
var includeDetectorTypes, excludeDetectorTypes map[detectorspb.DetectorType]config.DetectorID
// Build include and exclude detector sets for filtering on engine initialization.
var includeDetectorSet, excludeDetectorSet map[config.DetectorID]struct{}
{
includeList, err := config.ParseDetectors(*includeDetectors)
if err != nil {
Expand All @@ -246,40 +245,55 @@ func run(state overseer.State) {
// Exit if there was an error to inform the user of the misconfiguration.
logFatal(err, "invalid exclude list detector configuration")
}
includeDetectorTypes = detectorTypeToMap(includeList)
excludeDetectorTypes = detectorTypeToMap(excludeList)
includeDetectorSet = detectorTypeToSet(includeList)
excludeDetectorSet = detectorTypeToSet(excludeList)
}
includeFilter := func(d detectors.Detector) bool {
id, ok := includeDetectorTypes[d.Type()]
if id.Version == 0 {
return ok
}
versionD, ok := d.(detectors.Versioner)
if !ok {
// Error: version provided but not a detectors.Versioner

// Verify that all the user-provided detectors support the optional
// detector features.
{
isVersioner := engine.DefaultDetectorTypesImplementing[detectors.Versioner]()
for id := range includeDetectorSet {
if id.Version == 0 {
// Version not provided.
continue
}
if _, ok := isVersioner[id.ID]; ok {
// Version provided for a Versioner detector.
continue
}
// Version provided on a non-Versioner detector.
logFatal(
fmt.Errorf("version provided but detector does not have a version"),
"invalid include list detector configuration",
"detector", id,
)
}
return versionD.Version() == id.Version
}
excludeFilter := func(d detectors.Detector) bool {
id, ok := excludeDetectorTypes[d.Type()]
if id.Version == 0 {
return !ok
}
versionD, ok := d.(detectors.Versioner)
if !ok {
// Error: version provided but not a detectors.Versioner
for id := range excludeDetectorSet {
if id.Version == 0 {
// Version not provided.
continue
}
if _, ok := isVersioner[id.ID]; ok {
// Version provided for a Versioner detector.
continue
}
// Version provided on a non-Versioner detector.
logFatal(
fmt.Errorf("version provided but detector does not have a version"),
"invalid exclude list detector configuration",
"detector", id,
)
}
return versionD.Version() != id.Version
}

includeFilter := func(d detectors.Detector) bool {
_, ok := getWithDetectorID(d, includeDetectorSet)
return ok
}
excludeFilter := func(d detectors.Detector) bool {
_, ok := getWithDetectorID(d, excludeDetectorSet)
return !ok
}

e := engine.Start(ctx,
Expand Down Expand Up @@ -523,10 +537,27 @@ func logFatalFunc(logger logr.Logger) func(error, string, ...any) {
}
}

func detectorTypeToMap(detectors []config.DetectorID) map[detectorspb.DetectorType]config.DetectorID {
output := make(map[detectorspb.DetectorType]config.DetectorID, len(detectors))
// detectorTypeToSet is a helper function to convert a slice of detector IDs into a set.
func detectorTypeToSet(detectors []config.DetectorID) map[config.DetectorID]struct{} {
output := make(map[config.DetectorID]struct{}, len(detectors))
for _, d := range detectors {
output[d.ID] = d
output[d] = struct{}{}
}
return output
}

// getWithDetectorID is a helper function to get a value from a map using a
// detector's ID. This function behaves like a normal map lookup, with an extra
// step of checking for the non-specific version of a detector.
func getWithDetectorID[T any](d detectors.Detector, data map[config.DetectorID]T) (T, bool) {
key := config.GetDetectorID(d)
// Check if the specific ID is provided.
if t, ok := data[key]; ok || key.Version == 0 {
return t, ok
}
// Check if the generic type is provided without a version.
// This means "all" versions of a type.
key.Version = 0
t, ok := data[key]
return t, ok
}
4 changes: 4 additions & 0 deletions pkg/config/detectors.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,15 @@ func init() {
}
}

// DetectorID identifies a detector type and version. This struct is used as a
// way for users to identify detectors, whether unique or not. A DetectorID
// with Version = 0 indicates all possible versions of a detector.
type DetectorID struct {
ID dpb.DetectorType
Version int
}

// GetDetectorID extracts the DetectorID from a Detector.
func GetDetectorID(d detectors.Detector) DetectorID {
var version int
if v, ok := d.(detectors.Versioner); ok {
Expand Down
24 changes: 17 additions & 7 deletions pkg/engine/defaults.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/airship"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/airtableapikey"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/airvisual"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/alchemy"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/alconost"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/alegra"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/aletheiaapi"
Expand Down Expand Up @@ -75,6 +76,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blablabus"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blazemeter"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blitapp"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blocknative"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blogger"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/bombbomb"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/boostnote"
Expand All @@ -83,6 +85,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/brandfetch"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/browserstack"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/browshot"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/bscscan"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/buddyns"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/bugherd"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/bugsnag"
Expand Down Expand Up @@ -142,6 +145,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/coinbase"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/coinlayer"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/coinlib"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/coinmarketcap"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/collect2"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/column"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/commercejs"
Expand Down Expand Up @@ -213,6 +217,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/elasticemail"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/enablex"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/enigma"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/etherscan"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/ethplorer"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/etsyapikey"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/everhour"
Expand Down Expand Up @@ -319,6 +324,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/imagekit"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/imagga"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/impala"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/infura"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/insightly"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/instabot"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/integromat"
Expand Down Expand Up @@ -408,6 +414,7 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/monkeylearn"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/moonclerk"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/moosend"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/moralis"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/mrticktock"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/mux"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/myfreshworks"
Expand Down Expand Up @@ -734,13 +741,6 @@ import (
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/zonkafeedback"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/zulipchat"
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/etherscan"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/infura"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/alchemy"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/blocknative"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/moralis"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/bscscan"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/coinmarketcap"
)

// CustomDetectors returns a list of detectors that are enabled by default, but
Expand Down Expand Up @@ -1570,3 +1570,13 @@ func DefaultDetectors() []detectors.Detector {
}

}

func DefaultDetectorTypesImplementing[T any]() map[detectorspb.DetectorType]struct{} {
out := make(map[detectorspb.DetectorType]struct{})
for _, detector := range DefaultDetectors() {
if _, ok := detector.(T); ok {
out[detector.Type()] = struct{}{}
}
}
return out
}
50 changes: 50 additions & 0 deletions pkg/engine/defaults_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
package engine

import (
"testing"

"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
)

func TestDefaultDetectorsHaveUniqueVersions(t *testing.T) {
detectorTypeToVersions := make(map[detectorspb.DetectorType]map[int]struct{})
addVersion := func(versions map[int]struct{}, version int) map[int]struct{} {
if versions == nil {
versions = make(map[int]struct{})
}
versions[version] = struct{}{}
return versions
}
// Loop through all our default detectors and find the ones that
// implement Versioner. Of those, check each version number is unique.
for _, detector := range DefaultDetectors() {
v, ok := detector.(detectors.Versioner)
if !ok {
continue
}
version := v.Version()
key := detector.Type()
if set, ok := detectorTypeToVersions[key]; ok && set != nil {
if _, ok := set[version]; ok {
t.Errorf("detector %q has duplicate version: %d", detectorspb.DetectorType_name[int32(key)], version)
}
}
detectorTypeToVersions[key] = addVersion(detectorTypeToVersions[key], version)
}
}

func TestDefaultDetectorTypesImplementing(t *testing.T) {
isVersioner := DefaultDetectorTypesImplementing[detectors.Versioner]()
for _, detector := range DefaultDetectors() {
_, expectedOk := detector.(detectors.Versioner)
_, gotOk := isVersioner[detector.Type()]
if expectedOk == gotOk {
continue
}
t.Errorf(
"detector %q doesn't match expected",
detectorspb.DetectorType_name[int32(detector.Type())],
)
}
}

0 comments on commit 5a86c18

Please sign in to comment.