package push

import (
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/json"
	"io"
	"math/big"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func mustDecode(t *testing.T, s string) []byte {
	t.Helper()
	b, err := decode64(strings.Join(strings.Fields(s), ""))
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// TestEncryptRFC8291 is the worked example in RFC 8291 section 5 and appendix A.
func TestEncryptRFC8291(t *testing.T) {
	as, err := ecdh.P256().NewPrivateKey(mustDecode(t, "yfWPiYE-n46HLnH0KqZOF1fJJU3MYrct3AELtAQ-oRw"))
	if err != nil {
		t.Fatal(err)
	}
	got, err := encrypt(
		[]byte("When I grow up, I want to be a watermelon"),
		mustDecode(t, "BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4"),
		mustDecode(t, "BTBZMqHH6r4Tts7J_aSIgg"),
		as,
		mustDecode(t, "DGv6ra1nlYgDCS1FRnbzlw"),
	)
	if err != nil {
		t.Fatal(err)
	}
	want := append(
		mustDecode(t, `DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z 9KsN6nGRTbVYI_c7VJSPQTBtkgcy27ml
			mlMoZIIgDll6e3vCYLocInmYWAmS6Tlz AC8wEqKK6PBru3jl7A8`),
		mustDecode(t, `8pfeW0KbunFT06SuDKoJH9Ql87S1QUrd irN6GcG7sFz1y1sqLgVi1VhjVkHsUoEs bI_0LpXMuGvnzQ`)...)
	if !bytes.Equal(got, want) {
		t.Errorf("got  %x\nwant %x", got, want)
	}
}

func TestKeysRoundTrip(t *testing.T) {
	k, err := GenerateKeys()
	if err != nil {
		t.Fatal(err)
	}
	k2, err := ParseKeys(k.Private())
	if err != nil {
		t.Fatal(err)
	}
	if k2.Public() != k.Public() || len(mustDecode(t, k.Public())) != 65 {
		t.Errorf("public key changed or has the wrong length: %s", k.Public())
	}
}

// TestSend posts to a fake push service, then checks the VAPID signature and decrypts the body
// the way a browser would.
func TestSend(t *testing.T) {
	keys, _ := GenerateKeys()
	ua, _ := ecdh.P256().GenerateKey(rand.Reader)
	auth := []byte("0123456789abcdef")

	var got *http.Request
	var body []byte
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r
		body, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusCreated)
	}))
	defer srv.Close()

	s := &Sender{Keys: keys, Subject: "https://bans.example.com", Client: srv.Client()}
	sub := Subscription{Endpoint: srv.URL + "/send/abc", P256dh: b64.EncodeToString(ua.PublicKey().Bytes()), Auth: b64.EncodeToString(auth)}
	if err := sub.Check(); err != nil {
		t.Fatal(err)
	}
	if err := s.Send(context.Background(), sub, Message{Payload: []byte(`{"title":"hi"}`), TTL: time.Hour, Urgency: "high", Topic: "server"}); err != nil {
		t.Fatal(err)
	}
	if got.Header.Get("Content-Encoding") != "aes128gcm" || got.Header.Get("TTL") != "3600" || got.Header.Get("Urgency") != "high" || got.Header.Get("Topic") != "server" {
		t.Errorf("headers: %v", got.Header)
	}

	// VAPID: the JWT's audience is the push service's origin, signed by the key in k=.
	authz := got.Header.Get("Authorization")
	tok, k, ok := strings.Cut(strings.TrimPrefix(authz, "vapid t="), ", k=")
	if !ok || k != keys.Public() {
		t.Fatalf("Authorization: %s", authz)
	}
	parts := strings.Split(tok, ".")
	var claims struct {
		Aud string `json:"aud"`
		Sub string `json:"sub"`
		Exp int64  `json:"exp"`
	}
	_ = json.Unmarshal(mustDecode(t, parts[1]), &claims)
	if claims.Aud != srv.URL || claims.Sub != s.Subject || claims.Exp < time.Now().Unix() {
		t.Errorf("claims: %+v", claims)
	}
	sig := mustDecode(t, parts[2])
	digest := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	r, ss := new(big.Int).SetBytes(sig[:32]), new(big.Int).SetBytes(sig[32:])
	if !ecdsa.Verify(&keys.priv.PublicKey, digest[:], r, ss) {
		t.Error("VAPID signature doesn't verify")
	}

	// Decrypt as the browser: header is salt, record size, key id (the sender's one-off key).
	salt, asPub, ct := body[:16], body[21:86], body[86:]
	as, err := ecdh.P256().NewPublicKey(asPub)
	if err != nil {
		t.Fatal(err)
	}
	secret, _ := ua.ECDH(as)
	plain := decryptForTest(t, secret, auth, ua.PublicKey().Bytes(), asPub, salt, ct)
	if string(plain) != `{"title":"hi"}`+"\x02" {
		t.Errorf("decrypted %q", plain)
	}
}

func decryptForTest(t *testing.T, secret, auth, uaPub, asPub, salt, ct []byte) []byte {
	t.Helper()
	// The same derivation as encrypt, from the receiving side.
	prkKey := hkdfExtract(t, secret, auth)
	ikm := hkdfExpand(t, prkKey, "WebPush: info\x00"+string(uaPub)+string(asPub), 32)
	prk := hkdfExtract(t, ikm, salt)
	cek := hkdfExpand(t, prk, "Content-Encoding: aes128gcm\x00", 16)
	nonce := hkdfExpand(t, prk, "Content-Encoding: nonce\x00", 12)
	block, _ := aes.NewCipher(cek)
	gcm, _ := cipher.NewGCM(block)
	plain, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	return plain
}

func TestAllowedEndpoint(t *testing.T) {
	for ep, want := range map[string]bool{
		"https://fcm.googleapis.com/fcm/send/abc":                true,
		"https://updates.push.services.mozilla.com/wpush/v2/abc": true,
		"https://web.push.apple.com/abc":                         true,
		"https://wns2-par02p.notify.windows.com/w/?token=abc":    true,
		"http://fcm.googleapis.com/fcm/send/abc":                 false,
		"https://fcm.googleapis.com.evil.example/abc":            false,
		"https://evilfcm.googleapis.com/abc":                     false,
		"https://fcm.googleapis.com:8443/abc":                    false,
		"https://user@fcm.googleapis.com/abc":                    false,
		"https://127.0.0.1/abc":                                  false,
	} {
		if got := AllowedEndpoint(ep); got != want {
			t.Errorf("%s: got %v", ep, got)
		}
	}
}

func hkdfExtract(t *testing.T, secret, salt []byte) []byte {
	t.Helper()
	b, err := hkdf.Extract(sha256.New, secret, salt)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func hkdfExpand(t *testing.T, prk []byte, info string, n int) []byte {
	t.Helper()
	b, err := hkdf.Expand(sha256.New, prk, info, n)
	if err != nil {
		t.Fatal(err)
	}
	return b
}