Use fs.FS when explicitly given (#5312)

* Use the `templateFS` if it's there when calling `OpenFile` * Add a new constructor * More refactoring * Both of my use cases are working * Fix for legacy assumptions * minor update: remove gologger debug stmts --------- Co-authored-by: Tarun Koyalwar <[email protected]>
projectdiscovery · Jun 22, 2024 · e61ca0c · e61ca0c
1 parent 1c51a6b
commit e61ca0c
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 48 deletions.
diff --git a/lib/sdk_private.go b/lib/sdk_private.go
@@ -231,7 +231,10 @@ func (e *NucleiEngine) init(ctx context.Context) error {
 	// and also upgrade templates to latest version if available
 	installer.NucleiSDKVersionCheck()
 
-	return e.processUpdateCheckResults()
+	if DefaultConfig.CanCheckForUpdates() {
+		return e.processUpdateCheckResults()
+	}
+	return nil
 }
 
 type syncOnce struct {

diff --git a/pkg/catalog/disk/catalog.go b/pkg/catalog/disk/catalog.go
@@ -11,29 +11,41 @@ import (
 // DiskCatalog is a template catalog helper implementation based on disk
 type DiskCatalog struct {
 	templatesDirectory string
-	templatesFS        fs.FS // TODO: Refactor to use this
+	templatesFS        fs.FS // Due to issues with how Go has implemented fs.FS, we'll have to also implement normal os operations, as well.  See: https://github.com/golang/go/issues/44279
 }
 
 // NewCatalog creates a new Catalog structure using provided input items
 // using disk based items
 func NewCatalog(directory string) *DiskCatalog {
 	catalog := &DiskCatalog{templatesDirectory: directory}
-	if directory != "" {
-		catalog.templatesFS = os.DirFS(directory)
-	} else {
-		catalog.templatesFS = os.DirFS(config.DefaultConfig.GetTemplateDir())
+	if directory == "" {
+		catalog.templatesDirectory = config.DefaultConfig.GetTemplateDir()
+	}
+	return catalog
+}
+
+// NewFSCatalog creates a new Catalog structure using provided input items
+// using the fs.FS as its filesystem.
+func NewFSCatalog(fs fs.FS, directory string) *DiskCatalog {
+	catalog := &DiskCatalog{
+		templatesDirectory: directory,
+		templatesFS:        fs,
 	}
 	return catalog
 }
 
 // OpenFile opens a file and returns an io.ReadCloser to the file.
 // It is used to read template and payload files based on catalog responses.
 func (d *DiskCatalog) OpenFile(filename string) (io.ReadCloser, error) {
-	file, err := os.Open(filename)
-	if err != nil {
-		if file, errx := os.Open(BackwardsCompatiblePaths(d.templatesDirectory, filename)); errx == nil {
-			return file, nil
+	if d.templatesFS == nil {
+		file, err := os.Open(filename)
+		if err != nil {
+			if file, errx := os.Open(BackwardsCompatiblePaths(d.templatesDirectory, filename)); errx == nil {
+				return file, nil
+			}
 		}
+		return file, err
 	}
-	return file, err
+
+	return d.templatesFS.Open(filename)
 }
diff --git a/pkg/catalog/disk/find.go b/pkg/catalog/disk/find.go
@@ -79,17 +79,21 @@ func (c *DiskCatalog) GetTemplatePath(target string) ([]string, error) {
 	}
 
 	// try to handle deprecated template paths
-	absPath := BackwardsCompatiblePaths(c.templatesDirectory, target)
-	if absPath != target && strings.TrimPrefix(absPath, c.templatesDirectory+string(filepath.Separator)) != target {
-		if config.DefaultConfig.LogAllEvents {
-			gologger.DefaultLogger.Print().Msgf("[%v] requested Template path %s is deprecated, please update to %s\n", aurora.Yellow("WRN").String(), target, absPath)
+	absPath := target
+	if c.templatesFS == nil {
+		absPath = BackwardsCompatiblePaths(c.templatesDirectory, target)
+		if absPath != target && strings.TrimPrefix(absPath, c.templatesDirectory+string(filepath.Separator)) != target {
+			if config.DefaultConfig.LogAllEvents {
+				gologger.DefaultLogger.Print().Msgf("[%v] requested Template path %s is deprecated, please update to %s\n", aurora.Yellow("WRN").String(), target, absPath)
+			}
+			deprecatedPathsCounter++
 		}
-		deprecatedPathsCounter++
-	}
 
-	absPath, err := c.convertPathToAbsolute(absPath)
-	if err != nil {
-		return nil, errors.Wrapf(err, "could not find template file")
+		var err error
+		absPath, err = c.convertPathToAbsolute(absPath)
+		if err != nil {
+			return nil, errors.Wrapf(err, "could not find template file")
+		}
 	}
 
 	// Template input is either a file or a directory
@@ -143,24 +147,60 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 		if c.templatesDirectory == "" {
 			templateDir = "./"
 		}
-		matches, _ := fs.Glob(os.DirFS(filepath.Join(templateDir, "http")), inputGlob)
-		if len(matches) != 0 {
+
+		if c.templatesFS == nil {
+			matches, _ := fs.Glob(os.DirFS(filepath.Join(templateDir, "http")), inputGlob)
+			if len(matches) != 0 {
+				return matches
+			}
+
+			// condition to support network cve related globs
+			matches, _ = fs.Glob(os.DirFS(filepath.Join(templateDir, "network")), inputGlob)
+			return matches
+		} else {
+			sub, err := fs.Sub(c.templatesFS, filepath.Join(templateDir, "http"))
+			if err != nil {
+				return nil
+			}
+			matches, _ := fs.Glob(sub, inputGlob)
+			if len(matches) != 0 {
+				return matches
+			}
+
+			// condition to support network cve related globs
+			sub, err = fs.Sub(c.templatesFS, filepath.Join(templateDir, "network"))
+			if err != nil {
+				return nil
+			}
+			matches, _ = fs.Glob(sub, inputGlob)
 			return matches
 		}
-		// condition to support network cve related globs
-		matches, _ = fs.Glob(os.DirFS(filepath.Join(templateDir, "network")), inputGlob)
-		return matches
 	}
 
 	var matched []string
-	matches, err := fs.Glob(c.templatesFS, relPath)
-	if len(matches) != 0 {
-		matched = append(matched, matches...)
+	var matches []string
+	if c.templatesFS == nil {
+		var err error
+		matches, err = filepath.Glob(relPath)
+		if len(matches) != 0 {
+			matched = append(matched, matches...)
+		} else {
+			matched = append(matched, OldPathsResolver(relPath)...)
+		}
+		if err != nil && len(matched) == 0 {
+			return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err)
+		}
 	} else {
-		matched = append(matched, OldPathsResolver(relPath)...)
-	}
-	if err != nil && len(matched) == 0 {
-		return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err)
+		var err error
+		matches, err = fs.Glob(c.templatesFS, relPath)
+		if len(matches) != 0 {
+			matched = append(matched, matches...)
+		} else {
+			matched = append(matched, OldPathsResolver(relPath)...)
+		}
+		if err != nil && len(matched) == 0 {
+			return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err)
+		}
 	}
 	results := make([]string, 0, len(matches))
 	for _, match := range matches {
@@ -175,11 +215,23 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 // findFileMatches finds if a path is an absolute file. If the path
 // is a file, it returns true otherwise false with no errors.
 func (c *DiskCatalog) findFileMatches(absPath string, processed map[string]struct{}) (match string, matched bool, err error) {
-	info, err := os.Stat(absPath)
+	if c.templatesFS != nil {
+		absPath = strings.TrimPrefix(absPath, "/")
+	}
+	var info fs.File
+	if c.templatesFS == nil {
+		info, err = os.Open(absPath)
+	} else {
+		info, err = c.templatesFS.Open(absPath)
+	}
+	if err != nil {
+		return "", false, err
+	}
+	stat, err := info.Stat()
 	if err != nil {
 		return "", false, err
 	}
-	if !info.Mode().IsRegular() {
+	if !stat.Mode().IsRegular() {
 		return "", false, nil
 	}
 	if _, ok := processed[absPath]; !ok {
@@ -192,22 +244,43 @@ func (c *DiskCatalog) findFileMatches(absPath string, processed map[string]struc
 // findDirectoryMatches finds matches for templates from a directory
 func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]struct{}) ([]string, error) {
 	var results []string
-	err := filepath.WalkDir(
-		absPath,
-		func(path string, d fs.DirEntry, err error) error {
-			// continue on errors
-			if err != nil {
+	var err error
+	if c.templatesFS == nil {
+		err = filepath.WalkDir(
+			absPath,
+			func(path string, d fs.DirEntry, err error) error {
+				// continue on errors
+				if err != nil {
+					return nil
+				}
+				if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown {
+					if _, ok := processed[path]; !ok {
+						results = append(results, path)
+						processed[path] = struct{}{}
+					}
+				}
 				return nil
-			}
-			if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown {
-				if _, ok := processed[path]; !ok {
-					results = append(results, path)
-					processed[path] = struct{}{}
+			},
+		)
+	} else {
+		err = fs.WalkDir(
+			c.templatesFS,
+			absPath,
+			func(path string, d fs.DirEntry, err error) error {
+				// continue on errors
+				if err != nil {
+					return nil
 				}
-			}
-			return nil
-		},
-	)
+				if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown {
+					if _, ok := processed[path]; !ok {
+						results = append(results, path)
+						processed[path] = struct{}{}
+					}
+				}
+				return nil
+			},
+		)
+	}
 	return results, err
 }