add checkJavaClass for validate template, update tests.

Author: For-ACGN

generator.go

@@ -17,17 +17,22 @@ func GenerateExecute(template []byte, command, class string) ([]byte, error) {
 		uint16Size   = 2
 	)
 
+	err := checkJavaClass(template)
+	if err != nil {
+		return nil, err
+	}
+
 	// find three special strings
 	fileNameIdx := bytes.Index(template, []byte(fileNameFlag))
-	if fileNameIdx == -1 || fileNameIdx < 2 {
+	if fileNameIdx == -1 {
 		return nil, errors.New("failed to find file name in execute template")
 	}
 	commandIdx := bytes.Index(template, []byte(commandFlag))
-	if commandIdx == -1 || commandIdx < 2 {
+	if commandIdx == -1 {
 		return nil, errors.New("failed to find command flag in execute template")
 	}
 	classNameIdx := bytes.Index(template, []byte(className))
-	if classNameIdx == -1 || classNameIdx < 2 {
+	if classNameIdx == -1 {
 		return nil, errors.New("failed to find class name in execute template")
 	}
 
@@ -77,25 +82,30 @@ func GenerateReverseTCP(template []byte, host string, port uint16, token, class
 		uint16Size   = 2
 	)
 
+	err := checkJavaClass(template)
+	if err != nil {
+		return nil, err
+	}
+
 	// find three special strings
 	fileNameIdx := bytes.Index(template, []byte(fileNameFlag))
-	if fileNameIdx == -1 || fileNameIdx < 2 {
+	if fileNameIdx == -1 {
 		return nil, errors.New("failed to find file name in reverse_tcp template")
 	}
 	hostIdx := bytes.Index(template, []byte(hostFlag))
-	if hostIdx == -1 || hostIdx < 2 {
+	if hostIdx == -1 {
 		return nil, errors.New("failed to find host flag in reverse_tcp template")
 	}
 	portIdx := bytes.Index(template, []byte(portFlag))
-	if portIdx == -1 || portIdx < 2 {
+	if portIdx == -1 {
 		return nil, errors.New("failed to find port flag in reverse_tcp template")
 	}
 	tokenIdx := bytes.Index(template, []byte(tokenFlag))
-	if tokenIdx == -1 || tokenIdx < 2 {
+	if tokenIdx == -1 {
 		return nil, errors.New("failed to find token flag in reverse_tcp template")
 	}
 	classNameIdx := bytes.Index(template, []byte(className))
-	if classNameIdx == -1 || classNameIdx < 2 {
+	if classNameIdx == -1 {
 		return nil, errors.New("failed to find class name in reverse_tcp template")
 	}
 
@@ -152,6 +162,16 @@ func GenerateReverseTCP(template []byte, host string, port uint16, token, class
 	return output.Bytes(), nil
 }
 
+func checkJavaClass(template []byte) error {
+	if len(template) < 4 {
+		return errors.New("invalid Java class template file size")
+	}
+	if !bytes.Equal(template[:2], []byte{0xCA, 0xFE}) {
+		return errors.New("invalid Java class template file")
+	}
+	return nil
+}
+
 func beUint16ToBytes(n uint16) []byte {
 	b := make([]byte, 2)
 	binary.BigEndian.PutUint16(b, n)

generator_test.go

@@ -49,6 +49,20 @@ func TestGenerateExecute(t *testing.T) {
 		require.Equal(t, expected, class)
 	})
 
+	t.Run("invalid template", func(t *testing.T) {
+		t.Run("invalid size", func(t *testing.T) {
+			class, err := GenerateExecute(nil, "", "")
+			require.EqualError(t, err, "invalid Java class template file size")
+			require.Zero(t, class)
+		})
+
+		t.Run("invalid data", func(t *testing.T) {
+			class, err := GenerateExecute(bytes.Repeat([]byte{0x00}, 8), "", "")
+			require.EqualError(t, err, "invalid Java class template file")
+			require.Zero(t, class)
+		})
+	})
+
 	t.Run("empty command", func(t *testing.T) {
 		class, err := GenerateExecute(template, "", "Test")
 		require.EqualError(t, err, "empty command")
@@ -83,6 +97,20 @@ func TestGenerateReverseTCP(t *testing.T) {
 		require.Equal(t, expected, class)
 	})
 
+	t.Run("invalid template", func(t *testing.T) {
+		t.Run("invalid size", func(t *testing.T) {
+			class, err := GenerateReverseTCP(nil, "", 0, "", "")
+			require.EqualError(t, err, "invalid Java class template file size")
+			require.Zero(t, class)
+		})
+
+		t.Run("invalid data", func(t *testing.T) {
+			class, err := GenerateReverseTCP(bytes.Repeat([]byte{0x00}, 8), "", 0, "", "")
+			require.EqualError(t, err, "invalid Java class template file")
+			require.Zero(t, class)
+		})
+	})
+
 	t.Run("empty host", func(t *testing.T) {
 		class, err := GenerateReverseTCP(template, "", 1234, "", "")
 		require.EqualError(t, err, "empty host")