package net

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"lux/crypto"
	"lux/proto"
	"net"
)

const LUX_PROTO_PACKET_HDRLEN = proto.LUX_PROTO_ID_SIZE + 2

type LuxPacket struct {
	Target proto.LuxID
	Type   uint
	Buffer proto.LuxBuffer
}

func DecryptLuxPacket(dgram LuxDatagram, key crypto.LuxKey) (LuxPacket, error) {
	packet := LuxPacket{}

	cipherBlock, err := aes.NewCipher(key.Key)
	if err != nil {
		return packet, err
	}
	decrypter := cipher.NewCBCDecrypter(cipherBlock, key.IV)

	payloadLen := len(dgram.Payload)
	if payloadLen == 0 || payloadLen%cipherBlock.BlockSize() != 0 {
		return packet, errors.New("payload is not block aligned")
	}

	payload := make([]byte, payloadLen)
	decrypter.CryptBlocks(payload, dgram.Payload)

	packet.Buffer = proto.FromSlice(payload)
	if err := packet.Target.Read(&packet.Buffer); err != nil {
		return packet, err

	}

	if pType, err := packet.Buffer.ReadUint16(); err == nil {
		packet.Type = uint(pType)
	} else {
		return packet, err
	}

	return packet, nil
}

func EncryptLuxPacket(packet LuxPacket, key crypto.LuxKey, target *net.UDPAddr) (LuxDatagram, error) {
	dgram := LuxDatagram{}

	cipherBlock, err := aes.NewCipher(key.Key)
	if err != nil {
		return dgram, err
	}
	encrypter := cipher.NewCBCEncrypter(cipherBlock, key.IV)

	var paddingLen int
	packetLen := packet.Buffer.Length()
	if packetLen%encrypter.BlockSize() != 0 {
		paddingLen = ((packetLen / encrypter.BlockSize()) + 1) * encrypter.BlockSize()
	} else {
		paddingLen = 0
	}

	padding := make([]byte, paddingLen)
	_, err = rand.Read(padding)
	if err != nil {
		return dgram, err
	}

	wd := proto.AllocLuxBuffer(LUX_PROTO_PACKET_HDRLEN + packetLen + paddingLen)
	key.Id.Write(&wd)
	wd.WriteUint16(uint16(packet.Type))
	wd.WriteBytes(packet.Buffer.AllBytes())
	wd.WriteBytes(padding)

	dgram.Target = target
	dgram.Payload = make([]byte, wd.Length())

	encrypter.CryptBlocks(dgram.Payload, wd.AllBytes())
	return dgram, nil
}