package station

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/subtle"
	"encoding/binary"
	"errors"
)

func counterAAD(counter uint32) []byte {
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], counter)
	return b[:]
}

func registrationAAD(clientID uint16) []byte {
	b := counterAAD(0)
	return appendU16(b, clientID)
}

func ivFor(prefix []byte, clientID uint16, counter uint32) []byte {
	iv := make([]byte, 12)
	copy(iv, prefix)
	binary.BigEndian.PutUint16(iv[6:8], clientID)
	binary.BigEndian.PutUint32(iv[8:12], counter)
	return iv
}

func sealDatagram(key, prefix []byte, clientID uint16, counter uint32, aad, plain []byte) ([]byte, error) {
	tag, ct, err := sealTrunc(key, ivFor(prefix, clientID, counter), plain, aad)
	if err != nil {
		return nil, err
	}
	wire := make([]byte, 0, len(aad)+len(tag)+len(ct))
	wire = append(wire, aad...)
	wire = append(wire, tag...)
	return append(wire, ct...), nil
}

func openDatagram(key, prefix []byte, clientID uint16, counter uint32, aad, tag, ct []byte) ([]byte, error) {
	return openTrunc(key, ivFor(prefix, clientID, counter), tag, ct, aad)
}

// The protocol's tag is 8 bytes. crypto/cipher only builds GCM tags of 12 to 16
// bytes, and an 8-byte tag is the leading 8 bytes of the ordinary 16-byte tag.
func sealTrunc(key, nonce, plain, aad []byte) (tag, ct []byte, err error) {
	g, err := gcmFull(key)
	if err != nil {
		return nil, nil, err
	}
	out := g.Seal(nil, nonce, plain, aad)
	ct = append([]byte(nil), out[:len(plain)]...)
	tag = append([]byte(nil), out[len(plain):len(plain)+8]...)
	return tag, ct, nil
}

func openTrunc(key, nonce, tag, ct, aad []byte) ([]byte, error) {
	if len(tag) != 8 {
		return nil, errors.New("tag")
	}
	g, err := gcmFull(key)
	if err != nil {
		return nil, err
	}
	stream := g.Seal(nil, nonce, make([]byte, len(ct)), aad)
	plain := make([]byte, len(ct))
	subtle.XORBytes(plain, ct, stream[:len(ct)])
	check := g.Seal(nil, nonce, plain, aad)
	var got, want [8]byte
	copy(got[:], tag)
	copy(want[:], check[len(plain):len(plain)+8])
	if subtle.ConstantTimeCompare(got[:], want[:]) != 1 {
		return nil, errors.New("tag")
	}
	return plain, nil
}

func gcmFull(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}