// SPDX-License-Identifier: MIT

package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"regexp"
	"strings"
	"unicode"
)

// slotFormat names the label a store key's slot is made under. It is sent
// with every key this tool makes: the service cannot open a slot to check it,
// so a copy of this tool left behind by a change of label is refused by name
// rather than leaving a key that opens nothing.
const slotFormat = "dke-store-kek-v2"

const (
	customerSealInfo = "dke-store-customer-kek-v1"
	loginProofInfo   = "langsyn-login-proof-v1"
	customerKeyInfo  = "langsyn-customer-key-v1"
	pbkdf2Iterations = 600000

	keyLen   = 32
	nonceLen = 12
	tagLen   = 16

	storeKeyPrefix = "dke_s_"
)

var errWrongKey = errors.New("the key does not open this")

// deriveKey is HKDF-SHA256 to 32 bytes. A nil salt is the empty salt.
func deriveKey(ikm, salt []byte, info string) ([]byte, error) {
	return hkdf.Key(sha256.New, ikm, salt, info, keyLen)
}

// seal wraps plain under kek as base64(nonce || ciphertext || tag).
func seal(kek, plain []byte) (string, error) {
	gcm, err := newGCM(kek)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, nonceLen)
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(gcm.Seal(nonce, nonce, plain, nil)), nil
}

// open reverses seal, and fails closed on any wrap kek did not make.
func open(kek []byte, wrapped string) ([]byte, error) {
	blob, err := base64.StdEncoding.DecodeString(wrapped)
	if err != nil || len(blob) < nonceLen+tagLen {
		return nil, errWrongKey
	}
	gcm, err := newGCM(kek)
	if err != nil {
		return nil, err
	}
	plain, err := gcm.Open(nil, blob[:nonceLen], blob[nonceLen:], nil)
	if err != nil {
		return nil, errWrongKey
	}
	return plain, nil
}

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

// keyHash is what the service knows a store key by: the lowercase hex
// SHA-256 of the key's text.
func keyHash(storeKey string) string {
	sum := sha256.Sum256([]byte(storeKey))
	return hex.EncodeToString(sum[:])
}

// newStoreKey makes a store key from the system's random source: the prefix
// and 32 random bytes in unpadded URL-safe base64, as the dashboard does.
func newStoreKey() (string, error) {
	raw := make([]byte, keyLen)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	return storeKeyPrefix + base64.RawURLEncoding.EncodeToString(raw), nil
}

// slotKEK is the key a store key's slot is wrapped under.
func slotKEK(storeKey, customerID string) ([]byte, error) {
	return deriveKey([]byte(storeKey), []byte(customerID), slotFormat)
}

// makeSlot wraps a store's data key under a new store key.
func makeSlot(dek []byte, storeKey, customerID string) (string, error) {
	kek, err := slotKEK(storeKey, customerID)
	if err != nil {
		return "", err
	}
	defer clear(kek)
	return seal(kek, dek)
}

// openSlot opens a store key's slot to the store's data key.
func openSlot(wrapped, storeKey, customerID string) ([]byte, error) {
	kek, err := slotKEK(storeKey, customerID)
	if err != nil {
		return nil, err
	}
	defer clear(kek)
	return dataKey(open(kek, wrapped))
}

// openCustomerSeal opens a store's seal with the account's customer key.
func openCustomerSeal(sealed string, customerKey []byte, customerID string) ([]byte, error) {
	kek, err := deriveKey(customerKey, []byte(customerID), customerSealInfo)
	if err != nil {
		return nil, err
	}
	defer clear(kek)
	return dataKey(open(kek, sealed))
}

func dataKey(dek []byte, err error) ([]byte, error) {
	if err != nil {
		return nil, err
	}
	if len(dek) != keyLen {
		clear(dek)
		return nil, errWrongKey
	}
	return dek, nil
}

// masterKey is what a password and an email address derive. It never leaves
// this process: the account is shown a proof made from it, and the customer
// key is kept under a key made from it.
func masterKey(password, email string) ([]byte, error) {
	return pbkdf2.Key(sha256.New, password, []byte(normEmail(email)), pbkdf2Iterations, keyLen)
}

// loginProof is what the account checks in place of the password. It is
// one-way: nothing that holds it can work back to the master key.
func loginProof(master []byte) (string, error) {
	proof, err := deriveKey(master, nil, loginProofInfo)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(proof), nil
}

// unwrapCustomerKey opens the account's customer key with the master key.
func unwrapCustomerKey(master []byte, wrapped string) ([]byte, error) {
	kek, err := deriveKey(master, nil, customerKeyInfo)
	if err != nil {
		return nil, err
	}
	defer clear(kek)
	return dataKey(open(kek, wrapped))
}

var customerKeyText = regexp.MustCompile(`^dke_[ck]_([A-Za-z0-9_-]{43})$`)

// parseCustomerKey reads a saved customer key back to its 32 bytes. Both
// prefixes the dashboard has shown are read. A body that decodes but is not
// the one the key encodes to is refused, rather than read as a neighbouring
// key: 43 characters carry two bits more than 32 bytes.
func parseCustomerKey(text string) ([]byte, bool) {
	m := customerKeyText.FindStringSubmatch(text)
	if m == nil {
		return nil, false
	}
	key, err := base64.RawURLEncoding.DecodeString(m[1])
	if err != nil || len(key) != keyLen || base64.RawURLEncoding.EncodeToString(key) != m[1] {
		return nil, false
	}
	return key, true
}

// normEmail is the address as the key derivation sees it, which must be
// exactly what the dashboard derives from: JavaScript's trim, then its
// toLowerCase. Go's own functions differ from both at the edges, and a
// different address is a different master key, read by the account as a
// wrong password.
func normEmail(email string) string {
	return jsLower(strings.TrimFunc(email, jsSpace))
}

// jsSpace is JavaScript's white space and line terminators. It is not
// unicode.IsSpace: JavaScript trims U+FEFF, and does not trim U+0085.
func jsSpace(r rune) bool {
	switch r {
	case '\t', '\n', '\v', '\f', '\r', ' ', '\u00a0', '\ufeff', '\u2028', '\u2029':
		return true
	}
	return unicode.Is(unicode.Zs, r)
}

// jsLower is JavaScript's toLowerCase: Unicode's default lower case, with the
// two mappings a rune-by-rune lower case misses. A dotted capital I lowers
// to "i" and a combining dot, and a capital sigma ending a word lowers to the
// final sigma.
func jsLower(s string) string {
	rs := []rune(s)
	var b strings.Builder
	for i, r := range rs {
		switch {
		case r == '\u0130':
			b.WriteString("i\u0307")
		case r == '\u03a3' && finalSigma(rs, i):
			b.WriteRune('\u03c2')
		default:
			b.WriteRune(unicode.ToLower(r))
		}
	}
	return b.String()
}

// finalSigma is Unicode's Final_Sigma condition: a cased letter before the
// sigma, and none after it, with case-ignorable characters skipped on both
// sides.
func finalSigma(rs []rune, at int) bool {
	before := false
	for i := at - 1; i >= 0; i-- {
		if caseIgnorable(rs[i]) {
			continue
		}
		before = cased(rs[i])
		break
	}
	if !before {
		return false
	}
	for i := at + 1; i < len(rs); i++ {
		if caseIgnorable(rs[i]) {
			continue
		}
		return !cased(rs[i])
	}
	return true
}

func cased(r rune) bool {
	return unicode.IsUpper(r) || unicode.IsLower(r) || unicode.IsTitle(r) ||
		unicode.Is(unicode.Other_Lowercase, r) || unicode.Is(unicode.Other_Uppercase, r)
}

// caseIgnorable approximates Unicode's Case_Ignorable: the marks, format
// characters and modifiers, and the word-internal punctuation an address
// can carry.
func caseIgnorable(r rune) bool {
	switch r {
	case '\'', '.', ':', '^', '`', '\u00b7', '\u2018', '\u2019', '\u2024', '\u2027':
		return true
	}
	return unicode.In(r, unicode.Mn, unicode.Me, unicode.Cf, unicode.Lm, unicode.Sk)
}