✨ Add depth to file searching. (#4816)

Signed-off-by: Preslav <[email protected]>
mondoohq · Nov 5, 2024 · 3536493 · 3536493
1 parent 4a2b16b
commit 3536493
Show file tree

Hide file tree

Showing 6 changed files with 104 additions and 24 deletions.
diff --git a/providers/os/connection/shared/shared.go b/providers/os/connection/shared/shared.go
@@ -116,7 +116,7 @@ func (c Capabilities) String() []string {
 }
 
 type FileSearch interface {
-	Find(from string, r *regexp.Regexp, typ string, perm *uint32) ([]string, error)
+	Find(from string, r *regexp.Regexp, typ string, perm *uint32, depth *int) ([]string, error)
 }
 
 type PerfStats struct {

diff --git a/providers/os/fs/find_files.go b/providers/os/fs/find_files.go
@@ -5,20 +5,25 @@ package fs
 
 import (
 	"io/fs"
+	"os"
 	"regexp"
 	"strings"
 )
 
-func FindFiles(iofs fs.FS, from string, r *regexp.Regexp, typ string, perm *uint32) ([]string, error) {
-	matcher := createFindFilesMatcher(iofs, typ, r, perm)
+func FindFiles(iofs fs.FS, from string, r *regexp.Regexp, typ string, perm *uint32, depth *int) ([]string, error) {
+	matcher := createFindFilesMatcher(iofs, typ, from, r, perm, depth)
 	matchedPaths := []string{}
 	err := fs.WalkDir(iofs, from, func(p string, d fs.DirEntry, err error) error {
-		skip, err := handleFsError(err)
+		if d.IsDir() && matcher.DepthReached(p) {
+			return fs.SkipDir
+		}
+
+		skipFile, err := handleFsError(err)
 		if err != nil {
 			return err
 		}
 
-		if skip {
+		if skipFile {
 			return nil
 		}
 		if matcher.Match(p, d.Type()) {
@@ -36,9 +41,26 @@ type findFilesMatcher struct {
 	types []byte
 	r     *regexp.Regexp
 	perm  *uint32
+	depth *int
+	from  string
 	iofs  fs.FS
 }
 
+// Depth 0 means we only walk the current directory
+// Depth 1 means we walk the current directory and its children
+// Depth 2 means we walk the current directory, its children and their children
+func (m findFilesMatcher) DepthReached(p string) bool {
+	if m.depth == nil {
+		return false
+	}
+
+	trimmed := strings.TrimPrefix(p, m.from)
+	// WalkDir always uses slash for separating, ignoring the OS separator. This is why we need to replace it.
+	normalized := strings.ReplaceAll(trimmed, string(os.PathSeparator), "/")
+	depth := strings.Count(normalized, "/")
+	return depth > *m.depth
+}
+
 func (m findFilesMatcher) Match(path string, t fs.FileMode) bool {
 	matchesType := m.matchesType(t)
 	matchesRegex := m.matchesRegex(path)
@@ -101,7 +123,7 @@ func (m findFilesMatcher) matchesPerm(path string) bool {
 	return true
 }
 
-func createFindFilesMatcher(iofs fs.FS, typeStr string, r *regexp.Regexp, perm *uint32) findFilesMatcher {
+func createFindFilesMatcher(iofs fs.FS, typeStr string, from string, r *regexp.Regexp, perm *uint32, depth *int) findFilesMatcher {
 	allowed := []byte{}
 	types := strings.Split(typeStr, ",")
 	for _, t := range types {
@@ -120,5 +142,7 @@ func createFindFilesMatcher(iofs fs.FS, typeStr string, r *regexp.Regexp, perm *
 		r:     r,
 		perm:  perm,
 		iofs:  iofs,
+		depth: depth,
+		from:  from,
 	}
 }
diff --git a/providers/os/fs/find_files_test.go b/providers/os/fs/find_files_test.go
@@ -66,8 +66,8 @@ func TestFindFilesMatcher(t *testing.T) {
 			}
 			fs := afero.IOFS{Fs: afero.NewMemMapFs()}
 			t.Run(fmt.Sprintf("%s matcher", string(tc.matches)), func(t *testing.T) {
-				exclusionMatcher := createFindFilesMatcher(fs, strings.Join(excludeTypes, ","), nil, nil)
-				exactMatcher := createFindFilesMatcher(fs, string(tc.matches), nil, nil)
+				exclusionMatcher := createFindFilesMatcher(fs, strings.Join(excludeTypes, ","), "", nil, nil, nil)
+				exactMatcher := createFindFilesMatcher(fs, string(tc.matches), "", nil, nil, nil)
 				assert.True(t, exactMatcher.Match("/foo", tc.typ), "exact matcher failed to match")
 				assert.False(t, exclusionMatcher.Match("/foo", tc.typ), "exclusion matcher matched")
 			})
@@ -77,7 +77,7 @@ func TestFindFilesMatcher(t *testing.T) {
 	t.Run("regex", func(t *testing.T) {
 		t.Run("any type", func(t *testing.T) {
 			fs := afero.IOFS{Fs: afero.NewMemMapFs()}
-			exactMatcher := createFindFilesMatcher(fs, "", regexp.MustCompile("foo.*"), nil)
+			exactMatcher := createFindFilesMatcher(fs, "", "", regexp.MustCompile("foo.*"), nil, nil)
 
 			for _, m := range possibleModes {
 				t.Run(fmt.Sprintf("mode %s", m.String()), func(t *testing.T) {
@@ -89,7 +89,7 @@ func TestFindFilesMatcher(t *testing.T) {
 		})
 
 		t.Run("specific type", func(t *testing.T) {
-			exactMatcher := createFindFilesMatcher(afero.IOFS{Fs: afero.NewMemMapFs()}, "f", regexp.MustCompile("foo.*"), nil)
+			exactMatcher := createFindFilesMatcher(afero.IOFS{Fs: afero.NewMemMapFs()}, "f", "", regexp.MustCompile("foo.*"), nil, nil)
 
 			assert.False(t, exactMatcher.Match("foobar", fs.ModeDir))
 			assert.True(t, exactMatcher.Match("foobar", fs.ModePerm))
@@ -135,44 +135,98 @@ func TestFindFilesMatcher(t *testing.T) {
 				}
 			}
 			t.Run(fmt.Sprintf("%s matcher", string(tc.matches)), func(t *testing.T) {
-				exactMatcher := createFindFilesMatcher(fs, "", nil, nil)
+				exactMatcher := createFindFilesMatcher(fs, "", "", nil, nil, nil)
 				assert.True(t, exactMatcher.Match("/foo", tc.typ), "matcher failed to match")
 			})
 		}
 	})
 }
 
+func TestDepthMatcher(t *testing.T) {
+	t.Run("depth 0", func(t *testing.T) {
+		fs := afero.IOFS{Fs: afero.NewMemMapFs()}
+		depth := 0
+		depthMatcher := createFindFilesMatcher(fs, "", "root", nil, nil, &depth)
+		assert.True(t, depthMatcher.DepthReached("root/foo"))
+		assert.True(t, depthMatcher.DepthReached("root/foo/bar"))
+	})
+
+	t.Run("depth 1", func(t *testing.T) {
+		fs := afero.IOFS{Fs: afero.NewMemMapFs()}
+		depth := 1
+		depthMatcher := createFindFilesMatcher(fs, "", "root/foo", nil, nil, &depth)
+		assert.False(t, depthMatcher.DepthReached("root/foo"))
+		assert.False(t, depthMatcher.DepthReached("root/foo/bar"))
+		assert.True(t, depthMatcher.DepthReached("root/foo/bar/baz"))
+	})
+}
+
 func TestFindFiles(t *testing.T) {
 	fs := afero.NewMemMapFs()
 	mkDir(t, fs, "root/a")
 	mkDir(t, fs, "root/b")
 	mkDir(t, fs, "root/c")
+	mkDir(t, fs, "root/c/d")
+
+	mkFile(t, fs, "root/file0")
 	mkFile(t, fs, "root/a/file1")
 	mkFile(t, fs, "root/a/file2")
 	mkFile(t, fs, "root/b/file1")
 	mkFile(t, fs, "root/c/file4")
+	mkFile(t, fs, "root/c/d/file5")
 	require.NoError(t, fs.Chmod("root/c/file4", 0o002))
 
-	rootAFiles, err := FindFiles(afero.NewIOFS(fs), "root/a", nil, "f", nil)
+	rootAFiles, err := FindFiles(afero.NewIOFS(fs), "root/a", nil, "f", nil, nil)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, rootAFiles, []string{"root/a/file1", "root/a/file2"})
 
-	rootAFilesAndDir, err := FindFiles(afero.NewIOFS(fs), "root/a", nil, "f,d", nil)
+	rootAFilesAndDir, err := FindFiles(afero.NewIOFS(fs), "root/a", nil, "f,d", nil, nil)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, rootAFilesAndDir, []string{"root/a", "root/a/file1", "root/a/file2"})
 
-	rootBFiles, err := FindFiles(afero.NewIOFS(fs), "root", regexp.MustCompile("root/b.*"), "f", nil)
+	rootBFiles, err := FindFiles(afero.NewIOFS(fs), "root", regexp.MustCompile("root/b.*"), "f", nil, nil)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, rootBFiles, []string{"root/b/file1"})
 
-	file1Files, err := FindFiles(afero.NewIOFS(fs), "root", regexp.MustCompile(".*/file1"), "f", nil)
+	file1Files, err := FindFiles(afero.NewIOFS(fs), "root", regexp.MustCompile(".*/file1"), "f", nil, nil)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, file1Files, []string{"root/b/file1", "root/a/file1"})
 
 	perm := uint32(0o002)
-	permFiles, err := FindFiles(afero.NewIOFS(fs), "root", nil, "f", &perm)
+	permFiles, err := FindFiles(afero.NewIOFS(fs), "root", nil, "f", &perm, nil)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, permFiles, []string{"root/c/file4"})
+
+	depth := 0
+	depthFiles, err := FindFiles(afero.NewIOFS(fs), "root", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/file0"})
+
+	depth = 1
+	depthFiles, err = FindFiles(afero.NewIOFS(fs), "root", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/file0", "root/a/file1", "root/a/file2", "root/b/file1", "root/c/file4"})
+
+	depth = 2
+	depthFiles, err = FindFiles(afero.NewIOFS(fs), "root", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/file0", "root/a/file1", "root/a/file2", "root/b/file1", "root/c/file4", "root/c/d/file5"})
+
+	// relative roots
+	depth = 0
+	depthFiles, err = FindFiles(afero.NewIOFS(fs), "root/a", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/a/file1", "root/a/file2"})
+
+	depth = 1
+	depthFiles, err = FindFiles(afero.NewIOFS(fs), "root/c", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/c/file4", "root/c/d/file5"})
+
+	depth = 0
+	depthFiles, err = FindFiles(afero.NewIOFS(fs), "root/c/d", nil, "f", nil, &depth)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, depthFiles, []string{"root/c/d/file5"})
 }
 
 func mkFile(t *testing.T, fs afero.Fs, name string) {

diff --git a/providers/os/fs/fs.go b/providers/os/fs/fs.go
@@ -113,7 +113,7 @@ func (t *MountedFs) Chown(name string, uid, gid int) error {
 	return notSupported
 }
 
-func (t *MountedFs) Find(from string, r *regexp.Regexp, typ string, perm *uint32) ([]string, error) {
+func (t *MountedFs) Find(from string, r *regexp.Regexp, typ string, perm *uint32, depth *int) ([]string, error) {
 	iofs := afero.NewIOFS(t)
-	return FindFiles(iofs, from, r, typ, perm)
+	return FindFiles(iofs, from, r, typ, perm, depth)
 }
diff --git a/providers/os/resources/files.go b/providers/os/resources/files.go
@@ -87,7 +87,7 @@ func (l *mqlFilesFind) list() ([]interface{}, error) {
 			perm = &p
 		}
 
-		foundFiles, err = fsSearch.Find(l.From.Data, compiledRegexp, l.Type.Data, perm)
+		foundFiles, err = fsSearch.Find(l.From.Data, compiledRegexp, l.Type.Data, perm, nil)
 		if err != nil {
 			return nil, err
 		}

diff --git a/providers/os/resources/packages/windows_packages.go b/providers/os/resources/packages/windows_packages.go
@@ -8,6 +8,7 @@ import (
 	"encoding/xml"
 	"fmt"
 	"io"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"time"
@@ -363,13 +364,14 @@ func (w *WinPkgManager) getFsAppxPackages() ([]Package, error) {
 		return nil, errors.New("find file is not supported for your platform")
 	}
 
-	paths := []string{
-		`Windows\SystemApps`,
-		`Program Files\WindowsApps`,
+	paths := map[string]int{
+		filepath.Join("Windows", "SystemApps"):        1,
+		filepath.Join("Program Files", "WindowsApps"): 1,
+		"Windows": 1,
 	}
 	appxPaths := map[string]struct{}{}
-	for _, p := range paths {
-		res, err := fsSearch.Find(p, regexp.MustCompile(".*/[Aa]ppx[Mm]anifest.xml"), "f", nil)
+	for p, depth := range paths {
+		res, err := fsSearch.Find(p, regexp.MustCompile(".*/[Aa]ppx[Mm]anifest.xml"), "f", nil, &depth)
 		if err != nil {
 			continue
 		}