Compare commits

...

2 Commits

Author SHA1 Message Date
Miccah Castorina c0752eae2c Update test 2023-07-26 00:36:42 -05:00
Miccah Castorina c6597969eb Add support for storing handler metadata
Currently it's a unstructured map[string]string that gets copied into
a Result.ExtraData after detection, but we might want to revisit this
decision when we get more handlers.
2023-07-26 00:24:11 -05:00
6 changed files with 43 additions and 26 deletions
+8 -1
View File
@@ -349,8 +349,15 @@ func (e *Engine) detectorWorker(ctx context.Context) {
continue continue
} }
result.DecoderType = decoderType result.DecoderType = decoderType
if result.ExtraData == nil && chunk.HandleMetadata != nil {
result.ExtraData = chunk.HandleMetadata
} else {
for k, v := range chunk.HandleMetadata {
// TODO: Check key collisions.
result.ExtraData[k] = v
}
}
chunkResults = append(chunkResults, detectors.CopyMetadata(resultChunk, result)) chunkResults = append(chunkResults, detectors.CopyMetadata(resultChunk, result))
} }
if len(results) > 0 { if len(results) > 0 {
elapsed := time.Since(start) elapsed := time.Since(start)
+20 -15
View File
@@ -6,12 +6,14 @@ import (
"errors" "errors"
"fmt" "fmt"
"io" "io"
"strings"
"time" "time"
"github.com/mholt/archiver/v4" "github.com/mholt/archiver/v4"
"github.com/trufflesecurity/trufflehog/v3/pkg/common" "github.com/trufflesecurity/trufflehog/v3/pkg/common"
logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context" logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
) )
type ctxKey int type ctxKey int
@@ -52,14 +54,14 @@ func SetArchiveMaxTimeout(timeout time.Duration) {
} }
// FromFile extracts the files from an archive. // FromFile extracts the files from an archive.
func (d *Archive) FromFile(originalCtx context.Context, data io.Reader) chan ([]byte) { func (d *Archive) FromFile(originalCtx context.Context, data io.Reader) chan ChunkOpt {
archiveChan := make(chan ([]byte), 512) archiveChan := make(chan ChunkOpt, 512)
go func() { go func() {
ctx, cancel := context.WithTimeout(originalCtx, maxTimeout) ctx, cancel := context.WithTimeout(originalCtx, maxTimeout)
logger := logContext.AddLogger(ctx).Logger() logger := logContext.AddLogger(ctx).Logger()
defer cancel() defer cancel()
defer close(archiveChan) defer close(archiveChan)
err := d.openArchive(ctx, 0, data, archiveChan) err := d.openArchive(ctx, nil, data, archiveChan)
if err != nil { if err != nil {
if errors.Is(err, archiver.ErrNoMatch) { if errors.Is(err, archiver.ErrNoMatch) {
return return
@@ -71,18 +73,21 @@ func (d *Archive) FromFile(originalCtx context.Context, data io.Reader) chan ([]
} }
// openArchive takes a reader and extracts the contents up to the maximum depth. // openArchive takes a reader and extracts the contents up to the maximum depth.
func (d *Archive) openArchive(ctx context.Context, depth int, reader io.Reader, archiveChan chan ([]byte)) error { func (d *Archive) openArchive(ctx context.Context, depth []string, reader io.Reader, archiveChan chan ChunkOpt) error {
if depth >= maxDepth { if len(depth) >= maxDepth {
return fmt.Errorf("max archive depth reached") return fmt.Errorf("max archive depth reached")
} }
format, reader, err := archiver.Identify("", reader) format, reader, err := archiver.Identify("", reader)
if err != nil { if err != nil {
if errors.Is(err, archiver.ErrNoMatch) && depth > 0 { if errors.Is(err, archiver.ErrNoMatch) && len(depth) > 0 {
chunkSize := 10 * 1024 chunkSize := 10 * 1024
for { for {
chunk := make([]byte, chunkSize) data := make([]byte, chunkSize)
n, _ := reader.Read(chunk) n, _ := reader.Read(data)
archiveChan <- chunk archiveChan <- func(chunk *sources.Chunk) {
chunk.Data = data
chunk.HandleMetadata = map[string]string{"archive path": strings.Join(depth, " > ")}
}
if n < chunkSize { if n < chunkSize {
break break
} }
@@ -102,9 +107,9 @@ func (d *Archive) openArchive(ctx context.Context, depth int, reader io.Reader,
return err return err
} }
newReader := bytes.NewReader(fileBytes) newReader := bytes.NewReader(fileBytes)
return d.openArchive(ctx, depth+1, newReader, archiveChan) return d.openArchive(ctx, append(depth, "decompressed"), newReader, archiveChan)
case archiver.Extractor: case archiver.Extractor:
err := archive.Extract(context.WithValue(ctx, depthKey, depth+1), reader, nil, d.extractorHandler(archiveChan)) err := archive.Extract(context.WithValue(ctx, depthKey, depth), reader, nil, d.extractorHandler(archiveChan))
if err != nil { if err != nil {
return err return err
} }
@@ -129,12 +134,12 @@ func (d *Archive) IsFiletype(ctx context.Context, reader io.Reader) (io.Reader,
} }
// extractorHandler is applied to each file in an archiver.Extractor file. // extractorHandler is applied to each file in an archiver.Extractor file.
func (d *Archive) extractorHandler(archiveChan chan ([]byte)) func(context.Context, archiver.File) error { func (d *Archive) extractorHandler(archiveChan chan ChunkOpt) func(context.Context, archiver.File) error {
return func(ctx context.Context, f archiver.File) error { return func(ctx context.Context, f archiver.File) error {
logger := logContext.AddLogger(ctx).Logger() logger := logContext.AddLogger(ctx).Logger()
logger.V(5).Info("Handling extracted file.", "filename", f.Name()) logger.V(5).Info("Handling extracted file.", "filename", f.Name())
depth := 0 var depth []string
if ctxDepth, ok := ctx.Value(depthKey).(int); ok { if ctxDepth, ok := ctx.Value(depthKey).([]string); ok {
depth = ctxDepth depth = ctxDepth
} }
@@ -148,7 +153,7 @@ func (d *Archive) extractorHandler(archiveChan chan ([]byte)) func(context.Conte
} }
fileContent := bytes.NewReader(fileBytes) fileContent := bytes.NewReader(fileBytes)
err = d.openArchive(ctx, depth, fileContent, archiveChan) err = d.openArchive(ctx, append(depth, f.Name()), fileContent, archiveChan)
if err != nil { if err != nil {
return err return err
} }
+4 -2
View File
@@ -85,9 +85,11 @@ func TestArchiveHandler(t *testing.T) {
count := 0 count := 0
re := regexp.MustCompile(testCase.matchString) re := regexp.MustCompile(testCase.matchString)
matched := false matched := false
for chunk := range archiveChan { var chunk sources.Chunk
for opt := range archiveChan {
count++ count++
if re.Match(chunk) { opt(&chunk)
if re.Match(chunk.Data) {
matched = true matched = true
} }
} }
+5 -3
View File
@@ -13,8 +13,10 @@ func DefaultHandlers() []Handler {
} }
} }
type ChunkOpt func(*sources.Chunk)
type Handler interface { type Handler interface {
FromFile(context.Context, io.Reader) chan ([]byte) FromFile(context.Context, io.Reader) chan ChunkOpt
IsFiletype(context.Context, io.Reader) (io.Reader, bool) IsFiletype(context.Context, io.Reader) (io.Reader, bool)
New() New()
} }
@@ -38,13 +40,13 @@ func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, c
handlerChan := handler.FromFile(ctx, file) handlerChan := handler.FromFile(ctx, file)
for { for {
select { select {
case data, open := <-handlerChan: case opt, open := <-handlerChan:
if !open { if !open {
// We finished reading everything from handlerChan. // We finished reading everything from handlerChan.
return true return true
} }
chunk := *chunkSkel chunk := *chunkSkel
chunk.Data = data opt(&chunk)
// Send data on chunksChan. // Send data on chunksChan.
select { select {
case chunksChan <- &chunk: case chunksChan <- &chunk:
+1 -1
View File
@@ -96,7 +96,7 @@ func structToMap(obj interface{}) (m map[string]map[string]interface{}, err erro
} }
type outputFormat struct { type outputFormat struct {
DetectorType, DetectorType string
DecoderType string DecoderType string
Verified bool Verified bool
Raw string Raw string
+2 -1
View File
@@ -21,7 +21,8 @@ type Chunk struct {
SourceType sourcespb.SourceType SourceType sourcespb.SourceType
// SourceMetadata holds the context of where the Chunk was found. // SourceMetadata holds the context of where the Chunk was found.
SourceMetadata *source_metadatapb.MetaData SourceMetadata *source_metadatapb.MetaData
// HandleMetadata holds the metadata from a handler if one was used.
HandleMetadata map[string]string
// Data is the data to decode and scan. // Data is the data to decode and scan.
Data []byte Data []byte
// Verify specifies whether any secrets in the Chunk should be verified. // Verify specifies whether any secrets in the Chunk should be verified.