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 <tarun@projectdiscovery.io>
This commit is contained in:
Douglas Danger Manley 2024-06-22 15:42:01 -04:00 committed by GitHub
parent 1c51a6bef6
commit e61ca0c42d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 137 additions and 49 deletions

View File

@ -231,7 +231,10 @@ func (e *NucleiEngine) init(ctx context.Context) error {
// and also upgrade templates to latest version if available // and also upgrade templates to latest version if available
installer.NucleiSDKVersionCheck() installer.NucleiSDKVersionCheck()
return e.processUpdateCheckResults() if DefaultConfig.CanCheckForUpdates() {
return e.processUpdateCheckResults()
}
return nil
} }
type syncOnce struct { type syncOnce struct {

View File

@ -11,17 +11,25 @@ import (
// DiskCatalog is a template catalog helper implementation based on disk // DiskCatalog is a template catalog helper implementation based on disk
type DiskCatalog struct { type DiskCatalog struct {
templatesDirectory string 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 // NewCatalog creates a new Catalog structure using provided input items
// using disk based items // using disk based items
func NewCatalog(directory string) *DiskCatalog { func NewCatalog(directory string) *DiskCatalog {
catalog := &DiskCatalog{templatesDirectory: directory} catalog := &DiskCatalog{templatesDirectory: directory}
if directory != "" { if directory == "" {
catalog.templatesFS = os.DirFS(directory) catalog.templatesDirectory = config.DefaultConfig.GetTemplateDir()
} else { }
catalog.templatesFS = os.DirFS(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 return catalog
} }
@ -29,11 +37,15 @@ func NewCatalog(directory string) *DiskCatalog {
// OpenFile opens a file and returns an io.ReadCloser to the file. // 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. // It is used to read template and payload files based on catalog responses.
func (d *DiskCatalog) OpenFile(filename string) (io.ReadCloser, error) { func (d *DiskCatalog) OpenFile(filename string) (io.ReadCloser, error) {
file, err := os.Open(filename) if d.templatesFS == nil {
if err != nil { file, err := os.Open(filename)
if file, errx := os.Open(BackwardsCompatiblePaths(d.templatesDirectory, filename)); errx == nil { if err != nil {
return file, 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)
} }

View File

@ -79,17 +79,21 @@ func (c *DiskCatalog) GetTemplatePath(target string) ([]string, error) {
} }
// try to handle deprecated template paths // try to handle deprecated template paths
absPath := BackwardsCompatiblePaths(c.templatesDirectory, target) absPath := target
if absPath != target && strings.TrimPrefix(absPath, c.templatesDirectory+string(filepath.Separator)) != target { if c.templatesFS == nil {
if config.DefaultConfig.LogAllEvents { absPath = BackwardsCompatiblePaths(c.templatesDirectory, target)
gologger.DefaultLogger.Print().Msgf("[%v] requested Template path %s is deprecated, please update to %s\n", aurora.Yellow("WRN").String(), target, absPath) 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) var err error
if err != nil { absPath, err = c.convertPathToAbsolute(absPath)
return nil, errors.Wrapf(err, "could not find template file") if err != nil {
return nil, errors.Wrapf(err, "could not find template file")
}
} }
// Template input is either a file or a directory // 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 == "" { if c.templatesDirectory == "" {
templateDir = "./" 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 return matches
} }
// condition to support network cve related globs
matches, _ = fs.Glob(os.DirFS(filepath.Join(templateDir, "network")), inputGlob)
return matches
} }
var matched []string var matched []string
matches, err := fs.Glob(c.templatesFS, relPath) var matches []string
if len(matches) != 0 { if c.templatesFS == nil {
matched = append(matched, matches...) 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 { } else {
matched = append(matched, OldPathsResolver(relPath)...) var err error
} matches, err = fs.Glob(c.templatesFS, relPath)
if err != nil && len(matched) == 0 { if len(matches) != 0 {
return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err) 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)) results := make([]string, 0, len(matches))
for _, match := range 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 // findFileMatches finds if a path is an absolute file. If the path
// is a file, it returns true otherwise false with no errors. // 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) { 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 { if err != nil {
return "", false, err return "", false, err
} }
if !info.Mode().IsRegular() { stat, err := info.Stat()
if err != nil {
return "", false, err
}
if !stat.Mode().IsRegular() {
return "", false, nil return "", false, nil
} }
if _, ok := processed[absPath]; !ok { 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 // findDirectoryMatches finds matches for templates from a directory
func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]struct{}) ([]string, error) { func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]struct{}) ([]string, error) {
var results []string var results []string
err := filepath.WalkDir( var err error
absPath, if c.templatesFS == nil {
func(path string, d fs.DirEntry, err error) error { err = filepath.WalkDir(
// continue on errors absPath,
if err != nil { func(path string, d fs.DirEntry, err error) error {
return nil // continue on errors
} if err != nil {
if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown { return nil
if _, ok := processed[path]; !ok {
results = append(results, path)
processed[path] = struct{}{}
} }
} if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown {
return nil if _, ok := processed[path]; !ok {
}, results = append(results, path)
) processed[path] = struct{}{}
}
}
return nil
},
)
} else {
err = fs.WalkDir(
c.templatesFS,
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
},
)
}
return results, err return results, err
} }