improve tests about obfuscate, main.go: fix bug in generateClass().

Author: For-ACGN

cmd/main.go

@@ -142,10 +142,9 @@ func generateClass() {
 		generateExecute()
 	case "reverse_tcp":
 		generateReverseTCP()
-	case "":
-		fmt.Println("supported Java class template: execute, reverse_tcp")
-		return
 	default:
+		fmt.Println("supported Java class template: execute, reverse_tcp")
+		fmt.Println()
 		log.Fatalf("[error] unknown Java class template name: \"%s\"\n", genClass)
 	}
 	fmt.Println("Save generated Java class file to the path:", genOut)
@@ -163,7 +162,7 @@ func generateExecute() {
 
 	if command == "" {
 		args.PrintDefaults()
-		return
+		os.Exit(2)
 	}
 	if gnClass == "" {
 		gnClass = "Execute"
@@ -194,11 +193,11 @@ func generateReverseTCP() {
 
 	if host == "" {
 		args.PrintDefaults()
-		return
+		os.Exit(2)
 	}
 	if port > 65535 {
 		fmt.Println("[error]", "invalid port:", port)
-		return
+		os.Exit(2)
 	}
 	if gnClass == "" {
 		gnClass = "ReverseTCP"

obfuscate.go

@@ -102,8 +102,8 @@ func Obfuscate(raw string, token bool) (string, string) {
 
 		// generate useless data before section
 		obfuscated.WriteString("${")
-		n := 1 + rand.Intn(3) // 1-3 // #nosec
-		for i := 0; i < n; i++ {
+		round := 1 + rand.Intn(3) // 1-3 // #nosec
+		for i := 0; i < round; i++ {
 			front := randString(2 + rand.Intn(5)) // #nosec
 			end := randString(2 + rand.Intn(5))   // #nosec
 
@@ -122,7 +122,7 @@ func Obfuscate(raw string, token bool) (string, string) {
 		remaining -= size
 		index += size
 		lastObfuscated = true
-		// lastCharacter must be "}"
+		lastCharacter = '}' //  lastCharacter must be "}"
 	}
 
 	return obfuscated.String(), rwt

obfuscate_test.go

@@ -56,26 +56,28 @@ func TestObfuscate(t *testing.T) {
 
 	t.Run("fuzz", func(t *testing.T) {
 		t.Run("with token", func(t *testing.T) {
-			for i := 0; i < 100000; i++ {
+			for i := 0; i < 10000; i++ {
 				raw := "${" + randString(64) + "}"
 				obfuscated, rwt := Obfuscate(raw, true)
 				require.NotZero(t, rwt)
 				require.NotZero(t, obfuscated)
 
 				// check exist bug "$" with "${"
 				require.NotContains(t, obfuscated, "$${")
+				require.NotContains(t, obfuscated, " ")
 			}
 		})
 
 		t.Run("without token", func(t *testing.T) {
-			for i := 0; i < 100000; i++ {
+			for i := 0; i < 10000; i++ {
 				raw := "${" + randString(64) + "}"
 				obfuscated, rwt := Obfuscate(raw, false)
 				require.Zero(t, rwt)
 				require.NotZero(t, obfuscated)
 
 				// check exist bug "$" with "${"
 				require.NotContains(t, obfuscated, "$${")
+				require.NotContains(t, obfuscated, " ")
 			}
 		})
 	})
@@ -127,24 +129,26 @@ func TestObfuscateWithDollar(t *testing.T) {
 
 	t.Run("fuzz", func(t *testing.T) {
 		t.Run("with token", func(t *testing.T) {
-			for i := 0; i < 100000; i++ {
+			for i := 0; i < 10000; i++ {
 				raw := "${" + randString(64) + "}"
 				obfuscated, rwt := ObfuscateWithDollar(raw, true)
 				require.NotZero(t, rwt)
 				require.NotZero(t, obfuscated)
 
 				require.Equal(t, 1, strings.Count(obfuscated, "$${"))
+				require.NotContains(t, obfuscated, " ")
 			}
 		})
 
 		t.Run("without token", func(t *testing.T) {
-			for i := 0; i < 100000; i++ {
+			for i := 0; i < 10000; i++ {
 				raw := "${" + randString(64) + "}"
 				obfuscated, rwt := ObfuscateWithDollar(raw, false)
 				require.Zero(t, rwt)
 				require.NotZero(t, obfuscated)
 
 				require.NotContains(t, obfuscated, "$${")
+				require.NotContains(t, obfuscated, " ")
 			}
 		})
 	})

rand_test.go

@@ -1 +1,22 @@
 package log4shell
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRandString(t *testing.T) {
+	for i := 0; i < 10000; i++ {
+		str := randString(64)
+		require.False(t, strings.Contains(str, " "))
+	}
+}
+
+func TestRandSecret(t *testing.T) {
+	for i := 0; i < 10000; i++ {
+		str := randSecret()
+		require.False(t, strings.Contains(str, " "))
+	}
+}