package kdbx

import "encoding/binary"

// ChaCha20 (RFC 8439) is the inner stream cipher KDBX 4 uses to protect
// individual values inside the XML, and an optional outer cipher. The Go
// standard library does not export it, so it is implemented here: the
// algorithm is 60 lines and comes with published test vectors.
type chacha20 struct {
	state [16]uint32
	block [64]byte
	used  int // bytes of block already consumed
}

// newChaCha20 starts a keystream for a 32-byte key and 12-byte nonce.
func newChaCha20(key, nonce []byte, counter uint32) *chacha20 {
	c := &chacha20{used: 64}
	c.state[0], c.state[1], c.state[2], c.state[3] =
		0x61707865, 0x3320646e, 0x79622d32, 0x6b206574 // "expand 32-byte k"
	for i := 0; i < 8; i++ {
		c.state[4+i] = binary.LittleEndian.Uint32(key[i*4:])
	}
	c.state[12] = counter
	for i := 0; i < 3; i++ {
		c.state[13+i] = binary.LittleEndian.Uint32(nonce[i*4:])
	}
	return c
}

// XORKeyStream applies the keystream to src, continuing where the previous
// call stopped. KDBX protects its values as one continuous stream, so this
// must survive being called once per value.
func (c *chacha20) XORKeyStream(dst, src []byte) {
	for i := range src {
		if c.used == 64 {
			c.nextBlock()
		}
		dst[i] = src[i] ^ c.block[c.used]
		c.used++
	}
}

func (c *chacha20) nextBlock() {
	var x [16]uint32
	copy(x[:], c.state[:])
	for i := 0; i < 10; i++ { // 20 rounds, two per iteration
		quarter(&x, 0, 4, 8, 12)
		quarter(&x, 1, 5, 9, 13)
		quarter(&x, 2, 6, 10, 14)
		quarter(&x, 3, 7, 11, 15)
		quarter(&x, 0, 5, 10, 15)
		quarter(&x, 1, 6, 11, 12)
		quarter(&x, 2, 7, 8, 13)
		quarter(&x, 3, 4, 9, 14)
	}
	for i := range x {
		binary.LittleEndian.PutUint32(c.block[i*4:], x[i]+c.state[i])
	}
	c.state[12]++
	c.used = 0
}

func quarter(x *[16]uint32, a, b, cc, d int) {
	x[a] += x[b]
	x[d] = rotl(x[d]^x[a], 16)
	x[cc] += x[d]
	x[b] = rotl(x[b]^x[cc], 12)
	x[a] += x[b]
	x[d] = rotl(x[d]^x[a], 8)
	x[cc] += x[d]
	x[b] = rotl(x[b]^x[cc], 7)
}

func rotl(v uint32, n uint) uint32 { return v<<n | v>>(32-n) }