package cc20p1305ssh

import (
	`golang.org/x/crypto/chacha20`
	`golang.org/x/crypto/poly1305`
)

/*
	Decrypt decrypts and authenticates ciphertext returning the decrypted format of ciphertext.

	If tag is nil or empty, it will be assumed that the tag is appended to the end of ciphertext.

	If tag is specified but is <TagSize, an error ErrInvalidTagSize will be returned.
*/
func (c *ChaCha20Poly1305OpenSSH) Decrypt(ciphertext, tag []byte) (decrypted []byte, err error) {

	var tagIdx int
	var firstBlock []byte
	var polyKey [PolyKeySize]byte
	var realTag [TagSize]byte
	var cc20 *chacha20.Cipher

	if tag == nil || len(tag) == 0 {
		tagIdx = len(ciphertext) - TagSize
		tag = ciphertext[tagIdx:]
		ciphertext = ciphertext[:tagIdx]
	}
	if tag == nil || len(tag) != TagSize {
		err = ErrInvalidTagSize
		return
	}
	copy(realTag[:], tag)

	// We need the crypter.
	if cc20, err = chacha20.NewUnauthenticatedCipher(c.realKey[:], iv); err != nil {
		return
	}

	// First we need the poly1305 key. This also sets the counter to 1.
	firstBlock = make([]byte, len(initBlock))
	cc20.XORKeyStream(firstBlock, initBlock)
	copy(polyKey[:], firstBlock[:PolyKeySize])
	// We explicitly set the counter to 1, just in case.
	cc20.SetCounter(1)

	// And verify against the tag.
	if !poly1305.Verify(&realTag, ciphertext, &polyKey) {
		err = ErrInvalidTag
		return
	}

	// Finally, decrypt.
	decrypted = make([]byte, len(ciphertext))
	cc20.XORKeyStream(decrypted, ciphertext)

	return
}

/*
	Encrypt encrypts and authenticates plaintext returning the encrypted format of plaintext.

	If polyTag is nil or <TagSize bytes in length, an error ErrInvalidTagSize will be returned.
*/
func (c *ChaCha20Poly1305OpenSSH) Encrypt(plaintext []byte) (encrypted, tag []byte, err error) {

	var polyKey [PolyKeySize]byte
	var firstBlock []byte
	var cc20 *chacha20.Cipher
	var poly *poly1305.MAC
	var tagTmp []byte

	// TODO: check size of plaintext?
	/*
		if uint64(len(plaintext)) > (1<<38)-64 {
		                panic("chacha20poly1305: plaintext too large")
		        }
	*/

	// We need the crypter.
	if cc20, err = chacha20.NewUnauthenticatedCipher(c.realKey[:], iv); err != nil {
		return
	}

	// First we need the poly1305 key. This also sets the counter to 1.
	firstBlock = make([]byte, len(initBlock))
	cc20.XORKeyStream(firstBlock, initBlock)
	copy(polyKey[:], firstBlock[:PolyKeySize])
	// We explicitly set the counter to 1, just in case.
	cc20.SetCounter(1)
	encrypted = make([]byte, len(plaintext))
	cc20.XORKeyStream(encrypted, plaintext)

	poly = poly1305.New(&polyKey)
	poly.Write(encrypted)
	tagTmp = poly.Sum(nil)
	tag = make([]byte, TagSize)
	copy(tag, tagTmp[:TagSize])

	return
}