Merge pull request #5678 from alban-stourbe-wmx/feature/sdk-add-vars

Add vars with SDK

lib/config.go

@@ -392,6 +392,25 @@ func WithHeaders(headers []string) NucleiSDKOptions {
 	}
 }
 
+// WithVars allows setting custom variables to use in templates/workflows context
+func WithVars(vars []string) NucleiSDKOptions {
+	// Create a goflags.RuntimeMap
+	runtimeVars := goflags.RuntimeMap{}
+	for _, v := range vars {
+		err := runtimeVars.Set(v)
+		if err != nil {
+			return func(e *NucleiEngine) error {
+				return err
+			}
+		}
+	}
+
+	return func(e *NucleiEngine) error {
+		e.opts.Vars = runtimeVars
+		return nil
+	}
+}
+
 // EnablePassiveMode allows enabling passive HTTP response processing mode
 func EnablePassiveMode() NucleiSDKOptions {
 	return func(e *NucleiEngine) error {

lib/tests/sdk_test.go

@@ -133,3 +133,36 @@ func TestThreadSafeNuclei(t *testing.T) {
 		fn()
 	}
 }
+
+func TestWithVarsNuclei(t *testing.T) {
+	fn := func() {
+		defer func() {
+			// resources like leveldb have a delay to commit in-memory resources
+			// to disk, typically 1-2 seconds, so we wait for 2 seconds
+			time.Sleep(2 * time.Second)
+			goleak.VerifyNone(t, knownLeaks...)
+		}()
+		ne, err := nuclei.NewNucleiEngineCtx(
+			context.TODO(),
+			nuclei.WithTemplatesOrWorkflows(nuclei.TemplateSources{Templates: []string{"http/token-spray/api-1forge.yaml"}}),
+			nuclei.WithVars([]string{"token=foobar"}),
+			nuclei.WithVerbosity(nuclei.VerbosityOptions{Debug: true}),
+		)
+		require.Nil(t, err)
+		ne.LoadTargets([]string{"scanme.sh"}, true) // probe http/https target is set to true here
+		err = ne.ExecuteWithCallback(nil)
+		require.Nil(t, err)
+		defer ne.Close()
+	}
+	// this is shared test so needs to be run as seperate process
+	if env.GetEnvOrDefault("TestWithVarsNuclei", false) {
+		cmd := exec.Command(os.Args[0], "-test.run=TestWithVarsNuclei")
+		cmd.Env = append(os.Environ(), "TestWithVarsNuclei=true")
+		out, err := cmd.CombinedOutput()
+		if err != nil {
+			t.Fatalf("process ran with error %s, output: %s", err, out)
+		}
+	} else {
+		fn()
+	}
+}