diff --git a/cmd/nuclei/main.go b/cmd/nuclei/main.go index a4ba02ddd..e619e0822 100644 --- a/cmd/nuclei/main.go +++ b/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"), diff --git a/integration_tests/protocols/multi/dynamic-values.yaml b/integration_tests/protocols/multi/dynamic-values.yaml index 4872c6bf3..b6115169b 100644 --- a/integration_tests/protocols/multi/dynamic-values.yaml +++ b/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 \ No newline at end of file diff --git a/integration_tests/protocols/multi/evaluate-variables.yaml b/integration_tests/protocols/multi/evaluate-variables.yaml index 7d0338ca2..ffb75110a 100644 --- a/integration_tests/protocols/multi/evaluate-variables.yaml +++ b/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 \ No newline at end of file diff --git a/integration_tests/protocols/multi/exported-response-vars.yaml b/integration_tests/protocols/multi/exported-response-vars.yaml index b6b941723..569756733 100644 --- a/integration_tests/protocols/multi/exported-response-vars.yaml +++ b/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 \ No newline at end of file diff --git a/pkg/catalog/loader/ai_loader.go b/pkg/catalog/loader/ai_loader.go new file mode 100644 index 000000000..64af39939 --- /dev/null +++ b/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 +} diff --git a/pkg/catalog/loader/loader.go b/pkg/catalog/loader/loader.go index ae4aeeca6..31ac7f41e 100644 --- a/pkg/catalog/loader/loader.go +++ b/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 diff --git a/pkg/types/types.go b/pkg/types/types.go index aa43683ec..2d29f1c81 100644 --- a/pkg/types/types.go +++ b/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