// 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"
)

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
}