package kdbx

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/binary"
	"errors"
	"fmt"
)

// maxBlock is the payload block size written into the HMAC block stream.
const maxBlock = 1 << 20

// keys are everything derived from the credentials and the file's seeds.
type keys struct {
	payload []byte // 32 bytes, the cipher key for the payload
	hmac    []byte // 64 bytes, the base from which per-block keys come
}

// compositeKey is SHA-256 over the concatenated credential keys. With a
// password alone that is SHA-256(SHA-256(password)).
func compositeKey(password string) []byte {
	pw := sha256.Sum256([]byte(password))
	comp := sha256.Sum256(pw[:])
	return comp[:]
}

// deriveKeys runs the file's key derivation function and produces the
// payload and HMAC keys.
func deriveKeys(masterSeed []byte, kdf variantDict, password string) (keys, error) {
	transformed, err := transformKey(kdf, compositeKey(password))
	if err != nil {
		return keys{}, err
	}
	pk := sha256.Sum256(append(append([]byte{}, masterSeed...), transformed...))
	hk := sha512.Sum512(append(append(append([]byte{}, masterSeed...), transformed...), 0x01))
	return keys{payload: pk[:], hmac: hk[:]}, nil
}

// transformKey applies the file's KDF to the composite key.
func transformKey(kdf variantDict, composite []byte) ([]byte, error) {
	if err := checkKDF(kdf); err != nil {
		return nil, err
	}
	uuid := kdf.bytes("$UUID")
	switch {
	case bytes.Equal(uuid, kdfAES):
		seed := kdf.bytes("S")
		rounds, ok := kdf.uint64("R")
		if len(seed) != 32 || !ok {
			return nil, errors.New("kdbx: AES-KDF parameters are missing or malformed")
		}
		return aesKDF(composite, seed, rounds)
	case bytes.Equal(uuid, kdfArgon2d):
		return argon2KDF(kdf, composite, argon2d)
	case bytes.Equal(uuid, kdfArgon2id):
		return argon2KDF(kdf, composite, argon2id)
	case len(uuid) == 0:
		return nil, errors.New("kdbx: file names no key derivation function")
	default:
		return nil, unsupported("key derivation function %x", uuid)
	}
}

// aesKDF encrypts the composite key with itself, round after round. It is
// KeePass's original work factor and needs only the standard library.
func aesKDF(composite, seed []byte, rounds uint64) ([]byte, error) {
	block, err := aes.NewCipher(seed)
	if err != nil {
		return nil, err
	}
	data := make([]byte, 32)
	copy(data, composite)
	for i := uint64(0); i < rounds; i++ {
		block.Encrypt(data[0:16], data[0:16])
		block.Encrypt(data[16:32], data[16:32])
	}
	sum := sha256.Sum256(data)
	return sum[:], nil
}

// blockKey is the HMAC key for one block of the payload stream.
func blockKey(base []byte, index uint64) []byte {
	buf := make([]byte, 8, 8+len(base))
	binary.LittleEndian.PutUint64(buf, index)
	sum := sha512.Sum512(append(buf, base...))
	return sum[:]
}

// headerHMAC authenticates the outer header. KeePass uses the maximum block
// index for it, so it cannot be confused with a payload block.
func headerHMAC(base, header []byte) []byte {
	mac := hmac.New(sha256.New, blockKey(base, ^uint64(0)))
	mac.Write(header)
	return mac.Sum(nil)
}

// readBlocks verifies and concatenates the HMAC block stream.
func readBlocks(data, base []byte) ([]byte, error) {
	var out []byte
	for index := uint64(0); ; index++ {
		if len(data) < 36 {
			return nil, errors.New("kdbx: truncated payload block header")
		}
		gotMAC := data[:32]
		size := binary.LittleEndian.Uint32(data[32:36])
		if int(size) < 0 || 36+int(size) > len(data) {
			return nil, errors.New("kdbx: truncated payload block")
		}
		body := data[36 : 36+int(size)]

		mac := hmac.New(sha256.New, blockKey(base, index))
		var meta [12]byte
		binary.LittleEndian.PutUint64(meta[0:], index)
		binary.LittleEndian.PutUint32(meta[8:], size)
		mac.Write(meta[:])
		mac.Write(body)
		if subtle.ConstantTimeCompare(mac.Sum(nil), gotMAC) != 1 {
			return nil, fmt.Errorf("kdbx: payload block %d failed authentication", index)
		}
		data = data[36+int(size):]
		if size == 0 {
			return out, nil
		}
		out = append(out, body...)
	}
}

// writeBlocks produces the HMAC block stream, ending with the empty block
// that marks the end of the file.
func writeBlocks(payload, base []byte) []byte {
	out := []byte{}
	index := uint64(0)
	for len(payload) > 0 {
		n := min(len(payload), maxBlock)
		out = appendBlock(out, base, index, payload[:n])
		payload = payload[n:]
		index++
	}
	return appendBlock(out, base, index, nil)
}

func appendBlock(out, base []byte, index uint64, body []byte) []byte {
	var meta [12]byte
	binary.LittleEndian.PutUint64(meta[0:], index)
	binary.LittleEndian.PutUint32(meta[8:], uint32(len(body)))
	mac := hmac.New(sha256.New, blockKey(base, index))
	mac.Write(meta[:])
	mac.Write(body)
	out = append(out, mac.Sum(nil)...)
	out = append(out, meta[8:12]...)
	return append(out, body...)
}

// decryptPayload reverses the file's payload cipher.
func decryptPayload(cipherID, key, iv, data []byte) ([]byte, error) {
	switch {
	case bytes.Equal(cipherID, cipherAES256CBC):
		block, err := aes.NewCipher(key)
		if err != nil {
			return nil, err
		}
		if len(iv) != block.BlockSize() {
			return nil, errors.New("kdbx: encryption IV has the wrong length")
		}
		if len(data) == 0 || len(data)%block.BlockSize() != 0 {
			return nil, ErrWrongPassword
		}
		out := make([]byte, len(data))
		cipher.NewCBCDecrypter(block, iv).CryptBlocks(out, data)
		return unpad(out, block.BlockSize())
	case bytes.Equal(cipherID, cipherChaCha20Outer):
		if len(iv) != 12 {
			return nil, errors.New("kdbx: ChaCha20 nonce has the wrong length")
		}
		out := make([]byte, len(data))
		newChaCha20(key, iv, 0).XORKeyStream(out, data)
		return out, nil
	default:
		return nil, unsupported("payload cipher %x", cipherID)
	}
}

// encryptPayload applies the file's payload cipher.
func encryptPayload(cipherID, key, iv, data []byte) ([]byte, error) {
	switch {
	case bytes.Equal(cipherID, cipherAES256CBC):
		block, err := aes.NewCipher(key)
		if err != nil {
			return nil, err
		}
		padded := pad(data, block.BlockSize())
		out := make([]byte, len(padded))
		cipher.NewCBCEncrypter(block, iv).CryptBlocks(out, padded)
		return out, nil
	case bytes.Equal(cipherID, cipherChaCha20Outer):
		out := make([]byte, len(data))
		newChaCha20(key, iv, 0).XORKeyStream(out, data)
		return out, nil
	default:
		return nil, unsupported("payload cipher %x", cipherID)
	}
}

func pad(b []byte, size int) []byte {
	n := size - len(b)%size
	out := make([]byte, len(b)+n)
	copy(out, b)
	for i := len(b); i < len(out); i++ {
		out[i] = byte(n)
	}
	return out
}

func unpad(b []byte, size int) ([]byte, error) {
	if len(b) == 0 || len(b)%size != 0 {
		return nil, ErrWrongPassword
	}
	n := int(b[len(b)-1])
	if n == 0 || n > size || n > len(b) {
		// Wrong key almost always shows up here first.
		return nil, ErrWrongPassword
	}
	for _, v := range b[len(b)-n:] {
		if int(v) != n {
			return nil, ErrWrongPassword
		}
	}
	return b[:len(b)-n], nil
}