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

import (
	"crypto/hkdf"
	"crypto/hmac"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"io"
)

const (
	defaultIterations = 4718580 // from cmd/finditers

	oneSaltSize = 64
	hashSize    = 64
	certSize    = sha512.Size
)

// 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 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
}

func (k HMACKey) mac(out, msg []byte, info string) []byte {
	mac := hmac.New(sha512.New, k)
	buf := make([]byte, 0, binary.MaxVarintLen64)
	mac.Write(binary.AppendUvarint(buf, uint64(len(info))))
	io.WriteString(mac, info)
	mac.Write(msg)
	return mac.Sum(out)
}

// Sign generates an HMAC-SHA512 signature and appends it to msg. The info and
// additionalData are also authenticated, but are not included in the returned
// signed message.
func (k HMACKey) Sign(msg []byte, info string) SignedMessage {
	return k.mac(msg, msg, info)
}

// Verify checks whether the given message has a valid signature, and returns
// the raw message if it does. The additionalData must match the data passed
// to Sign.
func (k HMACKey) Verify(msg SignedMessage, info string) ([]byte, bool) {
	if len(msg) < certSize {
		return nil, false
	}
	msg, sig := msg[:len(msg)-certSize], msg[len(msg)-certSize:]
	if !hmac.Equal(sig, k.mac(nil, msg, info)) {
		return nil, false
	}
	return msg, true
}

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