From bdc1b337e0c79309039fbc13c561f8a3de424bd8 Mon Sep 17 00:00:00 2001 From: Sajad Parra Date: Wed, 2 Mar 2022 16:50:20 +0530 Subject: [PATCH] fix setting custom nuclei-templates directory #1611 --- v2/cmd/nuclei/main.go | 4 +--- v2/internal/runner/options.go | 2 +- v2/internal/runner/runner.go | 2 +- v2/internal/runner/update.go | 20 +++++++++++--------- v2/pkg/catalog/config/config.go | 18 ++++++++++++++---- v2/pkg/catalog/loader/loader.go | 5 +++-- v2/pkg/types/resume.go | 6 +++--- 7 files changed, 34 insertions(+), 23 deletions(-) diff --git a/v2/cmd/nuclei/main.go b/v2/cmd/nuclei/main.go index 9cd8a948..1c5eb062 100644 --- a/v2/cmd/nuclei/main.go +++ b/v2/cmd/nuclei/main.go @@ -70,8 +70,6 @@ func main() { } func readConfig() { - home, _ := os.UserHomeDir() - templatesDirectory := filepath.Join(home, "nuclei-templates") flagSet := goflags.NewFlagSet() flagSet.SetDescription(`Nuclei is a fast, template based vulnerability scanner focusing @@ -196,7 +194,7 @@ on extensive configurability, massive extensibility and ease of use.`) createGroup(flagSet, "update", "Update", flagSet.BoolVar(&options.UpdateNuclei, "update", false, "update nuclei engine to the latest released version"), flagSet.BoolVarP(&options.UpdateTemplates, "update-templates", "ut", false, "update nuclei-templates to latest released version"), - flagSet.StringVarP(&options.TemplatesDirectory, "update-directory", "ud", templatesDirectory, "overwrite the default directory to install nuclei-templates"), + flagSet.StringVarP(&options.TemplatesDirectory, "update-directory", "ud", "", "overwrite the default directory to install nuclei-templates"), flagSet.BoolVarP(&options.NoUpdateTemplates, "disable-update-check", "duc", false, "disable automatic nuclei/templates update check"), ) diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 22764c27..b2672036 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -41,7 +41,7 @@ func ParseOptions(options *types.Options) { // Show the user the banner showBanner() - if !filepath.IsAbs(options.TemplatesDirectory) { + if options.TemplatesDirectory != "" && !filepath.IsAbs(options.TemplatesDirectory) { cwd, _ := os.Getwd() options.TemplatesDirectory = filepath.Join(cwd, options.TemplatesDirectory) } diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 63c8490a..af2873b3 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -299,7 +299,7 @@ func (r *Runner) RunEnumeration() error { } executerOpts.WorkflowLoader = workflowLoader - store, err := loader.New(loader.NewConfig(r.options, r.catalog, executerOpts)) + store, err := loader.New(loader.NewConfig(r.options, r.templatesConfig, r.catalog, executerOpts)) if err != nil { return errors.Wrap(err, "could not load templates from config") } diff --git a/v2/internal/runner/update.go b/v2/internal/runner/update.go index 4966901e..78013506 100644 --- a/v2/internal/runner/update.go +++ b/v2/internal/runner/update.go @@ -50,19 +50,22 @@ var reVersion = regexp.MustCompile(`\d+\.\d+\.\d+`) // If the path exists but does not contain the latest version of public templates, // the new version is downloaded from GitHub to the templates' directory, overwriting the old content. func (r *Runner) updateTemplates() error { // TODO this method does more than just update templates. Should be refactored. - home, err := os.UserHomeDir() + configDir, err := config.GetConfigDir() if err != nil { return err } - configDir := filepath.Join(home, ".config", "nuclei") _ = os.MkdirAll(configDir, 0755) - if err := r.readInternalConfigurationFile(home, configDir); err != nil { + if err := r.readInternalConfigurationFile(configDir); err != nil { return errors.Wrap(err, "could not read configuration file") } // If the config doesn't exist, create it now. if r.templatesConfig == nil { + home, err := os.UserHomeDir() + if err != nil { + return err + } currentConfig := &config.Config{ TemplatesDirectory: filepath.Join(home, "nuclei-templates"), NucleiVersion: config.Version, @@ -72,6 +75,9 @@ func (r *Runner) updateTemplates() error { // TODO this method does more than ju } r.templatesConfig = currentConfig } + if r.options.TemplatesDirectory == "" { + r.options.TemplatesDirectory = r.templatesConfig.TemplatesDirectory + } if r.options.NoUpdateTemplates && !r.options.UpdateTemplates { return nil @@ -90,11 +96,7 @@ func (r *Runner) updateTemplates() error { // TODO this method does more than ju if r.templatesConfig.TemplateVersion == "" || (r.options.TemplatesDirectory != "" && r.templatesConfig.TemplatesDirectory != r.options.TemplatesDirectory) || noTemplatesFound { gologger.Info().Msgf("nuclei-templates are not installed, installing...\n") - // Use the custom location if the user has given a template directory - r.templatesConfig = &config.Config{ - TemplatesDirectory: filepath.Join(home, "nuclei-templates"), - } - if r.options.TemplatesDirectory != "" && r.options.TemplatesDirectory != filepath.Join(home, "nuclei-templates") { + if r.options.TemplatesDirectory != "" && r.templatesConfig.TemplatesDirectory != r.options.TemplatesDirectory { r.templatesConfig.TemplatesDirectory, _ = filepath.Abs(r.options.TemplatesDirectory) } r.fetchLatestVersionsFromGithub(configDir) // also fetch the latest versions @@ -192,7 +194,7 @@ func getVersions(runner *Runner) (semver.Version, semver.Version, error) { } // readInternalConfigurationFile reads the internal configuration file for nuclei -func (r *Runner) readInternalConfigurationFile(home, configDir string) error { +func (r *Runner) readInternalConfigurationFile(configDir string) error { templatesConfigFile := filepath.Join(configDir, nucleiConfigFilename) if _, statErr := os.Stat(templatesConfigFile); !os.IsNotExist(statErr) { configuration, readErr := config.ReadConfiguration() diff --git a/v2/pkg/catalog/config/config.go b/v2/pkg/catalog/config/config.go index caf5dd9e..4a6edec2 100644 --- a/v2/pkg/catalog/config/config.go +++ b/v2/pkg/catalog/config/config.go @@ -3,6 +3,7 @@ package config import ( "os" "path/filepath" + "strings" jsoniter "github.com/json-iterator/go" "github.com/pkg/errors" @@ -30,16 +31,26 @@ const nucleiConfigFilename = ".templates-config.json" const Version = `2.6.3-dev` func getConfigDetails() (string, error) { - homeDir, err := os.UserHomeDir() + configDir, err := GetConfigDir() if err != nil { return "", errors.Wrap(err, "could not get home directory") } - configDir := filepath.Join(homeDir, ".config", "nuclei") _ = os.MkdirAll(configDir, 0755) templatesConfigFile := filepath.Join(configDir, nucleiConfigFilename) return templatesConfigFile, nil } +// GetConfigDir returns the nuclei configuration directory +func GetConfigDir() (string, error) { + appName := filepath.Base(os.Args[0]) + appName = strings.TrimSuffix(appName, filepath.Ext(appName)) + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".config", appName), nil +} + // ReadConfiguration reads the nuclei configuration file from disk. func ReadConfiguration() (*Config, error) { templatesConfigFile, err := getConfigDetails() @@ -138,9 +149,8 @@ func getIgnoreFilePath() string { return defIgnoreFilePath } - home, err := os.UserHomeDir() + configDir, err := GetConfigDir() if err == nil { - configDir := filepath.Join(home, ".config", "nuclei") _ = os.MkdirAll(configDir, 0755) defIgnoreFilePath = filepath.Join(configDir, nucleiIgnoreFile) diff --git a/v2/pkg/catalog/loader/loader.go b/v2/pkg/catalog/loader/loader.go index 6323b832..0d9cf0a8 100644 --- a/v2/pkg/catalog/loader/loader.go +++ b/v2/pkg/catalog/loader/loader.go @@ -5,6 +5,7 @@ import ( "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/catalog" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" "github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader/filter" "github.com/projectdiscovery/nuclei/v2/pkg/model/types/severity" "github.com/projectdiscovery/nuclei/v2/pkg/parsers" @@ -57,7 +58,7 @@ type Store struct { } // NewConfig returns a new loader config -func NewConfig(options *types.Options, catalog *catalog.Catalog, executerOpts protocols.ExecuterOptions) *Config { +func NewConfig(options *types.Options, templateConfig *config.Config, catalog *catalog.Catalog, executerOpts protocols.ExecuterOptions) *Config { loaderConfig := Config{ Templates: options.Templates, Workflows: options.Workflows, @@ -74,7 +75,7 @@ func NewConfig(options *types.Options, catalog *catalog.Catalog, executerOpts pr IncludeTags: options.IncludeTags, IncludeIds: options.IncludeIds, ExcludeIds: options.ExcludeIds, - TemplatesDirectory: options.TemplatesDirectory, + TemplatesDirectory: templateConfig.TemplatesDirectory, Protocols: options.Protocols, ExcludeProtocols: options.ExcludeProtocols, Catalog: catalog, diff --git a/v2/pkg/types/resume.go b/v2/pkg/types/resume.go index 5308888e..06ce21e0 100644 --- a/v2/pkg/types/resume.go +++ b/v2/pkg/types/resume.go @@ -3,10 +3,10 @@ package types import ( "fmt" "math" - "os" "path/filepath" "sync" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" "github.com/rs/xid" ) @@ -14,11 +14,11 @@ import ( const DefaultResumeFileName = "resume-%s.cfg" func DefaultResumeFilePath() string { - home, err := os.UserHomeDir() + configDir, err := config.GetConfigDir() if err != nil { return fmt.Sprintf("resume-%s.cfg", xid.New().String()) } - resumeFile := filepath.Join(home, ".config", "nuclei", fmt.Sprintf("resume-%s.cfg", xid.New().String())) + resumeFile := filepath.Join(configDir, fmt.Sprintf("resume-%s.cfg", xid.New().String())) return resumeFile }