From 3d03cf4b63d3912343f43a55ebfd56cf2c18c31d Mon Sep 17 00:00:00 2001
From: CX330Blake <timchen975@gmail.com>
Date: Tue, 10 Dec 2024 17:35:58 +0800
Subject: [PATCH] Refactor payload testing and add output handling; move test
 functions to letsgo package

---
 main.go              | 49 +++----------------------------------
 pkg/letsgo/letsgo.go | 58 ++++++++++++++++++++++++++++++++++++++++++++
 pkg/output/output.go | 15 ++++++++++++
 3 files changed, 76 insertions(+), 46 deletions(-)
 create mode 100644 pkg/letsgo/letsgo.go
 create mode 100644 pkg/output/output.go

diff --git a/main.go b/main.go
index 374290e..cc72beb 100644
--- a/main.go
+++ b/main.go
@@ -3,15 +3,12 @@ package main
 import (
 	"bufio"
 	"flag"
-	"fmt"
 	"io"
 	"log"
-	"net/http"
 	"os"
-	"strings"
-	"sync"
 
 	"github.com/CX330Blake/letsgo/pkg/greet"
+	"github.com/CX330Blake/letsgo/pkg/letsgo"
 	"github.com/fatih/color"
 )
 
@@ -39,47 +36,6 @@ func loadPayloads(filePath string) ([]string, bool) {
 	return payloads, useDefault
 }
 
-// Send request and check response
-func testPayload(url string, param string, fileRoot string, wordlist string) {
-	if !strings.HasSuffix(fileRoot, "/") {
-		fileRoot += "/"
-	}
-
-	fullURL := fmt.Sprintf("%s?%s=%s%s", url, param, fileRoot, wordlist)
-	resp, err := http.Get(fullURL)
-	if err != nil {
-		// color.Red("[!] Request failed: %v\n", err)
-		return
-	}
-	defer resp.Body.Close()
-
-	// Simple check if the response might indicate a vulnerability
-	if resp.StatusCode == http.StatusOK {
-		color.Green("[+] Found Possible Vuln URL: %s (Status: %d)\n", fullURL, resp.StatusCode)
-	}
-}
-
-// Multi-threaded testing
-func testPathTraversal(url string, param string, fileRoot string, wordlist []string) {
-	_, err := http.Get(url)
-	if err != nil {
-		color.Red("[!] Host seems down...\n")
-		return
-	}
-
-	var wg sync.WaitGroup
-
-	for _, payload := range wordlist {
-		wg.Add(1)
-		go func(wl string) {
-			defer wg.Done()
-			testPayload(url, param, fileRoot, wl)
-		}(payload)
-	}
-
-	wg.Wait()
-}
-
 func main() {
 
 	log.SetOutput(io.Discard)
@@ -88,6 +44,7 @@ func main() {
 	param := flag.String("param", "file", "Parameter for testing (default is 'file')")
 	wordlistFile := flag.String("wordlist", "default.txt", "Wordlist path")
 	fileRoot := flag.String("root", "", "Root of the server file (e.g. https://example.com/image?filename=/var/www/images/1337.jpg, then root is `/var/www/images`, don't need to include the last `/`)")
+	extension := flag.String("extension", "", "File extension (e.g. jpg, png, txt, etc.), this will triger the null byte bypass mode")
 
 	flag.Parse()
 
@@ -106,6 +63,6 @@ func main() {
 		color.Magenta("[+] Using default wordlist...\n")
 	}
 
-	testPathTraversal(*url, *param, *fileRoot, payloads)
+	letsgo.Test(*url, *param, *fileRoot, *extension, payloads)
 	greet.End()
 }
diff --git a/pkg/letsgo/letsgo.go b/pkg/letsgo/letsgo.go
new file mode 100644
index 0000000..145034e
--- /dev/null
+++ b/pkg/letsgo/letsgo.go
@@ -0,0 +1,58 @@
+package letsgo
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+	"sync"
+
+	"github.com/CX330Blake/letsgo/pkg/output"
+)
+
+// Send request and check response
+func testEach(url string, param string, fileRoot string, extension string, wordlist string) {
+	// If the file root is set, asure it ends with a slash
+	if fileRoot != "" && !strings.HasSuffix(fileRoot, "/") {
+		fileRoot += "/"
+	}
+
+	fullURL := fmt.Sprintf("%s?%s=%s%s", url, param, fileRoot, wordlist)
+	if extension != "" && strings.HasPrefix(extension, ".") {
+		fullURL = fmt.Sprintf("%s?%s=%s%s%%00%s", url, param, fileRoot, wordlist, extension)
+	} else if extension != "" && !strings.HasPrefix(extension, ".") {
+		fullURL = fmt.Sprintf("%s?%s=%s%s%%00.%s", url, param, fileRoot, wordlist, extension)
+	}
+
+	resp, err := http.Get(fullURL)
+	if err != nil {
+		// color.Red("[!] Request failed: %v\n", err)
+		return
+	}
+	defer resp.Body.Close()
+
+	// Simple check if the response might indicate a vulnerability
+	if resp.StatusCode == http.StatusOK {
+		output.Good("[+] Found Possible Vuln URL: %s (Status: %d)\n", fullURL, resp.StatusCode)
+	}
+}
+
+// Multi-threaded testing
+func Test(url string, param string, fileRoot string, extension string, wordlist []string) {
+	_, err := http.Get(url)
+	if err != nil {
+		output.Err("[!] Host seems down...\n")
+		return
+	}
+
+	var wg sync.WaitGroup
+
+	for _, payload := range wordlist {
+		wg.Add(1)
+		go func(wl string) {
+			defer wg.Done()
+			testEach(url, param, fileRoot, extension, wl)
+		}(payload)
+	}
+
+	wg.Wait()
+}
diff --git a/pkg/output/output.go b/pkg/output/output.go
new file mode 100644
index 0000000..c826b17
--- /dev/null
+++ b/pkg/output/output.go
@@ -0,0 +1,15 @@
+package output
+
+import "github.com/fatih/color"
+
+func Good(format string, a ...interface{}) {
+	color.Green(format, a...)
+}
+
+func Err(format string, a ...interface{}) {
+	color.Red(format, a...)
+}
+
+func Warn(format string, a ...interface{}) {
+	color.Yellow(format, a...)
+}