Skip to content

Commit

Permalink
Add profile option to load template profile (#5125)
Browse files Browse the repository at this point in the history
* Add profile  option to load template profile

* Misc update

* Add profile-list option

* Misc update

* Add tests
  • Loading branch information
RamanaReddy0M authored May 4, 2024
1 parent 9784ca8 commit 902eb78
Show file tree
Hide file tree
Showing 6 changed files with 140 additions and 3 deletions.
1 change: 1 addition & 0 deletions cmd/integration-test/integration-test.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ var (
"dns": dnsTestCases,
"workflow": workflowTestcases,
"loader": loaderTestcases,
"profile-loader": profileLoaderTestcases,
"websocket": websocketTestCases,
"headless": headlessTestcases,
"whois": whoisTestCases,
Expand Down
53 changes: 53 additions & 0 deletions cmd/integration-test/profile-loader.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
package main

import (
"fmt"

"github.com/projectdiscovery/nuclei/v3/pkg/testutils"
errorutil "github.com/projectdiscovery/utils/errors"
)

var profileLoaderTestcases = []TestCaseInfo{
{Path: "profile-loader/load-with-filename", TestCase: &profileLoaderByRelFile{}},
{Path: "profile-loader/load-with-id", TestCase: &profileLoaderById{}},
{Path: "profile-loader/basic.yml", TestCase: &customProfileLoader{}},
}

type profileLoaderByRelFile struct{}

func (h *profileLoaderByRelFile) Execute(testName string) error {
results, err := testutils.RunNucleiWithArgsAndGetResults(false, "-tl", "-tp", "kev.yml")
if err != nil {
return errorutil.NewWithErr(err).Msgf("failed to load template with id")
}
if len(results) < 267 {
return fmt.Errorf("incorrect result: expected more results than %d, got %v", 267, len(results))
}
return nil
}

type profileLoaderById struct{}

func (h *profileLoaderById) Execute(testName string) error {
results, err := testutils.RunNucleiWithArgsAndGetResults(false, "-tl", "-tp", "kev")
if err != nil {
return errorutil.NewWithErr(err).Msgf("failed to load template with id")
}
if len(results) < 267 {
return fmt.Errorf("incorrect result: expected more results than %d, got %v", 267, len(results))
}
return nil
}

type customProfileLoader struct{}

func (h *customProfileLoader) Execute(filepath string) error {
results, err := testutils.RunNucleiWithArgsAndGetResults(false, "-tl", "-tp", filepath)
if err != nil {
return errorutil.NewWithErr(err).Msgf("failed to load template with id")
}
if len(results) < 267 {
return fmt.Errorf("incorrect result: expected more results than %d, got %v", 267, len(results))
}
return nil
}
58 changes: 55 additions & 3 deletions cmd/nuclei/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,10 @@ import (
)

var (
cfgFile string
memProfile string // optional profile file path
options = &types.Options{}
cfgFile string
templateProfile string
memProfile string // optional profile file path
options = &types.Options{}
)

func main() {
Expand Down Expand Up @@ -270,6 +271,8 @@ on extensive configurability, massive extensibility and ease of use.`)

flagSet.CreateGroup("configs", "Configurations",
flagSet.StringVar(&cfgFile, "config", "", "path to the nuclei configuration file"),
flagSet.StringVarP(&templateProfile, "profile", "tp", "", "template profile config file to run"),
flagSet.BoolVarP(&options.ListTemplateProfiles, "profile-list", "tpl", false, "list community template profiles"),
flagSet.BoolVarP(&options.FollowRedirects, "follow-redirects", "fr", false, "enable following redirects for http templates"),
flagSet.BoolVarP(&options.FollowHostRedirects, "follow-host-redirects", "fhr", false, "follow redirects on the same host"),
flagSet.IntVarP(&options.MaxRedirects, "max-redirects", "mr", 10, "max number of redirects to follow for http templates"),
Expand Down Expand Up @@ -497,6 +500,34 @@ Additional documentation is available at: https://docs.nuclei.sh/getting-started
config.DefaultConfig.SetTemplatesDir(options.NewTemplatesDirectory)
}

defaultProfilesPath := filepath.Join(config.DefaultConfig.GetTemplateDir(), "profiles")
if templateProfile != "" {
if filepath.Ext(templateProfile) == "" {
if tp := findProfilePathById(templateProfile, defaultProfilesPath); tp != "" {
templateProfile = tp
} else {
gologger.Fatal().Msgf("'%s' is not a profile-id or profile path", templateProfile)
}
}
if !filepath.IsAbs(templateProfile) {
if filepath.Dir(templateProfile) == "profiles" {
defaultProfilesPath = filepath.Join(config.DefaultConfig.GetTemplateDir())
}
currentDir, err := os.Getwd()
if err == nil && fileutil.FileExists(filepath.Join(currentDir, templateProfile)) {
templateProfile = filepath.Join(currentDir, templateProfile)
} else {
templateProfile = filepath.Join(defaultProfilesPath, templateProfile)
}
}
if !fileutil.FileExists(templateProfile) {
gologger.Fatal().Msgf("given template profile file '%s' does not exist", templateProfile)
}
if err := flagSet.MergeConfigFile(templateProfile); err != nil {
gologger.Fatal().Msgf("Could not read template profile: %s\n", err)
}
}

if len(options.SecretsFile) > 0 {
for _, secretFile := range options.SecretsFile {
if !fileutil.FileExists(secretFile) {
Expand Down Expand Up @@ -622,6 +653,27 @@ Note: Make sure you have backup of your custom nuclei-templates before proceedin
os.Exit(0)
}

func findProfilePathById(profileId, templatesDir string) string {
var profilePath string
err := filepath.WalkDir(templatesDir, func(iterItem string, d fs.DirEntry, err error) error {
ext := filepath.Ext(iterItem)
isYaml := ext == extensions.YAML || ext == extensions.YML
if err != nil || d.IsDir() || !isYaml {
// skip non yaml files
return nil
}
if strings.TrimSuffix(filepath.Base(iterItem), ext) == profileId {
profilePath = iterItem
return fmt.Errorf("FOUND")
}
return nil
})
if err != nil && err.Error() != "FOUND" {
gologger.Error().Msgf("%s\n", err)
}
return profilePath
}

func init() {
// print stacktrace of errors in debug mode
if strings.EqualFold(os.Getenv("DEBUG"), "true") {
Expand Down
2 changes: 2 additions & 0 deletions integration_tests/profile-loader/basic.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
tags:
- kev
27 changes: 27 additions & 0 deletions internal/runner/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package runner
import (
"bufio"
"fmt"
"io/fs"
"os"
"path/filepath"
"strconv"
Expand All @@ -25,6 +26,7 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/reporting/exporters/jsonl"
"github.com/projectdiscovery/nuclei/v3/pkg/reporting/exporters/markdown"
"github.com/projectdiscovery/nuclei/v3/pkg/reporting/exporters/sarif"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/extensions"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
"github.com/projectdiscovery/nuclei/v3/pkg/utils/yaml"
fileutil "github.com/projectdiscovery/utils/file"
Expand Down Expand Up @@ -74,6 +76,31 @@ func ParseOptions(options *types.Options) {
}
os.Exit(0)
}

defaultProfilesPath := filepath.Join(config.DefaultConfig.GetTemplateDir(), "profiles")
if options.ListTemplateProfiles {
gologger.Print().Msgf(
"\nListing available %v nuclei template profiles for %v",
config.DefaultConfig.TemplateVersion,
config.DefaultConfig.TemplatesDirectory,
)
templatesRootDir := config.DefaultConfig.GetTemplateDir()
err := filepath.WalkDir(defaultProfilesPath, func(iterItem string, d fs.DirEntry, err error) error {
ext := filepath.Ext(iterItem)
isYaml := ext == extensions.YAML || ext == extensions.YML
if err != nil || d.IsDir() || !isYaml {
return nil
}
if profileRelPath, err := filepath.Rel(templatesRootDir, iterItem); err == nil {
gologger.Print().Msgf("%s (%s)\n", profileRelPath, strings.TrimSuffix(filepath.Base(iterItem), ext))
}
return nil
})
if err != nil {
gologger.Error().Msgf("%s\n", err)
}
os.Exit(0)
}
if options.StoreResponseDir != DefaultDumpTrafficOutputFolder && !options.StoreResponse {
gologger.Debug().Msgf("Store response directory specified, enabling \"store-resp\" flag automatically\n")
options.StoreResponse = true
Expand Down
2 changes: 2 additions & 0 deletions pkg/types/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -393,6 +393,8 @@ type Options struct {
DAST bool
// HttpApiEndpoint is the experimental http api endpoint
HttpApiEndpoint string
// ListTemplateProfiles lists all available template profiles
ListTemplateProfiles bool
}

// ShouldLoadResume resume file
Expand Down

0 comments on commit 902eb78

Please sign in to comment.