// Package cryptoutil contains friendly wrappers around the algorithms from the
// standard library's crypto package.
package cryptoutil

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/hmac"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"slices"
)

const (
	defaultIterations = 4718580 // from cmd/finditers

	oneSaltSize = 64
	hashSize    = 64
	certSize    = sha512.Size
	aesKeySize  = 32
	nonceSize   = aes.BlockSize
)

// SaltSize is the expected salt length for HashIter.
const SaltSize = 2 * oneSaltSize

// A PasswordHash must be PasswordHashSize bytes.
const PasswordHashSize = SaltSize + hashSize

// A PasswordHash is derived from the user's password and can be passed to
// CheckPassword to verify if two passwords match. A PasswordHash must be
// PasswordHashSize bytes.
type PasswordHash []byte

func (h PasswordHash) String() string { return hex.EncodeToString(h) }

// A RawKey is generated from a password hash and can be used to derive further
// key material.
type RawKey struct {
	key []byte
}

// An HMACKey must be HMACKeySize bytes.
const HMACKeySize = sha512.BlockSize

// An HMAC key can be used to symmetrically sign and verify messages
// using HMAC-SHA512. An HMACKey must be HMACKeyLen bytes.
type HMACKey []byte

// A SignedMessage is a message that has been cryptographically signed
// with HMAC-SHA512.
type SignedMessage []byte

// An EncryptionKey must be EncryptionKeySize bytes.
const EncryptionKeySize = aesKeySize + HMACKeySize

// An EncryptionKey is used for encrypting and decrypting data. An EncryptionKey
// must be EncryptionKeySize bytes.
type EncryptionKey []byte

// An EncryptedMessage is encrypted with AES-256-CTR-HMAC-SHA512.
type EncryptedMessage []byte

// HashIter runs PBKDF2-SHA512 for iter iterations. Useful for benchmarking. The
// salt must be SaltSize bytes.
func HashIter(password string, salt []byte, iter int) ([]byte, error) {
	return pbkdf2.Key(sha512.New, password, salt, iter, sha512.Size)
}

func hashWithSalt(password string, salt []byte) (RawKey, []byte, error) {
	hash, err := HashIter(password, salt[:oneSaltSize], defaultIterations)
	if err != nil {
		return RawKey{}, nil, err
	}
	key, err := hkdf.Extract(sha512.New, hash, salt[oneSaltSize:])
	if err != nil {
		return RawKey{}, nil, err
	}
	pwHash, err := hkdf.Expand(sha512.New, key, "pwhash", hashSize)
	if err != nil {
		return RawKey{}, nil, err
	}
	return RawKey{key}, pwHash, nil
}

// Hash hashes a user password using PBKDF2-SHA512.
func Hash(password string) (PasswordHash, error) {
	salt := make([]byte, SaltSize, SaltSize+hashSize)
	rand.Read(salt)
	_, pwHash, err := hashWithSalt(password, salt)
	if err != nil {
		return nil, err
	}
	return append(salt, pwHash...), nil
}

// CheckPassword verifies password against a PasswordHash and returns an
// EncryptionKey derived from the password if successful.
func (h PasswordHash) CheckPassword(password string) (RawKey, error) {
	salt, expectedHash := h[:SaltSize], h[SaltSize:]
	key, pwHash, err := hashWithSalt(password, salt)
	if err != nil {
		return RawKey{}, err
	}
	if subtle.ConstantTimeCompare(pwHash, expectedHash) == 0 {
		return RawKey{}, errors.New("incorrect password")
	}
	return key, nil
}

// Sign generates an HMAC-SHA512 signature and appends it to msg.
func (k HMACKey) Sign(msg []byte) SignedMessage {
	mac := hmac.New(sha512.New, k)
	mac.Write(msg)
	return mac.Sum(msg)
}

// Verify checks whether the given message has a valid signature, and returns
// the raw message if it does.
func (k HMACKey) Verify(msg SignedMessage) ([]byte, bool) {
	if len(msg) < certSize {
		return nil, false
	}
	msg, sig := msg[:len(msg)-certSize], msg[len(msg)-certSize:]
	mac := hmac.New(sha512.New, k)
	mac.Write(msg)
	if !hmac.Equal(sig, mac.Sum(nil)) {
		return nil, false
	}
	return msg, true
}

// EncryptionKey derives an EncryptionKey.
func (k RawKey) EncryptionKey() (EncryptionKey, error) {
	return hkdf.Expand(sha512.New, k.key, "encrypt", EncryptionKeySize)
}

// Encrypt encrypts a message with AES-256-CTR-HMAC-SHA512.
func (k EncryptionKey) Encrypt(msg []byte) (EncryptedMessage, error) {
	aesKey, hmacKey := k[:aesKeySize], HMACKey(k[aesKeySize:])
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	cipherText := make([]byte, len(msg)+nonceSize, len(msg)+nonceSize+certSize)
	nonce := cipherText[len(msg) : len(msg)+nonceSize]
	rand.Read(nonce)
	cipher.NewCTR(block, nonce).XORKeyStream(cipherText, msg)
	return EncryptedMessage(hmacKey.Sign(cipherText)), nil
}

// Decrypt verifies and decrypts an encrypted message. It overwrites msg with
// the resulting plain text.
func (k EncryptionKey) Decrypt(msg EncryptedMessage) ([]byte, error) {
	aesKey, hmacKey := k[:aesKeySize], HMACKey(k[aesKeySize:])
	msg, ok := hmacKey.Verify(SignedMessage(msg))
	if !ok {
		return nil, errors.New("bad signature")
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	msg, nonce := msg[:len(msg)-nonceSize], msg[len(msg)-nonceSize:]
	cipher.NewCTR(block, nonce).XORKeyStream(msg, msg)
	return msg, nil
}

// An EncryptingWriter encrypts the output and writes it to the underlying
// writer. It's very important to call .Flush() to write the MAC after the data
// has been written.
type EncryptingWriter struct {
	w      io.Writer
	stream cipher.Stream
	mac    hash.Hash
	buf    []byte
}

// Writer returns a new writer that encrypts its output. Don't forget to call
// .Flush() to write the MAC.
func (k EncryptionKey) Writer(w io.Writer) (*EncryptingWriter, error) {
	aesKey, hmacKey := k[:aesKeySize], k[aesKeySize:]
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	mac := hmac.New(sha512.New, hmacKey)
	nonce := make([]byte, nonceSize)
	rand.Read(nonce)
	mac.Write(nonce)
	if _, err := w.Write(nonce); err != nil {
		return nil, err
	}
	return &EncryptingWriter{
		w:      w,
		stream: cipher.NewCTR(block, nonce),
		mac:    mac,
	}, nil
}

// Write encrypts buf and writes it to the underlying writer.
func (w *EncryptingWriter) Write(buf []byte) (int, error) {
	if len(w.buf) < len(buf) {
		w.buf = slices.Grow(w.buf, len(buf)-len(w.buf))
	}
	w.buf = w.buf[:len(buf)]
	// Encrypt-then-MAC
	w.stream.XORKeyStream(w.buf, buf)
	w.mac.Write(w.buf)
	return w.w.Write(w.buf)
}

// Flush writes the MAC for the encrypted message. Flush must be called after
// all data has been written.
func (w *EncryptingWriter) Flush() error {
	_, err := w.w.Write(w.mac.Sum(nil))
	return err
}

// A DecryptingReader decrypts data from an underlying reader.
type DecryptingReader struct {
	r cipher.StreamReader
}

// Reader returns a new DecryptingReader. The data is processed in two passes,
// first to verify the MAC, then the io.Seeker interface is used to reset the
// reader for decryption.
func (k EncryptionKey) Reader(r io.ReadSeeker) (*DecryptingReader, error) {
	aesKey, hmacKey := k[:aesKeySize], k[aesKeySize:]
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	totalLen, err := r.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, err
	}
	if totalLen < certSize {
		return nil, errors.New("file too short")
	}
	dataLen := totalLen - certSize
	if _, err := r.Seek(0, io.SeekStart); err != nil {
		return nil, err
	}
	mac := hmac.New(sha512.New, hmacKey)
	if _, err := io.Copy(mac, &io.LimitedReader{R: r, N: dataLen}); err != nil {
		return nil, err
	}
	expectedMAC := mac.Sum(nil)
	sig := make([]byte, certSize)
	if _, err := io.ReadFull(r, sig); err != nil {
		return nil, err
	}
	if !hmac.Equal(sig, expectedMAC) {
		return nil, fmt.Errorf("invalid mac (got %x, want %x)", sig, expectedMAC)
	}
	if _, err := r.Seek(0, io.SeekStart); err != nil {
		return nil, err
	}
	nonce := make([]byte, nonceSize)
	if _, err := io.ReadFull(r, nonce); err != nil {
		return nil, err
	}
	return &DecryptingReader{cipher.StreamReader{S: cipher.NewCTR(block, nonce), R: &io.LimitedReader{R: r, N: dataLen - nonceSize}}}, nil
}

// Read reads and decrypts data from the underlying reader.
func (r *DecryptingReader) Read(buf []byte) (int, error) {
	return r.r.Read(buf)
}