Added -ai option to generate and run nuclei templates on the fly for given prompt (#6041)

* Add ai flag

* Add AI flag 2

* fix stdin

* fix stdin 2

* minor

* print both url and path

* store ai generated templates in `$HOME/nuclei-templates/pdcp`

* todo

* do not remove all

* make it less restrictive

* use retryablehttp

* fix creds check

* return errs

* return more detailed err for non-ok status code

* add prompt validation

* fix integration tests

---------

Co-authored-by: Doğan Can Bakır <dogancanbakir@protonmail.com>

cmd/nuclei/main.go

@@ -252,6 +252,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.BoolVarP(&options.AutomaticScan, "automatic-scan", "as", false, "automatic web scan using wappalyzer technology detection to tags mapping"),
 		flagSet.StringSliceVarP(&options.Templates, "templates", "t", nil, "list of template or template directory to run (comma-separated, file)", goflags.FileCommaSeparatedStringSliceOptions),
 		flagSet.StringSliceVarP(&options.TemplateURLs, "template-url", "turl", nil, "template url or list containing template urls to run (comma-separated, file)", goflags.FileCommaSeparatedStringSliceOptions),
+		flagSet.StringVarP(&options.AITemplatePrompt, "prompt", "ai", "", "generate and run template using ai prompt"),
 		flagSet.StringSliceVarP(&options.Workflows, "workflows", "w", nil, "list of workflow or workflow directory to run (comma-separated, file)", goflags.FileCommaSeparatedStringSliceOptions),
 		flagSet.StringSliceVarP(&options.WorkflowURLs, "workflow-url", "wurl", nil, "workflow url or list containing workflow urls to run (comma-separated, file)", goflags.FileCommaSeparatedStringSliceOptions),
 		flagSet.BoolVar(&options.Validate, "validate", false, "validate the passed templates to nuclei"),

integration_tests/protocols/multi/dynamic-values.yaml

@@ -24,6 +24,6 @@ http:
     matchers:
       - type: dsl
         dsl:
-          - contains(body,'introduction') # check for http string
+          - contains(body,'home') # check for http string
           - blogid == 'cname' # check for cname (extracted information from dns response)
         condition: and

integration_tests/protocols/multi/evaluate-variables.yaml

@@ -24,7 +24,7 @@ http:
     matchers:
       - type: dsl
         dsl:
-          - contains(http_body,'introduction') # check for http string
+          - contains(http_body,'home') # check for http string
           - cname_filtered == 'cname' # check for cname (extracted information from dns response)
           - ssl_subject_cn == 'docs.projectdiscovery.io'
         condition: and

integration_tests/protocols/multi/exported-response-vars.yaml

@@ -20,7 +20,7 @@ http:
     matchers:
       - type: dsl
         dsl:
-          - contains(http_body,'introduction') # check for http string
+          - contains(http_body,'home') # check for http string
           - trim_suffix(dns_cname,'.vercel-dns.com') == 'cname' # check for cname (extracted information from dns response)
           - ssl_subject_cn == 'docs.projectdiscovery.io'
         condition: and

pkg/catalog/loader/ai_loader.go

@@ -0,0 +1,141 @@
+package loader
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/alecthomas/chroma/quick"
+	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
+	"github.com/projectdiscovery/nuclei/v3/pkg/types"
+	"github.com/projectdiscovery/retryablehttp-go"
+	pdcpauth "github.com/projectdiscovery/utils/auth/pdcp"
+	errorutil "github.com/projectdiscovery/utils/errors"
+)
+
+const (
+	aiTemplateGeneratorAPIEndpoint = "https://api.projectdiscovery.io/v1/template/ai"
+)
+
+type AITemplateResponse struct {
+	CanRun     bool   `json:"canRun"`
+	Comment    string `json:"comment"`
+	Completion string `json:"completion"`
+	Message    string `json:"message"`
+	Name       string `json:"name"`
+	TemplateID string `json:"template_id"`
+}
+
+func getAIGeneratedTemplates(prompt string, options *types.Options) ([]string, error) {
+	prompt = strings.TrimSpace(prompt)
+	if len(prompt) < 5 {
+		return nil, errorutil.New("Prompt is too short. Please provide a more descriptive prompt")
+	}
+
+	if len(prompt) > 3000 {
+		return nil, errorutil.New("Prompt is too long. Please limit to 3000 characters")
+	}
+
+	template, templateID, err := generateAITemplate(prompt)
+	if err != nil {
+		return nil, errorutil.New("Failed to generate template: %v", err)
+	}
+
+	pdcpTemplateDir := filepath.Join(config.DefaultConfig.GetTemplateDir(), "pdcp")
+	if err := os.MkdirAll(pdcpTemplateDir, 0755); err != nil {
+		return nil, errorutil.New("Failed to create pdcp template directory: %v", err)
+	}
+
+	templateFile := filepath.Join(pdcpTemplateDir, templateID+".yaml")
+	err = os.WriteFile(templateFile, []byte(template), 0644)
+	if err != nil {
+		return nil, errorutil.New("Failed to generate template: %v", err)
+	}
+
+	gologger.Info().Msgf("Generated template available at: https://cloud.projectdiscovery.io/templates/%s", templateID)
+	gologger.Info().Msgf("Generated template path: %s", templateFile)
+
+	// Check if we should display the template
+	// This happens when:
+	// 1. No targets are provided (-target/-list)
+	// 2. No stdin input is being used
+	hasNoTargets := len(options.Targets) == 0 && options.TargetsFilePath == ""
+	hasNoStdin := !options.Stdin
+
+	if hasNoTargets && hasNoStdin {
+		// Display the template content with syntax highlighting
+		if !options.NoColor {
+			var buf bytes.Buffer
+			err = quick.Highlight(&buf, template, "yaml", "terminal16m", "monokai")
+			if err == nil {
+				template = buf.String()
+			}
+		}
+		gologger.Silent().Msgf("\n%s", template)
+		// FIXME:
+		// we should not be exiting the program here
+		// but we need to find a better way to handle this
+		os.Exit(0)
+	}
+
+	return []string{templateFile}, nil
+}
+
+func generateAITemplate(prompt string) (string, string, error) {
+	reqBody := map[string]string{
+		"prompt": prompt,
+	}
+	jsonBody, err := json.Marshal(reqBody)
+	if err != nil {
+		return "", "", errorutil.New("Failed to marshal request body: %v", err)
+	}
+
+	req, err := http.NewRequest(http.MethodPost, aiTemplateGeneratorAPIEndpoint, bytes.NewBuffer(jsonBody))
+	if err != nil {
+		return "", "", errorutil.New("Failed to create HTTP request: %v", err)
+	}
+
+	ph := pdcpauth.PDCPCredHandler{}
+	creds, err := ph.GetCreds()
+	if err != nil {
+		return "", "", errorutil.New("Failed to get PDCP credentials: %v", err)
+	}
+
+	if creds == nil {
+		return "", "", errorutil.New("PDCP API Key not configured, Create one for free at https://cloud.projectdiscovery.io/")
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set(pdcpauth.ApiKeyHeaderName, creds.APIKey)
+
+	resp, err := retryablehttp.DefaultClient().Do(req)
+	if err != nil {
+		return "", "", errorutil.New("Failed to send HTTP request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusUnauthorized {
+		return "", "", errorutil.New("Invalid API Key or API Key not configured, Create one for free at https://cloud.projectdiscovery.io/")
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return "", "", errorutil.New("API returned status code %d: %s", resp.StatusCode, string(body))
+	}
+
+	var result AITemplateResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return "", "", errorutil.New("Failed to decode API response: %v", err)
+	}
+
+	if result.TemplateID == "" || result.Completion == "" {
+		return "", "", errorutil.New("Failed to generate template")
+	}
+
+	return result.Completion, result.TemplateID, nil
+}

pkg/catalog/loader/loader.go

@@ -50,6 +50,7 @@ type Config struct {
 	ExcludeTemplates         []string
 	IncludeTemplates         []string
 	RemoteTemplateDomainList []string
+	AITemplatePrompt         string
 
 	Tags              []string
 	ExcludeTags       []string
@@ -109,6 +110,7 @@ func NewConfig(options *types.Options, catalog catalog.Catalog, executerOpts pro
 		IncludeConditions:        options.IncludeConditions,
 		Catalog:                  catalog,
 		ExecutorOptions:          executerOpts,
+		AITemplatePrompt:         options.AITemplatePrompt,
 	}
 	loaderConfig.RemoteTemplateDomainList = append(loaderConfig.RemoteTemplateDomainList, TrustedTemplateDomains...)
 	return &loaderConfig
@@ -133,7 +135,6 @@ func New(cfg *Config) (*Store, error) {
 		return nil, err
 	}
 
-	// Create a tag filter based on provided configuration
 	store := &Store{
 		id:        cfg.StoreId,
 		config:    cfg,
@@ -164,7 +165,6 @@ func New(cfg *Config) (*Store, error) {
 		if _, err := urlutil.Parse(v); err == nil {
 			remoteTemplates = append(remoteTemplates, handleTemplatesEditorURLs(v))
 		} else {
-
 			templatesFinal = append(templatesFinal, v) // something went wrong, treat it as a file
 		}
 	}
@@ -181,6 +181,15 @@ func New(cfg *Config) (*Store, error) {
 		store.finalWorkflows = append(store.finalWorkflows, remoteWorkflows...)
 	}
 
+	// Handle AI template generation if prompt is provided
+	if len(cfg.AITemplatePrompt) > 0 {
+		aiTemplates, err := getAIGeneratedTemplates(cfg.AITemplatePrompt, cfg.ExecutorOptions.Options)
+		if err != nil {
+			return nil, err
+		}
+		store.finalTemplates = append(store.finalTemplates, aiTemplates...)
+	}
+
 	// Handle a dot as the current working directory
 	if len(store.finalTemplates) == 1 && store.finalTemplates[0] == "." {
 		currentDirectory, err := os.Getwd()
@@ -189,6 +198,7 @@ func New(cfg *Config) (*Store, error) {
 		}
 		store.finalTemplates = []string{currentDirectory}
 	}
+
 	// Handle a case with no templates or workflows, where we use base directory
 	if len(store.finalTemplates) == 0 && len(store.finalWorkflows) == 0 && !urlBasedTemplatesProvided {
 		store.finalTemplates = []string{config.DefaultConfig.TemplatesDirectory}
@@ -381,8 +391,8 @@ func (store *Store) areWorkflowOrTemplatesValid(filteredTemplatePaths map[string
 			// `ErrGlobalMatchersTemplate` during `templates.Parse` and checking it
 			// with `errors.Is`.
 			//
-			// However, I’m not sure if every reference to it should be handled
-			// that way. Returning a `templates.Template` pointer would mean it’s
+			// However, I'm not sure if every reference to it should be handled
+			// that way. Returning a `templates.Template` pointer would mean it's
 			// an active template (sending requests), and adding a specific field
 			// like `isGlobalMatchers` in `templates.Template` (then checking it
 			// with a `*templates.Template.IsGlobalMatchersEnabled` method) would

pkg/types/types.go

@@ -42,6 +42,8 @@ type Options struct {
 	Templates goflags.StringSlice
 	// TemplateURLs specifies URLs to a list of templates to use
 	TemplateURLs goflags.StringSlice
+	// AITemplatePrompt specifies prompt to generate template using AI
+	AITemplatePrompt string
 	// RemoteTemplates specifies list of allowed URLs to load remote templates from
 	RemoteTemplateDomainList goflags.StringSlice
 	// 	ExcludedTemplates  specifies the template/templates to exclude