501 lines
14 KiB
Go
501 lines
14 KiB
Go
package filesystem
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/go-logr/logr"
|
|
"github.com/kylelemons/godebug/pretty"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/protobuf/types/known/anypb"
|
|
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/sourcestest"
|
|
)
|
|
|
|
func TestSource_Scan(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
|
|
defer cancel()
|
|
|
|
type init struct {
|
|
name string
|
|
verify bool
|
|
connection *sourcespb.Filesystem
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
init init
|
|
wantSourceMetadata *source_metadatapb.MetaData
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "get a chunk",
|
|
init: init{
|
|
name: "this repo",
|
|
connection: &sourcespb.Filesystem{
|
|
Paths: []string{"."},
|
|
},
|
|
verify: true,
|
|
},
|
|
wantSourceMetadata: &source_metadatapb.MetaData{
|
|
Data: &source_metadatapb.MetaData_Filesystem{
|
|
Filesystem: &source_metadatapb.Filesystem{
|
|
File: "filesystem.go",
|
|
},
|
|
},
|
|
},
|
|
wantErr: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
s := Source{}
|
|
|
|
conn, err := anypb.New(tt.init.connection)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
err = s.Init(ctx, tt.init.name, 0, 0, tt.init.verify, conn, 5)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("Source.Init() error = %v, wantErr %v", err, tt.wantErr)
|
|
return
|
|
}
|
|
chunksCh := make(chan *sources.Chunk, 1)
|
|
// TODO: this is kind of bad, if it errors right away we don't see it as a test failure.
|
|
// Debugging this usually requires setting a breakpoint on L78 and running test w/ debug.
|
|
go func() {
|
|
defer close(chunksCh)
|
|
err = s.Chunks(ctx, chunksCh)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("Source.Chunks() error = %v, wantErr %v", err, tt.wantErr)
|
|
return
|
|
}
|
|
}()
|
|
var counter int
|
|
for chunk := range chunksCh {
|
|
if chunk.SourceMetadata.GetFilesystem().GetFile() == "filesystem.go" {
|
|
counter++
|
|
if diff := pretty.Compare(chunk.SourceMetadata, tt.wantSourceMetadata); diff != "" {
|
|
t.Errorf("Source.Chunks() %s diff: (-got +want)\n%s", tt.name, diff)
|
|
}
|
|
}
|
|
}
|
|
assert.Equal(t, 1, counter)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestScanFile(t *testing.T) {
|
|
chunkSize := sources.DefaultChunkSize
|
|
secretPart1 := "SECRET"
|
|
secretPart2 := "SPLIT"
|
|
// Split the secret into two parts and pad the rest of the chunk with A's.
|
|
data := strings.Repeat("A", chunkSize-len(secretPart1)) + secretPart1 + secretPart2 + strings.Repeat("A", chunkSize-len(secretPart2))
|
|
|
|
tmpfile, cleanup, err := createTempFile("", data)
|
|
assert.Nil(t, err)
|
|
defer cleanup()
|
|
|
|
source := &Source{}
|
|
chunksChan := make(chan *sources.Chunk, 2)
|
|
|
|
ctx := context.WithLogger(context.Background(), logr.Discard())
|
|
go func() {
|
|
defer close(chunksChan)
|
|
err = source.scanFile(ctx, tmpfile.Name(), chunksChan)
|
|
assert.Nil(t, err)
|
|
}()
|
|
|
|
// Read from the channel and validate the secrets.
|
|
foundSecret := ""
|
|
for chunkCh := range chunksChan {
|
|
foundSecret += string(chunkCh.Data)
|
|
}
|
|
|
|
assert.Contains(t, foundSecret, secretPart1+secretPart2)
|
|
}
|
|
|
|
func TestScanBinaryFile(t *testing.T) {
|
|
tmpfile, err := os.CreateTemp("", "example.bin")
|
|
require.NoError(t, err)
|
|
defer os.Remove(tmpfile.Name())
|
|
|
|
// binary data that decodes to "TuffleHog"
|
|
fileContents := []byte{0x54, 0x75, 0x66, 0x66, 0x6C, 0x65, 0x48, 0x6F, 0x67}
|
|
_, err = tmpfile.Write(fileContents)
|
|
require.NoError(t, err)
|
|
require.NoError(t, tmpfile.Close())
|
|
|
|
source := &Source{}
|
|
chunksChan := make(chan *sources.Chunk, 2)
|
|
errChan := make(chan error, 1)
|
|
|
|
ctx := context.WithLogger(context.Background(), logr.Discard())
|
|
|
|
go func() {
|
|
defer close(chunksChan)
|
|
errChan <- source.scanFile(ctx, tmpfile.Name(), chunksChan)
|
|
}()
|
|
|
|
err = <-errChan
|
|
require.NoError(t, err)
|
|
|
|
var data string
|
|
for chunk := range chunksChan {
|
|
require.NotNil(t, chunk)
|
|
data += string(chunk.Data)
|
|
}
|
|
|
|
assert.Contains(t, data, "TuffleHog")
|
|
}
|
|
|
|
func TestEnumerate(t *testing.T) {
|
|
// TODO: refactor to allow a virtual filesystem.
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// Setup the connection to test enumeration.
|
|
dir, err := os.MkdirTemp("", "trufflehog-test-enumerate")
|
|
assert.NoError(t, err)
|
|
defer os.RemoveAll(dir)
|
|
|
|
units := []string{
|
|
"/one", "/two", "/three",
|
|
"/path/to/dir/", "/path/to/another/dir/",
|
|
}
|
|
// Prefix the units with the tempdir and create files on disk.
|
|
for i, unit := range units {
|
|
fullPath := filepath.Join(dir, unit)
|
|
units[i] = fullPath
|
|
if i < 3 {
|
|
f, err := os.Create(fullPath)
|
|
assert.NoError(t, err)
|
|
f.Close()
|
|
} else {
|
|
assert.NoError(t, os.MkdirAll(fullPath, 0755))
|
|
// Create a file in the directory for enumeration to find.
|
|
f, err := os.CreateTemp(fullPath, "file")
|
|
assert.NoError(t, err)
|
|
units[i] = f.Name()
|
|
f.Close()
|
|
}
|
|
}
|
|
conn, err := anypb.New(&sourcespb.Filesystem{
|
|
Paths: units[0:3],
|
|
Directories: units[3:],
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "test enumerate", 0, 0, true, conn, 1)
|
|
assert.NoError(t, err)
|
|
|
|
reporter := sourcestest.TestReporter{}
|
|
err = s.Enumerate(ctx, &reporter)
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, len(units), len(reporter.Units))
|
|
assert.Equal(t, 0, len(reporter.UnitErrs))
|
|
for _, unit := range reporter.Units {
|
|
path, _ := unit.SourceUnitID()
|
|
assert.Contains(t, units, path)
|
|
}
|
|
for _, unit := range units {
|
|
assert.Contains(t, reporter.Units, sources.CommonSourceUnit{ID: unit})
|
|
}
|
|
}
|
|
|
|
func TestChunkUnit(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// Setup test file to chunk.
|
|
fileContents := "TestChunkUnit"
|
|
tmpfile, cleanup, err := createTempFile("", fileContents)
|
|
assert.NoError(t, err)
|
|
defer cleanup()
|
|
|
|
tmpdir, cleanup, err := createTempDir("", "foo", "bar", "baz")
|
|
assert.NoError(t, err)
|
|
defer cleanup()
|
|
|
|
conn, err := anypb.New(&sourcespb.Filesystem{})
|
|
assert.NoError(t, err)
|
|
|
|
// Initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "test chunk unit", 0, 0, true, conn, 1)
|
|
assert.NoError(t, err)
|
|
|
|
// Happy path single file.
|
|
reporter := sourcestest.TestReporter{}
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: tmpfile.Name(),
|
|
}, &reporter)
|
|
assert.NoError(t, err)
|
|
|
|
// Happy path directory.
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: tmpdir,
|
|
}, &reporter)
|
|
assert.NoError(t, err)
|
|
|
|
// Error path.
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: "/file/not/found",
|
|
}, &reporter)
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, 4, len(reporter.Chunks))
|
|
assert.Equal(t, 1, len(reporter.ChunkErrs))
|
|
dataFound := make(map[string]struct{}, 4)
|
|
for _, chunk := range reporter.Chunks {
|
|
dataFound[string(chunk.Data)] = struct{}{}
|
|
}
|
|
assert.Contains(t, dataFound, fileContents)
|
|
assert.Contains(t, dataFound, "foo")
|
|
assert.Contains(t, dataFound, "bar")
|
|
assert.Contains(t, dataFound, "baz")
|
|
}
|
|
|
|
func TestEnumerateReporterErr(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// Setup the connection to test enumeration.
|
|
units := []string{
|
|
"/one", "/two", "/three",
|
|
"/path/to/dir/", "/path/to/another/dir/",
|
|
}
|
|
conn, err := anypb.New(&sourcespb.Filesystem{
|
|
Paths: units[0:3],
|
|
Directories: units[3:],
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "test enumerate", 0, 0, true, conn, 1)
|
|
assert.NoError(t, err)
|
|
|
|
// Enumerate should always return an error if the reporter returns an
|
|
// error.
|
|
reporter := sourcestest.ErrReporter{}
|
|
err = s.Enumerate(ctx, &reporter)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestChunkUnitReporterErr(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// Setup test file to chunk.
|
|
tmpfile, err := os.CreateTemp("", "example.txt")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer os.Remove(tmpfile.Name())
|
|
|
|
fileContents := []byte("TestChunkUnit")
|
|
_, err = tmpfile.Write(fileContents)
|
|
assert.NoError(t, err)
|
|
assert.NoError(t, tmpfile.Close())
|
|
|
|
conn, err := anypb.New(&sourcespb.Filesystem{})
|
|
assert.NoError(t, err)
|
|
|
|
// Initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "test chunk unit", 0, 0, true, conn, 1)
|
|
assert.NoError(t, err)
|
|
|
|
// Happy path. ChunkUnit should always return an error if the reporter
|
|
// returns an error.
|
|
reporter := sourcestest.ErrReporter{}
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: tmpfile.Name(),
|
|
}, &reporter)
|
|
assert.Error(t, err)
|
|
|
|
// Error path. ChunkUnit should always return an error if the reporter
|
|
// returns an error.
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: "/file/not/found",
|
|
}, &reporter)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestSkipDir(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// create a temp directory with files
|
|
ignoreDir, cleanupDir, err := createTempDir("", "ignore1", "ignore2", "ignore3")
|
|
require.NoError(t, err)
|
|
defer cleanupDir()
|
|
|
|
// create an ExcludePathsFile that contains the ignoreDir path
|
|
excludeFile, cleanupFile, err := createTempFile("", ignoreDir+"\n")
|
|
require.NoError(t, err)
|
|
defer cleanupFile()
|
|
|
|
conn, err := anypb.New(&sourcespb.Filesystem{
|
|
ExcludePathsFile: excludeFile.Name(),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "exclude directory", 0, 0, true, conn, 1)
|
|
require.NoError(t, err)
|
|
|
|
reporter := sourcestest.TestReporter{}
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: ignoreDir,
|
|
}, &reporter)
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, 0, len(reporter.Chunks), "Expected no chunks from excluded directory")
|
|
require.Equal(t, 0, len(reporter.ChunkErrs), "Expected no errors for excluded directory")
|
|
}
|
|
|
|
func TestScanSubDirFile(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
// create a temp directory with files
|
|
parentDir, cleanupParentDir, err := createTempDir("", "file1")
|
|
require.NoError(t, err)
|
|
defer cleanupParentDir()
|
|
|
|
childDir, cleanupChildDir, err := createTempDir(parentDir, "file2")
|
|
require.NoError(t, err)
|
|
defer cleanupChildDir()
|
|
|
|
// create a file in child directory
|
|
file, cleanupFile, err := createTempFile(childDir, "should scan this file")
|
|
require.NoError(t, err)
|
|
defer cleanupFile()
|
|
|
|
// create an IncludePathsFile that contains the file path
|
|
includeFile, cleanupFile, err := createTempFile("", file.Name()+"\n")
|
|
require.NoError(t, err)
|
|
defer cleanupFile()
|
|
|
|
conn, err := anypb.New(&sourcespb.Filesystem{
|
|
IncludePathsFile: includeFile.Name(),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// initialize the source.
|
|
s := Source{}
|
|
err = s.Init(ctx, "include sub directory file", 0, 0, true, conn, 1)
|
|
require.NoError(t, err)
|
|
|
|
reporter := sourcestest.TestReporter{}
|
|
err = s.ChunkUnit(ctx, sources.CommonSourceUnit{
|
|
ID: parentDir,
|
|
}, &reporter)
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, 1, len(reporter.Chunks), "Expected chunks from included file")
|
|
require.Equal(t, 0, len(reporter.ChunkErrs), "Expected no errors")
|
|
}
|
|
|
|
func TestSkipBinaries(t *testing.T) {
|
|
// Create a temporary directory for testing
|
|
tempDir, err := os.MkdirTemp("", "trufflehog_test")
|
|
require.NoError(t, err)
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
// Create a binary file (executable)
|
|
binaryFile := filepath.Join(tempDir, "test.exe")
|
|
err = os.WriteFile(binaryFile, []byte{0x4D, 0x5A, 0x90, 0x00}, 0644)
|
|
require.NoError(t, err)
|
|
|
|
// Create a text file
|
|
textFile := filepath.Join(tempDir, "test.txt")
|
|
err = os.WriteFile(textFile, []byte("This is a text file"), 0644)
|
|
require.NoError(t, err)
|
|
|
|
// Test with skipBinaries = true
|
|
source := &Source{
|
|
paths: []string{textFile, binaryFile}, // Test individual files
|
|
skipBinaries: true,
|
|
log: logr.Discard(),
|
|
}
|
|
|
|
chunks := make(chan *sources.Chunk, 10)
|
|
ctx := context.Background()
|
|
|
|
// Run the scan
|
|
go func() {
|
|
err := source.Chunks(ctx, chunks)
|
|
require.NoError(t, err)
|
|
close(chunks)
|
|
}()
|
|
|
|
// Collect chunks
|
|
var chunkCount int
|
|
var processedFiles []string
|
|
for chunk := range chunks {
|
|
chunkCount++
|
|
metadata := chunk.SourceMetadata.GetFilesystem()
|
|
require.NotNil(t, metadata)
|
|
processedFiles = append(processedFiles, metadata.File)
|
|
}
|
|
|
|
// Should have exactly one chunk from the text file
|
|
require.Equal(t, 1, chunkCount, "Should have processed exactly one text file")
|
|
require.Contains(t, processedFiles, textFile, "Should have processed the text file")
|
|
require.NotContains(t, processedFiles, binaryFile, "Binary file should be skipped")
|
|
}
|
|
|
|
// createTempFile is a helper function to create a temporary file in the given
|
|
// directory with the provided contents. If dir is "", the operating system's
|
|
// temp directory is used.
|
|
func createTempFile(dir string, contents string) (*os.File, func(), error) {
|
|
tmpfile, err := os.CreateTemp(dir, "trufflehogtest")
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
if _, err := tmpfile.Write([]byte(contents)); err != nil {
|
|
_ = os.Remove(tmpfile.Name())
|
|
return nil, nil, err
|
|
}
|
|
if err := tmpfile.Close(); err != nil {
|
|
_ = os.Remove(tmpfile.Name())
|
|
return nil, nil, err
|
|
}
|
|
return tmpfile, func() { _ = os.Remove(tmpfile.Name()) }, nil
|
|
}
|
|
|
|
// createTempDir is a helper function to create a temporary directory in the
|
|
// given directory with files containing the provided contents. If dir is "",
|
|
// the operating system's temp directory is used.
|
|
func createTempDir(dir string, contents ...string) (string, func(), error) {
|
|
tmpdir, err := os.MkdirTemp(dir, "trufflehogtest")
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
for _, content := range contents {
|
|
if _, _, err := createTempFile(tmpdir, content); err != nil {
|
|
_ = os.RemoveAll(tmpdir)
|
|
return "", nil, err
|
|
}
|
|
}
|
|
return tmpdir, func() { _ = os.RemoveAll(tmpdir) }, nil
|
|
}
|