package cryptoutil

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"errors"
	"io"
)

// Online Authenticated Encryption from https://eprint.iacr.org/2015/189.pdf

const (
	aeadOverhead         = 16
	aesKeySize           = 32
	noncePrefixSize      = 3
	gcmNonceSize         = 12
	headerSize           = aesKeySize + noncePrefixSize
	cacheSize            = 192 * 1024
	encryptedSegmentSize = cacheSize - 1
	segmentSize          = encryptedSegmentSize - aeadOverhead
)

// An EncryptionKey is used for encrypting and decrypting data. A key can be any
// length, but at least 64 bytes would be recommended.
type EncryptionKey []byte

func (k EncryptionKey) deriveAESKey(salt []byte) (cipher.AEAD, error) {
	derivedKey, err := hkdf.Key(sha512.New, k, salt, "", aesKeySize)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(derivedKey)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

func makeNonce(nonce []byte, noncePrefix []byte, i *uint64) error {
	copy(nonce, noncePrefix)
	binary.BigEndian.PutUint64(nonce[noncePrefixSize:], *i)
	*i++
	if *i == 0 {
		return errors.New("counter overflowed (64 bits??)")
	}
	return nil
}

// An EncryptingWriter encrypts data in segments using the STREAM construction
// described in https://eprint.iacr.org/2015/189.pdf. The writer buffers data up
// to the segment size, so it's important to call Close to flush the
// final segment.
type EncryptingWriter struct {
	w              io.Writer
	key            EncryptionKey
	aead           cipher.AEAD
	additionalData []byte
	noncePrefix    [noncePrefixSize]byte
	i              uint64
	initialized    bool
	err            error
	bufN           int
	buf            [encryptedSegmentSize]byte
}

// NewWriter returns a new EncryptingWriter that writes to w. The additionalData
// will be authenticated with the first segment, but is not written to w. The
// same additional data must be provided when decrypting.
func (k EncryptionKey) NewWriter(w io.Writer, additionalData []byte) *EncryptingWriter {
	return &EncryptingWriter{
		w:              w,
		key:            k,
		additionalData: additionalData,
		i:              1,
	}
}

func (w *EncryptingWriter) initialize() error {
	w.initialized = true
	header := make([]byte, headerSize)
	rand.Read(header)
	copy(w.noncePrefix[:], header[aesKeySize:])
	if w.aead, w.err = w.key.deriveAESKey(header[:aesKeySize]); w.err != nil {
		return w.err
	}
	_, w.err = w.w.Write(header)
	return w.err
}

func (w *EncryptingWriter) nonce(nonce []byte) error {
	w.err = makeNonce(nonce, w.noncePrefix[:], &w.i)
	return w.err
}

func (w *EncryptingWriter) writeBuf() error {
	nonce := make([]byte, gcmNonceSize)
	if err := w.nonce(nonce); err != nil {
		return err
	}
	if _, w.err = w.w.Write(w.aead.Seal(w.buf[:0], nonce, w.buf[:w.bufN], w.additionalData)); w.err != nil {
		return w.err
	}
	w.bufN = 0
	w.additionalData = nil
	return nil
}

func (w *EncryptingWriter) Write(buf []byte) (int, error) {
	if !w.initialized {
		if err := w.initialize(); err != nil {
			return 0, err
		}
	}
	if w.err != nil {
		return 0, w.err
	}
	nn := 0
	for len(buf) > 0 {
		if w.bufN == segmentSize {
			if err := w.writeBuf(); err != nil {
				return nn, err
			}
		}
		n := copy(w.buf[w.bufN:segmentSize], buf)
		w.bufN += n
		nn += n
		buf = buf[n:]
	}
	return nn, nil
}

var errClosed = errors.New("closed")

// Close encrypts writes the final segment. It does not close the
// underlying writer.
func (w *EncryptingWriter) Close() error {
	if !w.initialized {
		if err := w.initialize(); err != nil {
			return err
		}
	}
	if w.err == errClosed {
		return nil
	}
	if w.err != nil {
		return w.err
	}
	if w.bufN > 0 {
		nonce := make([]byte, gcmNonceSize)
		if err := w.nonce(nonce); err != nil {
			return err
		}
		nonce[gcmNonceSize-1] = 1
		if _, w.err = w.w.Write(w.aead.Seal(w.buf[:0], nonce, w.buf[:w.bufN], w.additionalData)); w.err != nil {
			return w.err
		}
	}
	w.err = errClosed
	return nil
}

// A DecryptingReader decrypts data using the STREAM construction.
type DecryptingReader struct {
	r              io.Reader
	key            EncryptionKey
	aead           cipher.AEAD
	additionalData []byte
	noncePrefix    [noncePrefixSize]byte
	i              uint64
	initialized    bool
	err            error
	bufRead, bufN  int
	peekedByte     bool
	buf            [encryptedSegmentSize + 1]byte
}

// NewReader returns a new DecryptingWriter that decrypts data from r. The
// additionaData must be the same that was provided when encrypting.
func (k EncryptionKey) NewReader(r io.Reader, additionalData []byte) *DecryptingReader {
	return &DecryptingReader{
		r:              r,
		key:            k,
		additionalData: additionalData,
		i:              1,
	}
}

func (r *DecryptingReader) initialize() error {
	r.initialized = true
	header := make([]byte, headerSize)
	if _, r.err = io.ReadFull(r.r, header); r.err != nil {
		return r.err
	}
	copy(r.noncePrefix[:], header[aesKeySize:])
	r.aead, r.err = r.key.deriveAESKey(header[:aesKeySize])
	return r.err
}

func (r *DecryptingReader) nonce(nonce []byte) error {
	if err := makeNonce(nonce, r.noncePrefix[:], &r.i); err != nil {
		r.err = err
		return err
	}
	return nil
}

func (r *DecryptingReader) fillBuf() error {
	n := 0
	if r.peekedByte {
		n = 1
		r.buf[0] = r.buf[encryptedSegmentSize]
		r.peekedByte = false
	}
	for n < len(r.buf) && r.err == nil {
		var m int
		m, r.err = r.r.Read(r.buf[n:])
		n += m
	}
	if n == 0 {
		return r.err
	}
	if n == encryptedSegmentSize+1 {
		r.peekedByte = true
		n = encryptedSegmentSize
	}
	nonce := make([]byte, gcmNonceSize)
	if err := r.nonce(nonce); err != nil {
		return err
	}
	if r.err == io.EOF {
		nonce[gcmNonceSize-1] = 1
	}
	result, err := r.aead.Open(r.buf[:0], nonce, r.buf[:n], r.additionalData)
	if err != nil {
		r.err = err
		return err
	}
	r.bufRead = 0
	r.bufN = len(result)
	r.additionalData = nil
	return nil
}

func (r *DecryptingReader) Read(buf []byte) (int, error) {
	if !r.initialized {
		if err := r.initialize(); err != nil {
			return 0, err
		}
	}
	if r.bufRead == r.bufN {
		if err := r.fillBuf(); err != nil {
			return 0, err
		}
	}
	n := copy(buf, r.buf[r.bufRead:r.bufN])
	r.bufRead += n
	if r.bufRead == r.bufN {
		return n, r.err
	}
	return n, nil
}