package auth

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"

	"upspin.io/upspin"
)

// testIssuer generates an ephemeral P-256 key and returns the upspin public-key
// string plus a mintToken function that signs tokens exactly as the vault issuer
// will (this is the executable spec for vault's mintToken).
func testIssuer(t *testing.T) (upspin.PublicKey, func(claims vaultClaims) string) {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	pub := upspin.PublicKey(fmt.Sprintf("p256\n%s\n%s\n", priv.X.String(), priv.Y.String()))

	mint := func(claims vaultClaims) string {
		payload, err := json.Marshal(claims)
		if err != nil {
			t.Fatal(err)
		}
		payloadB64 := base64.RawURLEncoding.EncodeToString(payload)
		hash := sha256.Sum256([]byte("v1." + payloadB64))
		r, s, err := ecdsa.Sign(rand.Reader, priv, hash[:])
		if err != nil {
			t.Fatal(err)
		}
		return "v1." + payloadB64 + "." +
			base64.RawURLEncoding.EncodeToString(r.Bytes()) + "." +
			base64.RawURLEncoding.EncodeToString(s.Bytes())
	}
	return pub, mint
}

func TestVaultVerify(t *testing.T) {
	pub, mint := testIssuer(t)
	const issuer = "vault@cosmocr.at"
	lookup := func(u upspin.UserName) (upspin.PublicKey, error) {
		if u != issuer {
			return "", fmt.Errorf("unexpected lookup for %q", u)
		}
		return pub, nil
	}
	v := NewVaultVerifier(issuer, "gitweb", lookup)

	valid := vaultClaims{Iss: issuer, Sub: "alice@d.com", Aud: "gitweb",
		Iat: time.Now().Unix(), Exp: time.Now().Add(time.Hour).Unix()}

	sub, err := v.Verify(mint(valid))
	if err != nil {
		t.Fatalf("Verify valid token: %v", err)
	}
	if sub != "alice@d.com" {
		t.Fatalf("sub = %q, want alice@d.com", sub)
	}
}

func TestVaultVerifyRejects(t *testing.T) {
	pub, mint := testIssuer(t)
	const issuer = "vault@cosmocr.at"
	v := NewVaultVerifier(issuer, "gitweb", func(upspin.UserName) (upspin.PublicKey, error) {
		return pub, nil
	})

	now := time.Now()
	base := vaultClaims{Iss: issuer, Sub: "alice@d.com", Aud: "gitweb",
		Iat: now.Unix(), Exp: now.Add(time.Hour).Unix()}

	// Expired.
	exp := base
	exp.Exp = now.Add(-time.Minute).Unix()
	if _, err := v.Verify(mint(exp)); err == nil {
		t.Error("expired token accepted")
	}

	// Wrong audience.
	aud := base
	aud.Aud = "someone-else"
	if _, err := v.Verify(mint(aud)); err == nil {
		t.Error("wrong-audience token accepted")
	}

	// Wrong issuer.
	iss := base
	iss.Iss = "attacker@evil.com"
	if _, err := v.Verify(mint(iss)); err == nil {
		t.Error("wrong-issuer token accepted")
	}

	// Tampered signature (flip the payload but keep the old signature).
	good := mint(base)
	tampered := tamper(t, good)
	if _, err := v.Verify(tampered); err == nil {
		t.Error("tampered token accepted")
	}

	// Signed by a different key.
	otherPub, otherMint := testIssuer(t)
	_ = otherPub
	if _, err := v.Verify(otherMint(base)); err == nil {
		t.Error("token signed by wrong key accepted")
	}
}

// tamper swaps in a forged payload (a different subject) while keeping the
// original signature segments, so the signature no longer matches.
func tamper(t *testing.T, token string) string {
	t.Helper()
	parts := strings.Split(token, ".") // v1 . payload . r . s
	if len(parts) != 4 {
		t.Fatalf("unexpected token shape: %q", token)
	}
	forged := vaultClaims{Iss: "vault@cosmocr.at", Sub: "attacker@d.com", Aud: "gitweb",
		Iat: time.Now().Unix(), Exp: time.Now().Add(time.Hour).Unix()}
	payload, _ := json.Marshal(forged)
	parts[1] = base64.RawURLEncoding.EncodeToString(payload)
	return strings.Join(parts, ".")
}