mirror of
https://github.com/trufflesecurity/trufflehog.git
synced 2024-11-10 15:14:38 +00:00
123 lines
3.2 KiB
Go
123 lines
3.2 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"regexp"
|
|
"strings"
|
|
"testing"
|
|
|
|
diskbufferreader "github.com/bill-rich/disk-buffer-reader"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
|
)
|
|
|
|
func TestArchiveHandler(t *testing.T) {
|
|
tests := map[string]struct {
|
|
archiveURL string
|
|
expectedChunks int
|
|
matchString string
|
|
}{
|
|
"gzip-single": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/one-zip.gz",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"gzip-nested": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/double-zip.gz",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"gzip-too-deep": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/six-zip.gz",
|
|
0,
|
|
"",
|
|
},
|
|
"tar-single": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/one.tar",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"tar-nested": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/two.tar",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"tar-too-deep": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/six.tar",
|
|
0,
|
|
"",
|
|
},
|
|
"targz-single": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/tar-archive.tar.gz",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"gzip-large": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/FifteenMB.gz",
|
|
1543,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
"zip-single": {
|
|
"https://raw.githubusercontent.com/bill-rich/bad-secrets/master/aws-canary-creds.zip",
|
|
1,
|
|
"AKIAYVP4CIPPH5TNP3SW",
|
|
},
|
|
}
|
|
|
|
for name, testCase := range tests {
|
|
resp, err := http.Get(testCase.archiveURL)
|
|
if err != nil || resp.StatusCode != http.StatusOK {
|
|
t.Error(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
archive := Archive{}
|
|
archive.New()
|
|
|
|
newReader, err := diskbufferreader.New(resp.Body)
|
|
if err != nil {
|
|
t.Errorf("error creating reusable reader: %s", err)
|
|
}
|
|
archiveChan := archive.FromFile(context.TODO(), newReader)
|
|
|
|
count := 0
|
|
re := regexp.MustCompile(testCase.matchString)
|
|
matched := false
|
|
for chunk := range archiveChan {
|
|
count++
|
|
if re.Match(chunk) {
|
|
matched = true
|
|
}
|
|
}
|
|
if !matched && len(testCase.matchString) > 0 {
|
|
t.Errorf("%s: Expected string not found in archive.", name)
|
|
}
|
|
if count != testCase.expectedChunks {
|
|
t.Errorf("%s: Unexpected number of chunks. Got %d, expected: %d", name, count, testCase.expectedChunks)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestHandleFile(t *testing.T) {
|
|
ch := make(chan *sources.Chunk, 2)
|
|
|
|
// Context cancels the operation.
|
|
canceledCtx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
assert.False(t, HandleFile(canceledCtx, strings.NewReader("file"), &sources.Chunk{}, ch))
|
|
|
|
// Only one chunk is sent on the channel.
|
|
// TODO: Embed a zip without making an HTTP request.
|
|
resp, err := http.Get("https://raw.githubusercontent.com/bill-rich/bad-secrets/master/aws-canary-creds.zip")
|
|
assert.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
archive := Archive{}
|
|
archive.New()
|
|
reader, err := diskbufferreader.New(resp.Body)
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, 0, len(ch))
|
|
assert.True(t, HandleFile(context.Background(), reader, &sources.Chunk{}, ch))
|
|
assert.Equal(t, 1, len(ch))
|
|
}
|