package aviso

import (
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"math/big"
	"strings"
	"testing"
)

func TestGenerateKeyRoundTrips(t *testing.T) {
	priv, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	raw, err := base64.RawURLEncoding.DecodeString(priv)
	if err != nil || len(raw) != 32 {
		t.Fatalf("private key = %q: want 32 unpadded base64url bytes (err %v)", priv, err)
	}
	pub, id, err := parsePrivateKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	pubRaw, err := base64.RawURLEncoding.DecodeString(pub)
	if err != nil || len(pubRaw) != 65 || pubRaw[0] != 0x04 {
		t.Fatalf("public key = %q: want 65-byte uncompressed point", pub)
	}
	if id == "" || strings.ContainsAny(id, "+/=") {
		t.Fatalf("key id = %q: want unpadded base64url", id)
	}
	// The same key must yield the same id after a restart — Sweep and
	// Send match rows on it.
	_, id2, _ := parsePrivateKey(priv)
	if id != id2 {
		t.Fatal("key id not deterministic")
	}
}

func TestParsePrivateKeyRefusesBadInput(t *testing.T) {
	if _, _, err := parsePrivateKey(""); !errors.Is(err, ErrEmptyPrivateKey) {
		t.Fatalf("empty: got %v, want ErrEmptyPrivateKey", err)
	}
	enc := base64.RawURLEncoding.EncodeToString
	zero := enc(make([]byte, 32))
	// The P-256 group order n, and n-1 (the largest valid scalar).
	n, _ := new(big.Int).SetString("ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551", 16)
	nMinus1 := new(big.Int).Sub(n, big.NewInt(1))
	nPlus1 := new(big.Int).Add(n, big.NewInt(1))
	one := make([]byte, 32)
	one[31] = 1
	good := enc(one)
	if _, _, err := parsePrivateKey(enc(nMinus1.FillBytes(make([]byte, 32)))); err != nil {
		t.Errorf("n-1 refused: %v", err)
	}
	for name, in := range map[string]string{
		"not base64":    "!!!",
		"short":         enc([]byte("short")),
		"zero scalar":   zero,
		"order n":       enc(n.FillBytes(make([]byte, 32))),
		"above n":       enc(nPlus1.FillBytes(make([]byte, 32))),
		"padded":        good + "=",
		"newline":       good[:10] + "\n" + good[10:],
		"trailing bits": good[:len(good)-1] + "F", // canonical ending is "E"
	} {
		if _, _, err := parsePrivateKey(in); !errors.Is(err, ErrInvalidPrivateKey) {
			t.Errorf("%s: got %v, want ErrInvalidPrivateKey", name, err)
		}
	}
}

// New must wire parsePrivateKey's outputs unchanged: two Services on
// one key agree on the public key and the key id rows are stamped with.
func TestNewWiresTheSameKeyIdentityEveryTime(t *testing.T) {
	key, _ := GenerateKey()
	mk := func() *Service {
		s, err := New(Config{DB: openInternalDB(t), PrivateKey: key, Contact: "mailto:x@y", Origin: "https://a"})
		if err != nil {
			t.Fatal(err)
		}
		return s
	}
	a, b := mk(), mk()
	pub, id, _ := parsePrivateKey(key)
	if a.pub != pub || b.pub != pub || a.keyID != id || b.keyID != id {
		t.Fatalf("New drifted from parsePrivateKey: %q/%q vs %q/%q", a.pub, a.keyID, pub, id)
	}
}

func TestKeyIDIsSHA256OfThePublicPoint(t *testing.T) {
	priv, _ := GenerateKey()
	pub, id, err := parsePrivateKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	point, _ := base64.RawURLEncoding.DecodeString(pub)
	sum := sha256.Sum256(point)
	if want := base64.RawURLEncoding.EncodeToString(sum[:]); id != want {
		t.Fatalf("key id = %s, want %s", id, want)
	}
}

func TestParsePrivateKeyAgreesWithECDH(t *testing.T) {
	k, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	priv := base64.RawURLEncoding.EncodeToString(k.Bytes())
	pub, _, err := parsePrivateKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	want := base64.RawURLEncoding.EncodeToString(k.PublicKey().Bytes())
	if pub != want {
		t.Fatalf("public key = %s, want %s", pub, want)
	}
}