update GenerateReverseTCP, add test to generate fake class for pass TestGenerateReverseTCP.

Author: For-ACGN

generator.go

@@ -106,12 +106,12 @@ func GenerateReverseTCP(template []byte, host string, port uint16, token, class
 	if port == 0 {
 		return nil, errors.New("zero port")
 	}
-	if class == "" {
-		class = "ReverseTCP"
-	}
 	if token == "" {
 		token = randString(8)
 	}
+	if class == "" {
+		class = "ReverseTCP"
+	}
 
 	// generate output class file
 	output := bytes.NewBuffer(make([]byte, 0, len(template)+128))

generator_test.go

@@ -1,6 +1,8 @@
 package log4shell
 
 import (
+	"bytes"
+	"encoding/binary"
 	"os"
 	"testing"
 
@@ -93,3 +95,91 @@ func TestGenerateReverseTCP(t *testing.T) {
 		require.Zero(t, class)
 	})
 }
+
+func TestGenerateReverseTCP_Fake(t *testing.T) {
+	const (
+		fileNameFlag = "ReverseTCP.java"
+		hostFlag     = "${host}"
+		portFlag     = "${port}"
+		tokenFlag    = "${token}"
+		className    = "ReverseTCP\x0C"
+	)
+
+	buf := bytes.NewBuffer(make([]byte, 0, 128))
+	buf.Write([]byte{0xCA, 0xFE})
+	buf.Write([]byte{0x00, 0x00})
+
+	size := make([]byte, 2)
+
+	binary.BigEndian.PutUint16(size, uint16(len(fileNameFlag)))
+	buf.Write(size)
+	buf.WriteString(fileNameFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(hostFlag)))
+	buf.Write(size)
+	buf.WriteString(hostFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(portFlag)))
+	buf.Write(size)
+	buf.WriteString(portFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(tokenFlag)))
+	buf.Write(size)
+	buf.WriteString(tokenFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(className)))
+	buf.Write(size)
+	buf.WriteString(className)
+	buf.Write([]byte{0x00, 0x00})
+
+	err := os.WriteFile("testdata/template/ReverseTCP.class", buf.Bytes(), 0600)
+	require.NoError(t, err)
+}
+
+func TestGenerateReverseTCP_Fake_Compare(t *testing.T) {
+	const (
+		fileNameFlag = "ReTCP.java"
+		hostFlag     = "127.0.0.1"
+		portFlag     = "9979"
+		tokenFlag    = "test"
+		className    = "ReTCP\x0C"
+	)
+
+	buf := bytes.NewBuffer(make([]byte, 0, 128))
+	buf.Write([]byte{0xCA, 0xFE})
+	buf.Write([]byte{0x00, 0x00})
+
+	size := make([]byte, 2)
+
+	binary.BigEndian.PutUint16(size, uint16(len(fileNameFlag)))
+	buf.Write(size)
+	buf.WriteString(fileNameFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(hostFlag)))
+	buf.Write(size)
+	buf.WriteString(hostFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(portFlag)))
+	buf.Write(size)
+	buf.WriteString(portFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(tokenFlag)))
+	buf.Write(size)
+	buf.WriteString(tokenFlag)
+	buf.Write([]byte{0x00, 0x00})
+
+	binary.BigEndian.PutUint16(size, uint16(len(className)-1))
+	buf.Write(size)
+	buf.WriteString(className)
+	buf.Write([]byte{0x00, 0x00})
+
+	err := os.WriteFile("testdata/template/compare/ReTCP.class", buf.Bytes(), 0600)
+	require.NoError(t, err)
+}

log4shell_test.go

@@ -19,7 +19,7 @@ func testGenerateConfig() *Config {
 	return &Config{
 		Logger:      os.Stdout,
 		Hostname:    "127.0.0.1",
-		PayloadDir:  "testdata",
+		PayloadDir:  "testdata/payload",
 		HTTPNetwork: "tcp",
 		HTTPAddress: "127.0.0.1:8088",
 		LDAPNetwork: "tcp",