Skip to content

Commit

Permalink
Implemented only "tl" to list available templates
Browse files Browse the repository at this point in the history
  • Loading branch information
vzamanillo committed Aug 29, 2020
1 parent 099a76a commit 4484ba8
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 46 deletions.
4 changes: 1 addition & 3 deletions internal/runner/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@ type Options struct {
ProxySocksURL string // ProxySocksURL is the URL for the proxy socks server
CustomHeaders requests.CustomHeaders // Custom global headers
TemplatesDirectory string // TemplatesDirectory is the directory to use for storing templates
TemplateSearch string // Search templates
}

type multiStringFlag []string
Expand Down Expand Up @@ -79,7 +78,6 @@ func ParseOptions() *Options {
flag.BoolVar(&options.JSONRequests, "json-requests", false, "Write requests/responses for matches in JSON output")
flag.BoolVar(&options.EnableProgressBar, "pbar", false, "Enable the progress bar")
flag.BoolVar(&options.TemplateList, "tl", false, "List available templates")
flag.StringVar(&options.TemplateSearch, "ts", "", "Search templates")

flag.Parse()

Expand Down Expand Up @@ -127,7 +125,7 @@ func (options *Options) validateOptions() error {
return errors.New("both verbose and silent mode specified")
}

if !options.TemplateList && options.TemplateSearch == "" {
if !options.TemplateList {
// Check if a list of templates was provided and it exists
if len(options.Templates) == 0 && !options.UpdateTemplates {
return errors.New("no template/templates provided")
Expand Down
4 changes: 2 additions & 2 deletions internal/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ func New(options *Options) (*Runner, error) {
gologger.Warningf("Could not update templates: %s\n", err)
}

if options.TemplateList || options.TemplateSearch != "" {
runner.listAvailableTemplates(options.TemplateSearch)
if options.TemplateList {
runner.listAvailableTemplates()
os.Exit(0)
}

Expand Down
89 changes: 48 additions & 41 deletions internal/runner/templates.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,13 @@ import (
"github.com/projectdiscovery/nuclei/v2/pkg/workflows"
)

var severityMap = map[string]string{
"info": aurora.Cyan("info").String(),
"low": aurora.Green("low").String(),
"medium": aurora.Yellow("medium").String(),
"high": aurora.Red("high").String(),
}

// getTemplatesFor parses the specified input template definitions and returns a list of unique, absolute template paths.
func (r *Runner) getTemplatesFor(definitions []string) []string {
// keeps track of processed dirs and files
Expand Down Expand Up @@ -136,19 +143,17 @@ func (r *Runner) getParsedTemplatesFor(templatePaths []string, severities string
t, err := r.parseTemplateFile(match)
switch tp := t.(type) {
case *templates.Template:
id := tp.ID

// only include if severity matches or no severity filtering
sev := strings.ToLower(tp.Info.Severity)
if !filterBySeverity || hasMatchingSeverity(sev, allSeverities) {
parsedTemplates = append(parsedTemplates, tp)
r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
gologger.Infof("%s\n", r.templateLogMsg(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity))
} else {
gologger.Warningf("Excluding template %s due to severity filter (%s not in [%s])", id, sev, severities)
gologger.Warningf("Excluding template %s due to severity filter (%s not in [%s])", tp.ID, sev, severities)
}
case *workflows.Workflow:
parsedTemplates = append(parsedTemplates, tp)
r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
gologger.Infof("%s\n", r.templateLogMsg(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity))
workflowCount++
default:
gologger.Errorf("Could not parse file '%s': %s\n", match, err)
Expand Down Expand Up @@ -182,56 +187,58 @@ func (r *Runner) parseTemplateFile(file string) (interface{}, error) {
return nil, errors.New("unknown error occurred")
}

// LogTemplateLoaded logs a message for loaded template
func (r *Runner) logTemplateLoaded(id, name, author, severity string) {
func (r *Runner) templateLogMsg(id, name, author, severity string) string {
// Display the message for the template
message := fmt.Sprintf("[%s] %s (%s)",
r.colorizer.BrightBlue(id).String(),
r.colorizer.Bold(name).String(),
r.colorizer.BrightYellow("@"+author).String())

if severity != "" {
message += " [" + r.colorizer.Yellow(severity).String() + "]"
message += " [" + severityMap[severity] + "]"
}

gologger.Infof("%s\n", message)
return message
}

func (r *Runner) logAvailableTemplate(tplPath string) {
t, err := r.parseTemplateFile(tplPath)
if t != nil {
switch tp := t.(type) {
case *templates.Template:
gologger.Silentf("%s\n", r.templateLogMsg(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity))
case *workflows.Workflow:
gologger.Silentf("%s\n", r.templateLogMsg(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity))
default:
gologger.Errorf("Could not parse file '%s': %s\n", tplPath, err)
}
}
}

// ListAvailableTemplates prints available templates to stdout
func (r *Runner) listAvailableTemplates(criteria string) {
if criteria == "" {
gologger.Infof("Listing available templates...")
} else {
gologger.Infof("Searching available templates for %s", criteria)
func (r *Runner) listAvailableTemplates() {
if r.templatesConfig == nil {
return
}
if r.templatesConfig != nil {
r.colorizer = aurora.NewAurora(true)
err := directoryWalker(
r.templatesConfig.TemplatesDirectory,
func(path string, d *godirwalk.Dirent) error {
if d.IsDir() {
gologger.Silentf("%s\n", d.Name())
} else if strings.HasSuffix(path, ".yaml") {
t, err := r.parseTemplateFile(path)
if t != nil {
switch tp := t.(type) {
case *templates.Template:
r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
case *workflows.Workflow:
r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
default:
gologger.Errorf("Could not parse file '%s': %s\n", path, err)
}
}
}

return nil
},
)
gologger.Silentf("\nListing available v.%s nuclei templates for %s", r.templatesConfig.CurrentVersion, r.templatesConfig.TemplatesDirectory)
r.colorizer = aurora.NewAurora(true)
err := directoryWalker(
r.templatesConfig.TemplatesDirectory,
func(path string, d *godirwalk.Dirent) error {
if d.IsDir() && path != r.templatesConfig.TemplatesDirectory {
gologger.Silentf("\n%s:\n\n", r.colorizer.Bold(r.colorizer.BgBrightBlue(strings.Title(d.Name()))).String())
} else if strings.HasSuffix(path, ".yaml") {
r.logAvailableTemplate(path)
}

// directory couldn't be walked
if err != nil {
gologger.Labelf("Could not find templates in directory '%s': %s\n", r.templatesConfig.TemplatesDirectory, err)
}
return nil
},
)

// directory couldn't be walked
if err != nil {
gologger.Labelf("Could not find templates in directory '%s': %s\n", r.templatesConfig.TemplatesDirectory, err)
}
}

Expand Down
2 changes: 2 additions & 0 deletions internal/runner/update.go
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,7 @@ func (r *Runner) downloadReleaseAndUnzip(ctx context.Context, downloadURL string

templateDirectory := path.Join(r.templatesConfig.TemplatesDirectory, finalPath)
err = os.MkdirAll(templateDirectory, os.ModePerm)

if err != nil {
return fmt.Errorf("failed to create template folder %s : %s", templateDirectory, err)
}
Expand All @@ -264,6 +265,7 @@ func (r *Runner) downloadReleaseAndUnzip(ctx context.Context, downloadURL string
f.Close()
return fmt.Errorf("could not write template file: %s", err)
}

f.Close()
}

Expand Down

0 comments on commit 4484ba8

Please sign in to comment.