From c7ac1648424fbba6dda6fe9c3749254f97b3f7c9 Mon Sep 17 00:00:00 2001
From: Bar Hofesh <bar.hofesh@gmail.com>
Date: Wed, 3 Jul 2024 12:50:47 +0300
Subject: [PATCH] feat(test): Allow overwriting the original value in function
 testing

---
 README.md               | 25 +++++++++++++++++++++++++
 spec/sec_tester_spec.cr | 20 ++++++++++++++++++++
 src/sec_tester/test.cr  |  4 ++--
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 03e4349..0187943 100644
--- a/README.md
+++ b/README.md
@@ -107,6 +107,31 @@ tester.run_check(scan_name: "UnitTestingScan - XSS - function only", tests: ["xs
 end
 ```
 
+> **Note**
+>
+> You also have an optional "param_overwrite" parameter that allows you to overwrite the parameters in the request.
+> This is useful when your function is expecting a specific data object like JSON or JWT etc..
+
+You can use the `param_overwrite` to overwrite the value to be attacked like:
+
+```crystal
+jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+tester.run_check(scan_name: "jwt-testing", tests: ["jwt"], param_overwrite: jwt) do |payload, response|
+  spawn do
+    while payload_data = payload.receive?
+      # This is where we send the payload to the function and send back a response
+      # In this example we just want to send back the payload
+      # as we are testing reflection
+      # my_function is a demo function that returns the payload
+      response_data = my_JWT_verification(payload_data)
+
+      # we end up sending the response back to the channel
+      response.send(response_data)
+    end
+  end
+end
+```
+
 There is also a variant of this interface that accepts target and yields back the whole HTTP::Server::Context.
 This is useful if you want to do something with the response body or headers.
 
diff --git a/spec/sec_tester_spec.cr b/spec/sec_tester_spec.cr
index 3ab059a..ed71394 100644
--- a/spec/sec_tester_spec.cr
+++ b/spec/sec_tester_spec.cr
@@ -286,6 +286,26 @@ describe SecTester::Test do
     end
   end
 
+  it "starts a function oriented test for XSS including param_overwrite" do
+    tester = SecTester::Test.new
+    expect_raises(SecTester::IssueFound) do
+      tester.run_check(scan_name: "xss", tests: ["xss"], param_overwrite: "abcdefu") do |payload, response|
+        spawn do
+          while payload_data = payload.receive?
+            # This is where we send the payload to the function and send back a response
+            # In this example we just want to send back the payload
+            # as we are testing reflection
+            # my_function is a demo function that returns the payload
+            response_data = my_function(payload_data)
+
+            # we end up sending the response back to the channel
+            response.send(response_data)
+          end
+        end
+      end
+    end
+  end
+
   it "starts a request/response oriented test for XSS" do
     tester = SecTester::Test.new
     target = SecTester::Target.new(
diff --git a/src/sec_tester/test.cr b/src/sec_tester/test.cr
index 5ff977a..65a7f5f 100644
--- a/src/sec_tester/test.cr
+++ b/src/sec_tester/test.cr
@@ -71,7 +71,7 @@ module SecTester
       server.try &.close
     end
 
-    def run_check(scan_name : String, tests : String | Array(String)?, severity_threshold : Severity = :low, options : Options = Options.new, on_issue : Bool = true, &)
+    def run_check(scan_name : String, tests : String | Array(String)?, severity_threshold : Severity = :low, options : Options = Options.new, on_issue : Bool = true, param_overwrite : String? = nil, &)
       # Start a server for the user, in this form we can test specific functions.
       payload = Channel(String).new
       response = Channel(String).new
@@ -95,7 +95,7 @@ module SecTester
       end
 
       target = Target.new(
-        url: "http://#{addr}?artificial=dummydata",
+        url: "http://#{addr}?artificial=#{param_overwrite || "dummydata"}",
       )
 
       yield payload, response