// `winbox.go` contains the logic for sending unencrypted M2 messages to
// the RouterOS Winbox port (8291)
package mikrotik

import (
	"encoding/binary"
	"net"

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

func SendM2(conn net.Conn, msg *M2Message) bool {
	// craft m2 message
	m2Payload := []byte("M2")
	m2Payload = append(m2Payload, (msg.Serialize())...)

	// add size in front of the message
	outputSize := make([]byte, 2)
	binary.BigEndian.PutUint16(outputSize, uint16(len(m2Payload)))
	sizeAndPayload := []byte{}
	sizeAndPayload = append(sizeAndPayload, outputSize...)
	sizeAndPayload = append(sizeAndPayload, m2Payload...)

	// add the pkt header (1 byte length, 1 byte messges)
	// not that this doesn't work for >0xff, and should be reworked
	pktHeader := make([]byte, 2)
	pktHeader[0] = byte(len(sizeAndPayload))
	pktHeader[1] = 1
	finalMsg := []byte{}
	finalMsg = append(finalMsg, pktHeader...)
	finalMsg = append(finalMsg, sizeAndPayload...)

	return protocol.TCPWrite(conn, finalMsg)
}

func ReceiveM2(conn net.Conn, msg *M2Message) bool {
	msgSize, ok := protocol.TCPReadAmount(conn, 4)
	if !ok {
		return false
	}

	msgSize = msgSize[2:]
	readSize := int(binary.BigEndian.Uint16(msgSize))
	if readSize == 0 {
		output.PrintfFrameworkError("The server provided an invalid message length")

		return false
	}

	finalBuffer := []byte{}
	loop := 0

	for readSize > 0xff {
		step := 0xff
		if loop == 0 {
			step -= 2
		}

		msgBytes, ok := protocol.TCPReadAmount(conn, step)
		if !ok {
			return false
		}

		loop++
		readSize -= step
		finalBuffer = append(finalBuffer, msgBytes...)

		// this should be two bytes of padding
		msgBytes, ok = protocol.TCPReadAmount(conn, 2)
		if !ok {
			return false
		}

		if msgBytes[1] != 0xff {
			output.PrintFrameworkError("Padding is off")

			return false
		}
	}

	msgBytes, ok := protocol.TCPReadAmount(conn, readSize)
	if !ok {
		return false
	}
	finalBuffer = append(finalBuffer, msgBytes...)

	if finalBuffer[0] != 'M' || finalBuffer[1] != '2' {
		output.PrintFrameworkError("Missing m2 in response")

		return false
	}

	return ParseM2Message(finalBuffer, msg)
}