From d0463e8ca39cf28e5e25b6426d041eab4c30b569 Mon Sep 17 00:00:00 2001
From: "alban.stourbe stourbe" <alban.stourbe@winamax.fr>
Date: Mon, 30 Sep 2024 16:52:54 +0200
Subject: [PATCH 1/2] Add vars with SDK

---
 lib/config.go | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/lib/config.go b/lib/config.go
index 7a9419aa50..e2a56c1375 100644
--- a/lib/config.go
+++ b/lib/config.go
@@ -392,6 +392,20 @@ 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 {
+		runtimeVars.Set(v)
+	}
+
+	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 {

From ed3e06a7011a92da046549916a2e499b3a741e4b Mon Sep 17 00:00:00 2001
From: "alban.stourbe stourbe" <alban.stourbe@winamax.fr>
Date: Tue, 1 Oct 2024 09:29:43 +0200
Subject: [PATCH 2/2] Add checkerr on runtimeVars & Add a new TestSDK

---
 lib/config.go         |  7 ++++++-
 lib/tests/sdk_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/lib/config.go b/lib/config.go
index e2a56c1375..20b3449c4c 100644
--- a/lib/config.go
+++ b/lib/config.go
@@ -397,7 +397,12 @@ func WithVars(vars []string) NucleiSDKOptions {
 	// Create a goflags.RuntimeMap
 	runtimeVars := goflags.RuntimeMap{}
 	for _, v := range vars {
-		runtimeVars.Set(v)
+		err := runtimeVars.Set(v)
+		if err != nil {
+			return func(e *NucleiEngine) error {
+				return err
+			}
+		}
 	}
 
 	return func(e *NucleiEngine) error {
diff --git a/lib/tests/sdk_test.go b/lib/tests/sdk_test.go
index 97ec489abc..395acfd515 100644
--- a/lib/tests/sdk_test.go
+++ b/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()
+	}
+}