// SPDX-License-Identifier: MIT

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"os"
	"strings"
	"testing"
)

// vectors are values made by the dashboard's own crypto, so passing them
// means this tool and the dashboard agree, not only that this tool agrees
// with itself.
type vectors struct {
	StoreKeySlot struct {
		RawKey     string   `json:"raw_key"`
		Salt       string   `json:"salt"`
		Info       string   `json:"info"`
		KEKHex     string   `json:"kek_hex"`
		DEKHex     string   `json:"dek_hex"`
		KeyHashHex string   `json:"key_hash_hex"`
		SlotsB64   []string `json:"slots_b64"`
	} `json:"store_key_slot"`
	CustomerSeal struct {
		CustomerKeyHex string `json:"customer_key_hex"`
		Salt           string `json:"salt"`
		DEKHex         string `json:"dek_hex"`
		SealB64        string `json:"seal_b64"`
	} `json:"customer_seal"`
	Password struct {
		Email              string `json:"email"`
		Password           string `json:"password"`
		NormalizedEmail    string `json:"normalized_email"`
		MasterKeyHex       string `json:"master_key_hex"`
		LoginProofHex      string `json:"login_proof_hex"`
		CustomerKeyHex     string `json:"customer_key_hex"`
		CustomerKeyWrapB64 string `json:"customer_key_wrap_b64"`
		CustomerKeyText    string `json:"customer_key_text"`
	} `json:"password"`
	NormalizedEmails [][2]string `json:"normalized_emails"`
}

func loadVectors(t *testing.T) vectors {
	t.Helper()
	raw, err := os.ReadFile("testdata/vectors.json")
	if err != nil {
		t.Fatal(err)
	}
	var v vectors
	if err := json.Unmarshal(raw, &v); err != nil {
		t.Fatal(err)
	}
	return v
}

func unhex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func TestTheSlotFormatIsTheOneTheVectorsWereMadeIn(t *testing.T) {
	v := loadVectors(t)
	if slotFormat != v.StoreKeySlot.Info {
		t.Fatalf("this tool makes %q, the vectors are %q", slotFormat, v.StoreKeySlot.Info)
	}
}

func TestAStoreKeyDerivesTheWrappingKeyTheServiceAndTheBrowserDerive(t *testing.T) {
	v := loadVectors(t).StoreKeySlot
	kek, err := slotKEK(v.RawKey, v.Salt)
	if err != nil {
		t.Fatal(err)
	}
	if got := hex.EncodeToString(kek); got != v.KEKHex {
		t.Fatalf("kek %s, want %s", got, v.KEKHex)
	}
	if got := keyHash(v.RawKey); got != v.KeyHashHex {
		t.Fatalf("key hash %s, want %s", got, v.KeyHashHex)
	}
}

func TestEverySlotTheOthersMadeOpensToTheSameDataKey(t *testing.T) {
	v := loadVectors(t).StoreKeySlot
	if len(v.SlotsB64) < 2 {
		t.Fatal("the vectors carry fewer slots than the implementations that made them")
	}
	for i, s := range v.SlotsB64 {
		dek, err := openSlot(s, v.RawKey, v.Salt)
		if err != nil {
			t.Fatalf("slot %d: %v", i, err)
		}
		if !bytes.Equal(dek, unhex(t, v.DEKHex)) {
			t.Fatalf("slot %d opened to %x", i, dek)
		}
	}
}

func TestASlotMadeHereHasTheSharedLayoutAndOpensUnderTheSharedKey(t *testing.T) {
	v := loadVectors(t).StoreKeySlot
	dek := unhex(t, v.DEKHex)
	s, err := makeSlot(dek, v.RawKey, v.Salt)
	if err != nil {
		t.Fatal(err)
	}
	blob, err := base64.StdEncoding.DecodeString(s)
	if err != nil || len(blob) != nonceLen+keyLen+tagLen {
		t.Fatalf("slot is %d bytes (%v), want nonce, data key and tag", len(blob), err)
	}
	// Opened with the vector's key directly, not through this tool's own
	// derivation, so the two halves are not checking each other.
	plain, err := open(unhex(t, v.KEKHex), s)
	if err != nil || !bytes.Equal(plain, dek) {
		t.Fatalf("opened to %x, %v", plain, err)
	}
	again, _ := makeSlot(dek, v.RawKey, v.Salt)
	if again == s {
		t.Fatal("two slots of one key are identical: the nonce is not fresh")
	}
}

func TestASlotThatWasTamperedWithOrIsShortOrUnderAnotherKeyDoesNotOpen(t *testing.T) {
	v := loadVectors(t).StoreKeySlot
	good := v.SlotsB64[0]
	blob, _ := base64.StdEncoding.DecodeString(good)
	blob[len(blob)-1] ^= 1
	cases := map[string]string{
		"tampered":   base64.StdEncoding.EncodeToString(blob),
		"short":      base64.StdEncoding.EncodeToString(blob[:nonceLen+tagLen-1]),
		"not base64": "%%%",
	}
	for name, s := range cases {
		if _, err := openSlot(s, v.RawKey, v.Salt); !errors.Is(err, errWrongKey) {
			t.Fatalf("%s: %v", name, err)
		}
	}
	if _, err := openSlot(good, v.RawKey+"x", v.Salt); !errors.Is(err, errWrongKey) {
		t.Fatalf("another key: %v", err)
	}
	if _, err := openSlot(good, v.RawKey, v.Salt+"x"); !errors.Is(err, errWrongKey) {
		t.Fatalf("another account: %v", err)
	}
}

func TestTheCustomerKeyOpensTheSealTheBrowserMade(t *testing.T) {
	v := loadVectors(t).CustomerSeal
	dek, err := openCustomerSeal(v.SealB64, unhex(t, v.CustomerKeyHex), v.Salt)
	if err != nil || !bytes.Equal(dek, unhex(t, v.DEKHex)) {
		t.Fatalf("opened to %x, %v", dek, err)
	}
	// The seal's label is not a slot's: the customer key presented as if it
	// were a store key opens nothing.
	asKey := storeKeyPrefix + base64.RawURLEncoding.EncodeToString(unhex(t, v.CustomerKeyHex))
	if _, err := openSlot(v.SealB64, asKey, v.Salt); err == nil {
		t.Fatal("a store's seal opened as a slot")
	}
}

func TestAPasswordDerivesTheProofAndOpensTheCustomerKeyAsInTheBrowser(t *testing.T) {
	v := loadVectors(t).Password
	if got := normEmail(v.Email); got != v.NormalizedEmail {
		t.Fatalf("email %q, want %q", got, v.NormalizedEmail)
	}
	master, err := masterKey(v.Password, v.Email)
	if err != nil {
		t.Fatal(err)
	}
	if got := hex.EncodeToString(master); got != v.MasterKeyHex {
		t.Fatalf("master key %s, want %s", got, v.MasterKeyHex)
	}
	proof, err := loginProof(master)
	if err != nil || proof != v.LoginProofHex {
		t.Fatalf("proof %s (%v), want %s", proof, err, v.LoginProofHex)
	}
	ck, err := unwrapCustomerKey(master, v.CustomerKeyWrapB64)
	if err != nil || hex.EncodeToString(ck) != v.CustomerKeyHex {
		t.Fatalf("customer key %x (%v)", ck, err)
	}
	other, _ := masterKey(v.Password+"!", v.Email)
	if _, err := unwrapCustomerKey(other, v.CustomerKeyWrapB64); !errors.Is(err, errWrongKey) {
		t.Fatalf("a wrong password opened the customer key: %v", err)
	}
}

func TestAnAddressIsNormalizedAsJavaScriptNormalizesIt(t *testing.T) {
	for _, c := range loadVectors(t).NormalizedEmails {
		if got := normEmail(c[0]); got != c[1] {
			t.Errorf("normEmail(%+q) = %+q, want %+q", c[0], got, c[1])
		}
	}
}

func TestASavedCustomerKeyIsReadInBothFormsAndOnlyInItsOwnSpelling(t *testing.T) {
	v := loadVectors(t).Password
	want := unhex(t, v.CustomerKeyHex)
	for _, text := range []string{v.CustomerKeyText, "dke_k_" + strings.TrimPrefix(v.CustomerKeyText, "dke_c_")} {
		got, ok := parseCustomerKey(text)
		if !ok || !bytes.Equal(got, want) {
			t.Fatalf("%s read as %x, %v", text, got, ok)
		}
	}
	// The last character carries two bits the key does not have. Setting one
	// spells a different text for the same bytes, which is refused.
	body := strings.TrimPrefix(v.CustomerKeyText, "dke_c_")
	const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
	last := strings.IndexByte(alphabet, body[len(body)-1])
	neighbour := "dke_c_" + body[:len(body)-1] + string(alphabet[last^1])
	refused := []string{neighbour, "dke_c_" + body[:42], "dke_s_" + body, "dke_c_" + body + "A", ""}
	for _, text := range refused {
		if _, ok := parseCustomerKey(text); ok {
			t.Fatalf("%q was read as a customer key", text)
		}
	}
}

func TestANewStoreKeyIsThirtyTwoRandomBytesInTheDashboardsSpelling(t *testing.T) {
	a, err := newStoreKey()
	if err != nil {
		t.Fatal(err)
	}
	b, _ := newStoreKey()
	if a == b {
		t.Fatal("two keys are the same")
	}
	body, ok := strings.CutPrefix(a, storeKeyPrefix)
	if !ok || len(body) != 43 {
		t.Fatalf("key %q", a)
	}
	raw, err := base64.RawURLEncoding.DecodeString(body)
	if err != nil || len(raw) != keyLen {
		t.Fatalf("key body decodes to %d bytes (%v)", len(raw), err)
	}
}