// Package rocketmq is a very basic (and incomplete) implementation of RocketMQ remoting protocol
package rocketmq

import (
	"encoding/binary"
	"encoding/json"
	"math"
	"net"

	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/protocol"
)

// Creates a RocketMQ remoting message. Format described here:
// https://github.com/apache/rocketmq/blob/bd0e9c09db9748f7f74a0c707579142dccf30afc/remoting/src/main/java/org/apache/rocketmq/remoting/protocol/RemotingCommand.java#L473
//
// [4 byte total length]
// [4 byte header length] (includes protocol type flag)
// [n bytes header data]
// [n bytes body data]
//
// The code can be found here:
// https://github.com/apache/rocketmq/blob/c78061bf6ca5f35452510ec4107c46735c51c316/test/src/test/resources/schema/protocol/common.protocol.RequestCode.schema
//
// This can return 'nil' if the requested payload is too large or marshalling the json header fails.
func CreateMqRemotingMessage(payload string, code int, version int) []byte {
	// Check the payload isn't overly large. I don't really see a scenario where 16 million bytes is
	// needed. Heck, I don't see needing more than 1k for the current exploitation use case.
	if len(payload) > int(math.Pow(2, 24)) {
		output.PrintFrameworkError("Payload size exceeded artificial cap.")

		return nil
	}
	mqHeader := map[string]any{
		"code":                    code,
		"flag":                    0,
		"language":                "JAVA",
		"opaque":                  0,
		"serializeTypeCurrentRPC": "JSON",
		"version":                 version,
	}
	mqHeaderBytes, err := json.Marshal(mqHeader)
	if err != nil {
		// this should never be reached
		output.PrintFrameworkError("Failed to marshal mqHeader")

		return nil
	}

	// build the total length field
	totalLengthField := make([]byte, 4)
	length := len(payload) + len(mqHeaderBytes) + 4
	binary.BigEndian.PutUint32(totalLengthField, uint32(length))

	// build the header length field
	headerLengthField := make([]byte, 4)
	binary.BigEndian.PutUint32(headerLengthField, uint32(len(mqHeaderBytes)))

	// create the full message
	message := make([]byte, 0)
	message = append(message, totalLengthField...)
	message = append(message, headerLengthField...)
	message = append(message, mqHeaderBytes...)
	message = append(message, []byte(payload)...)

	return message
}

// Reads a remoting message from the remote target.
func ReadMqRemotingResponse(conn net.Conn) ([]byte, []byte, bool) {
	msgSize, ok := protocol.TCPReadAmount(conn, 4)
	if !ok {
		return nil, nil, false
	}

	readSize := int(binary.BigEndian.Uint32(msgSize))
	if readSize == 0 {
		output.PrintFrameworkError("The server provided an invalid message length")

		return nil, nil, false
	}

	msg, ok := protocol.TCPReadAmount(conn, readSize)
	if !ok {
		return nil, nil, false
	}

	headerSize := int(binary.BigEndian.Uint32(msg[:4]))
	if (headerSize + 4) > len(msg) {
		output.PrintFrameworkError("Invalid remoting response")

		return nil, nil, false
	}

	header := msg[4 : headerSize+4]
	body := msg[headerSize+4:]

	return header, body, true
}