// `msg.go` contains the logic for building, reading, accessing, and
// serializing RouterOS M2 messages.
package mikrotik

import (
	"encoding/binary"

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

type M2Message struct {
	Bools    map[uint32]bool
	U32s     map[uint32]uint32
	Strings  map[uint32][]byte
	Raw      map[uint32][]byte
	ArrayU32 map[uint32][]uint32
	ArrayM2  map[uint32][]*M2Message
}

func NewM2Message() *M2Message {
	var m2 M2Message
	m2.Bools = make(map[uint32]bool)
	m2.U32s = make(map[uint32]uint32)
	m2.Strings = make(map[uint32][]byte)
	m2.Raw = make(map[uint32][]byte)
	m2.ArrayU32 = make(map[uint32][]uint32)
	m2.ArrayM2 = make(map[uint32][]*M2Message)

	return &m2
}

const (
	typeBoolean      uint32 = 0
	typeShortLength  uint32 = 0x01000000
	typeUint32       uint32 = 0x08000000
	typeString       uint32 = 0x20000000
	typeRaw          uint32 = 0x30000000
	typeUint32Array  uint32 = 0x88000000
	typeMessageArray uint32 = 0xa8000000
)

const (
	VarSysTo         uint32 = 0x00ff0001
	VarFrom          uint32 = 0x00ff0002
	VarReplyExpected uint32 = 0x00ff0005
	VarRequestID     uint32 = 0x00ff0006
	VarCommand       uint32 = 0x00ff0007
	VarErrorCode     uint32 = 0x00ff0008
	VarErrorString   uint32 = 0x00ff0009
	VarSessionID     uint32 = 0x00fe0001
)

func (msg M2Message) SetTo(to uint32, handler uint32) {
	uint32Slice := make([]uint32, 2)
	uint32Slice[0] = to
	uint32Slice[1] = handler
	msg.AddU32Array(VarSysTo, uint32Slice)
}

func (msg M2Message) SetCommand(command uint32) {
	msg.AddU32(VarCommand, command)
}

func (msg M2Message) SetRequestID(id uint32) {
	msg.AddU32(VarRequestID, id)
}

func (msg M2Message) SetReplyExpected(expected bool) {
	msg.AddBool(VarReplyExpected, expected)
}

func (msg M2Message) SetSessionID(id uint32) {
	msg.AddU32(VarSessionID, id)
}

func (msg M2Message) GetSessionID() uint32 {
	return msg.U32s[VarSessionID]
}

func (msg M2Message) AddBool(varname uint32, data bool) {
	msg.Bools[varname] = data
}

func (msg M2Message) AddU32(varname uint32, data uint32) {
	msg.U32s[varname] = data
}

func (msg M2Message) AddString(varname uint32, data []byte) {
	msg.Strings[varname] = make([]byte, len(data))
	copy(msg.Strings[varname], data)
}

func (msg M2Message) AddRaw(varname uint32, data []byte) {
	msg.Raw[varname] = make([]byte, len(data))
	copy(msg.Raw[varname], data)
}

func (msg M2Message) AddU32Array(varname uint32, data []uint32) {
	msg.ArrayU32[varname] = append(msg.ArrayU32[varname], data...)
}

func (msg M2Message) Serialize() []byte {
	serialized := []byte{}

	for varname, value := range msg.Bools {
		binaryBool := make([]byte, 4)
		if value {
			varname |= typeShortLength
		}
		binary.LittleEndian.PutUint32(binaryBool, varname)
		serialized = append(serialized, binaryBool...)
	}

	for varname, value := range msg.U32s {
		varname |= typeUint32
		binaryUint32 := make([]byte, 4)

		if value > 255 {
			binary.LittleEndian.PutUint32(binaryUint32, varname)
			serialized = append(serialized, binaryUint32...)
			binary.LittleEndian.PutUint32(binaryUint32, value)
			serialized = append(serialized, binaryUint32...)
		} else {
			varname |= typeShortLength
			binary.LittleEndian.PutUint32(binaryUint32, varname)
			serialized = append(serialized, binaryUint32...)
			serialized = append(serialized, byte(value&0x000000ff))
		}
	}

	for varname, value := range msg.Strings {
		varname |= typeString
		binaryString := make([]byte, 4)

		if len(value) > 255 {
			binary.LittleEndian.PutUint32(binaryString, varname)
			serialized = append(serialized, binaryString...)
			binaryLength := make([]byte, 2)
			binary.LittleEndian.PutUint16(binaryLength, uint16(len(value)))
			serialized = append(serialized, binaryLength...)
		} else {
			varname |= typeShortLength
			binary.LittleEndian.PutUint32(binaryString, varname)
			serialized = append(serialized, binaryString...)
			serialized = append(serialized, byte(len(value)))
		}
		serialized = append(serialized, value...)
	}

	for varname, value := range msg.Raw {
		varname |= typeRaw
		binaryString := make([]byte, 4)

		if len(value) > 255 {
			binary.LittleEndian.PutUint32(binaryString, varname)
			serialized = append(serialized, binaryString...)
			binaryLength := make([]byte, 2)
			binary.LittleEndian.PutUint16(binaryLength, uint16(len(value)))
			serialized = append(serialized, binaryLength...)
		} else {
			varname |= typeShortLength
			binary.LittleEndian.PutUint32(binaryString, varname)
			serialized = append(serialized, binaryString...)
			serialized = append(serialized, byte(len(value)))
		}
		serialized = append(serialized, value...)
	}

	for varname, value := range msg.ArrayU32 {
		varname |= typeUint32Array
		binaryArray := make([]byte, 4)
		binary.LittleEndian.PutUint32(binaryArray, varname)
		serialized = append(serialized, binaryArray...)

		binaryLength := make([]byte, 2)
		binary.LittleEndian.PutUint16(binaryLength, uint16(len(value)))
		serialized = append(serialized, binaryLength...)

		for _, entry := range value {
			binaryEntry := make([]byte, 4)
			binary.LittleEndian.PutUint32(binaryEntry, entry)
			serialized = append(serialized, binaryEntry...)
		}
	}

	// Skipping array of m2. I'm not sure I ever used this before
	// and doing it correctly takes some thought

	return serialized
}

func handleStringorRaw(varTypeName uint32, varName uint32, data *[]byte, storage *map[uint32][]byte) bool {
	if len(*data) <= 2 {
		return false
	}
	length := int(binary.LittleEndian.Uint16(*data))
	if (varTypeName & typeShortLength) != 0 {
		length = int((*data)[0])
		*data = (*data)[1:]
	} else {
		*data = (*data)[2:]
	}
	if len(*data) < length {
		return false
	}

	(*storage)[varName] = make([]byte, length)
	copy((*storage)[varName], *data)
	*data = (*data)[length:]

	return true
}

// this switch is cursed. it needs to be broke into manageable functions.
func ParseM2Message(data []byte, msg *M2Message) bool {
	if len(data) < 4 {
		return false
	}

	if data[2] == 'M' && data[3] == '2' {
		// the message was passed in with the length + M2 header. Truncate
		// and skip
		m2length := int(data[0])
		data = data[4:]
		data = data[:m2length]
	}

	for len(data) > 4 {
		varTypeName := binary.LittleEndian.Uint32(data)
		varType := varTypeName & 0xf8000000
		varName := varTypeName & 0x00ffffff
		data = data[4:]

		switch varType {
		case typeBoolean:
			msg.Bools[varName] = (varTypeName & typeShortLength) != 0
		case typeUint32:
			if (varTypeName & typeShortLength) != 0 {
				if len(data) == 0 {
					return false
				}
				msg.U32s[varName] = uint32(data[0])
				data = data[1:]
			} else {
				if len(data) < 4 {
					return false
				}
				msg.U32s[varName] = binary.LittleEndian.Uint32(data)
				data = data[4:]
			}
		case typeString:
			ok := handleStringorRaw(varTypeName, varName, &data, &msg.Strings)
			if !ok {
				return false
			}
		case typeRaw:
			ok := handleStringorRaw(varTypeName, varName, &data, &msg.Raw)
			if !ok {
				return false
			}
		case typeUint32Array:
			if len(data) <= 2 {
				return false
			}
			arrayEntries := int(binary.LittleEndian.Uint16(data))
			data = data[2:]
			if len(data) < (arrayEntries * 4) {
				return false
			}

			for i := 0; i < arrayEntries; i++ {
				msg.ArrayU32[varName] = append(msg.ArrayU32[varName], binary.LittleEndian.Uint32(data))
				data = data[4:]
			}
		case typeMessageArray:
			if len(data) <= 2 {
				return false
			}
			length := int(binary.LittleEndian.Uint16(data))
			if (varTypeName & typeShortLength) != 0 {
				length = int(data[0])
				data = data[1:]
			} else {
				data = data[2:]
			}
			subdata := make([]byte, length)
			copy(subdata, data)
			subMsg := NewM2Message()
			msg.ArrayM2[varName] = append(msg.ArrayM2[varName], subMsg)

			// recursion... technically not safe. Not checking the return here either is
			// a bold choice. Mostly avoiding handling empty entries.
			_ = ParseM2Message(subdata, subMsg)

			data = data[length:]
		default:
			output.PrintfFrameworkError("Unhandled %x named %x\n", varType, varName)
		}
	}

	return true
}