// Package kmac implements the KMAC128 and KMAC256 keyed hash functions based
// on sha3.
//
// The KECCAK Message Authentication Code (KMAC) algorithm is a PRF and keyed
// hash function based on KECCAK. It provides variable-length output, and unlike
// SHAKE and cSHAKE, altering the requested output length generates a new,
// unrelated output. KMAC has two variants, KMAC128 and KMAC256, built from
// cSHAKE128 and cSHAKE256, respectively. The two variants differ somewhat in
// their technical security properties. Nonetheless, for most applications, both
// variants can support any security strength up to 256 bits of security,
// provided that the length of the input key is not less than the required
// security strength.
//
// It's possible to use KMAC as a key derivation function. For the input data,
// use a string containing information related to the derived keying material.
// It may include the identities of the parties who are deriving and/or using
// the derived keying material and, optionally, a nonce known by the parties who
// derive the keys. Use the string "KDF" for optionalCustomizationString.
//
// # Disclaimer
//
// This package has not been reviewed by security experts. Consider using
// crypto/hmac instead.
package kmac

import (
	"crypto/sha3"
	"encoding/binary"
	"math/bits"
)

// from https://cs.opensource.google/go/go/+/master:src/crypto/internal/fips140/sha3/shake.go;l=36;drc=4a3cef2036097d323b6cc0bbe90fc4d8c7588660
func leftEncode(x uint64) []byte {
	// Let n be the smallest positive integer for which 2^(8n) > x.
	n := (bits.Len64(x) + 7) / 8
	if n == 0 {
		n = 1
	}
	// Return n || x with n as a byte and x an n bytes in big-endian order.
	b := make([]byte, 9)
	binary.BigEndian.PutUint64(b[1:], x)
	b = b[9-n-1:]
	b[0] = byte(n)
	return b
}

func rightEncode(x uint64) []byte {
	// Let n be the smallest positive integer for which 2^(8n) > x.
	n := (bits.Len64(x) + 7) / 8
	if n == 0 {
		n = 1
	}
	// Return x || n with n as a byte and x an n bytes in big-endian order.
	b := make([]byte, 9)
	binary.BigEndian.PutUint64(b, x)
	b = b[9-n-1:]
	b[n] = byte(n)
	return b
}

func encodeString(s []byte) []byte {
	return append(leftEncode(uint64(len(s))), s...)
}

func bytepad(data []byte, rate int) []byte {
	out := make([]byte, 0, 9+len(data)+rate-1)
	out = append(out, leftEncode(uint64(rate))...)
	out = append(out, data...)
	if padlen := rate - len(out)%rate; padlen < rate {
		out = append(out, make([]byte, padlen)...)
	}
	return out
}

// KMAC128 computes the KMAC128 message authentication code for data using key
// as the secret key. KMAC provides variable-length output and fills out with
// the resulting MAC. It's recommended to use at least 8 bytes of output.
func KMAC128(out, key, data []byte, optionalCustomizationString string) {
	shake := sha3.NewCSHAKE128([]byte("KMAC"), []byte(optionalCustomizationString))
	shake.Write(bytepad(encodeString(key), 168))
	shake.Write(data)
	shake.Write(rightEncode(uint64(len(out))))
	shake.Read(out)
}

// KMAC256 computes the KMAC256 message authentication code for data using key
// as the secret key. KMAC provides variable-length output and fills out with
// the resulting MAC. It's recommended to use at least 8 bytes of output.
func KMAC256(out, key, data []byte, optionalCustomizationString string) {
	shake := sha3.NewCSHAKE256([]byte("KMAC"), []byte(optionalCustomizationString))
	shake.Write(bytepad(encodeString(key), 136))
	shake.Write(data)
	shake.Write(rightEncode(uint64(len(out))))
	shake.Read(out)
}