package wristband

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"net/http"
	"regexp"
	"strings"
	"testing"
	"time"
)

// memberOrder returns the member names of a JSON object in wire order.
func memberOrder(t *testing.T, raw []byte) []string {
	t.Helper()
	dec := json.NewDecoder(strings.NewReader(string(raw)))
	if tok, err := dec.Token(); err != nil || tok != json.Delim('{') {
		t.Fatalf("not an object: %s", raw)
	}
	var names []string
	for dec.More() {
		tok, err := dec.Token()
		if err != nil {
			t.Fatal(err)
		}
		names = append(names, tok.(string))
		var skip json.RawMessage
		if err := dec.Decode(&skip); err != nil {
			t.Fatal(err)
		}
	}
	return names
}

// TestJWKSShapes pins the JWKS members and their order for RSA and EC keys,
// the cache header, and the fixed 500 for a key-store failure or a key whose
// algorithm does not match its type (T-19.4-08, T-19.4-09).
func TestJWKSShapes(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	f.keys.rotate(t, AlgRS256, time.Now())
	f.keys.rotate(t, AlgES256, time.Now())
	rec := f.get("/oauth/jwks")
	if rec.Code != http.StatusOK || rec.Header().Get("Cache-Control") != "max-age=3600, public" {
		t.Fatalf("jwks = %d %v", rec.Code, rec.Header())
	}
	var doc struct {
		Keys []json.RawMessage `json:"keys"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &doc); err != nil {
		t.Fatal(err)
	}
	if len(doc.Keys) != 2 {
		t.Fatalf("%d keys", len(doc.Keys))
	}
	if got := strings.Join(memberOrder(t, doc.Keys[0]), ","); got != "kty,use,alg,kid,n,e" {
		t.Fatalf("RSA members = %s", got)
	}
	if got := strings.Join(memberOrder(t, doc.Keys[1]), ","); got != "kty,use,alg,kid,crv,x,y" {
		t.Fatalf("EC members = %s", got)
	}
	var rsaKey, ecKey map[string]string
	_ = json.Unmarshal(doc.Keys[0], &rsaKey)
	_ = json.Unmarshal(doc.Keys[1], &ecKey)
	if rsaKey["kty"] != "RSA" || rsaKey["use"] != "sig" || rsaKey["alg"] != "RS256" || rsaKey["e"] != "AQAB" || !regexp.MustCompile(`^[0-9a-f]{16}$`).MatchString(rsaKey["kid"]) {
		t.Fatalf("RSA JWK = %v", rsaKey)
	}
	if ecKey["kty"] != "EC" || ecKey["crv"] != "P-256" || ecKey["alg"] != "ES256" || len(ecKey["x"]) != 43 || len(ecKey["y"]) != 43 {
		t.Fatalf("EC JWK = %v", ecKey)
	}
	if strings.ContainsAny(rec.Body.String(), "=+/") {
		t.Fatalf("JWK members are not unpadded base64url: %s", rec.Body.String())
	}

	// Key-store failures answer the fixed 500 without the error text.
	for _, failActive := range []bool{true, false} {
		srv := NewOIDCServer(Options{Issuer: testOIDCIssuer}, DefaultOIDCOptions())
		if err := srv.SetOIDCBackend(newMemoryOIDCBackend(), failingKeyStore{failActive: failActive, inner: newMemoryKeyStore()}, &fakeOIDCHost{}); err != nil {
			t.Fatal(err)
		}
		rec := serve(srv.JWKS, "/oauth/jwks")
		if rec.Code != http.StatusInternalServerError || rec.Body.String() != `{"message":"Signing keys are unavailable."}` {
			t.Fatalf("failing key store (active=%v) = %d %s", failActive, rec.Code, rec.Body.String())
		}
	}
	bare := NewOIDCServer(Options{Issuer: testOIDCIssuer}, DefaultOIDCOptions())
	if rec := serve(bare.JWKS, "/oauth/jwks"); rec.Code != http.StatusInternalServerError || rec.Body.String() != `{"message":"Signing keys are unavailable."}` {
		t.Fatalf("unconfigured JWKS = %d %s", rec.Code, rec.Body.String())
	}

	// A published key whose Alg does not match its type is refused.
	f.keys.mu.Lock()
	f.keys.keys[len(f.keys.keys)-1].alg = AlgRS256
	f.keys.mu.Unlock()
	if rec := f.get("/oauth/jwks"); rec.Code != http.StatusInternalServerError || rec.Body.String() != `{"message":"Signing keys are unavailable."}` {
		t.Fatalf("mislabelled key = %d %s", rec.Code, rec.Body.String())
	}
	rsaPriv, _ := rsa.GenerateKey(rand.Reader, 2048)
	p384, _ := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	for name, k := range map[string]PublicKey{
		"RSA labelled ES256": {Kid: "a", Alg: AlgES256, Key: &rsaPriv.PublicKey},
		"P-384 key":          {Kid: "b", Alg: AlgES256, Key: &p384.PublicKey},
		"unknown key type":   {Kid: "c", Alg: AlgRS256, Key: "not a key"},
	} {
		if _, err := jwkFor(k); err == nil {
			t.Fatalf("%s: jwkFor accepted it", name)
		}
	}
}

func pemOf(typ string, der []byte) []byte {
	return pem.EncodeToMemory(&pem.Block{Type: typ, Bytes: der})
}

// TestKeyPEMParsing pins the key parsers: PKCS#8 RSA and EC, PKCS#1 RSA,
// SEC1 EC, PKIX public keys, and fixed errors for garbage, other curves and
// a key of the wrong kind.
func TestKeyPEMParsing(t *testing.T) {
	for _, alg := range []string{AlgRS256, AlgES256} {
		kid, priv, pub, err := GenerateSigningKey(alg)
		if err != nil {
			t.Fatal(err)
		}
		if !regexp.MustCompile(`^[0-9a-f]{16}$`).MatchString(kid) || !strings.Contains(string(priv), "BEGIN PRIVATE KEY") || !strings.Contains(string(pub), "BEGIN PUBLIC KEY") {
			t.Fatalf("%s: kid %q, PEM types %q / %q", alg, kid, priv[:30], pub[:30])
		}
		signer, err := ParsePrivateKeyPEM(priv)
		if err != nil {
			t.Fatalf("%s PKCS#8: %v", alg, err)
		}
		public, err := ParsePublicKeyPEM(pub)
		if err != nil {
			t.Fatalf("%s PKIX: %v", alg, err)
		}
		if eq, ok := signer.Public().(interface{ Equal(crypto.PublicKey) bool }); !ok || !eq.Equal(public) {
			t.Fatalf("%s: the public half does not match the private key", alg)
		}
	}

	rsaKey, _ := rsa.GenerateKey(rand.Reader, 2048)
	if k, err := ParsePrivateKeyPEM(pemOf("RSA PRIVATE KEY", x509.MarshalPKCS1PrivateKey(rsaKey))); err != nil || !rsaKey.Equal(k) {
		t.Fatalf("PKCS#1: %v", err)
	}
	ecKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	sec1, _ := x509.MarshalECPrivateKey(ecKey)
	if k, err := ParsePrivateKeyPEM(pemOf("EC PRIVATE KEY", sec1)); err != nil || !ecKey.Equal(k) {
		t.Fatalf("SEC1: %v", err)
	}

	p384, _ := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	p384der, _ := x509.MarshalPKCS8PrivateKey(p384)
	p384pub, _ := x509.MarshalPKIXPublicKey(&p384.PublicKey)
	_, edKey, _ := ed25519.GenerateKey(rand.Reader)
	edPriv, _ := x509.MarshalPKCS8PrivateKey(edKey)
	rsaPub, _ := x509.MarshalPKIXPublicKey(&rsaKey.PublicKey)
	rsaPKCS8, _ := x509.MarshalPKCS8PrivateKey(rsaKey)

	for name, data := range map[string][]byte{
		"empty":                   nil,
		"no PEM block":            []byte("not a key"),
		"garbage DER":             pemOf("PRIVATE KEY", []byte("garbage")),
		"P-384 private key":       pemOf("PRIVATE KEY", p384der),
		"Ed25519 private key":     pemOf("PRIVATE KEY", edPriv),
		"a public key as private": pemOf("PUBLIC KEY", rsaPub),
	} {
		if _, err := ParsePrivateKeyPEM(data); err == nil || !strings.HasPrefix(err.Error(), "wristband: ") {
			t.Fatalf("private %s: %v", name, err)
		}
	}
	for name, data := range map[string][]byte{
		"empty":                   nil,
		"no PEM block":            []byte("not a key"),
		"garbage DER":             pemOf("PUBLIC KEY", []byte("garbage")),
		"P-384 public key":        pemOf("PUBLIC KEY", p384pub),
		"a private key as public": pemOf("PRIVATE KEY", rsaPKCS8),
	} {
		if _, err := ParsePublicKeyPEM(data); err == nil || !strings.HasPrefix(err.Error(), "wristband: ") {
			t.Fatalf("public %s: %v", name, err)
		}
	}
}

// TestES256Ring runs the flow on a ring whose only key is ES256: the
// id_token is ES256 with the active kid, the JWKS has one EC key and
// discovery lists ES256 only.
func TestES256Ring(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	f.keys.alg = AlgES256
	tok := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid profile"), oidcTestVerifier))
	key, _ := f.keys.ActiveSigningKey(context.Background())
	if h := jwtHeader(t, tok.IDToken); h["alg"] != "ES256" || h["kid"] != key.Kid {
		t.Fatalf("header = %v", h)
	}
	c := f.parseIDToken(t, tok.IDToken)
	if c["aud"] != f.client.ClientID || c["name"] != "Ada Example" {
		t.Fatalf("claims = %v", c)
	}
	var jwks struct {
		Keys []map[string]any `json:"keys"`
	}
	_ = json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks)
	if len(jwks.Keys) != 1 || jwks.Keys[0]["kty"] != "EC" || jwks.Keys[0]["kid"] != key.Kid {
		t.Fatalf("jwks = %v", jwks.Keys)
	}
	if body := f.get("/.well-known/openid-configuration").Body.String(); !strings.Contains(body, `"id_token_signing_alg_values_supported":["ES256"]`) {
		t.Fatalf("discovery = %s", body)
	}
	// The ES256 signature is a JOSE R||S pair of 64 bytes.
	if sig := strings.Split(tok.IDToken, ".")[2]; len(sig) != 86 {
		t.Fatalf("ES256 signature has %d base64url characters, want 86", len(sig))
	}
}

// TestMixedRingJWKS rotates between RS256 and ES256 in both directions: the
// JWKS serves both kinds in ring order inside the grace window and discovery
// always lists RS256 before ES256.
func TestMixedRingJWKS(t *testing.T) {
	for _, order := range [][2]string{{AlgES256, AlgRS256}, {AlgRS256, AlgES256}} {
		f := newOIDCFixture(t, testOIDCIssuer, nil)
		f.keys.alg = order[0]
		before := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid"), oidcTestVerifier))
		f.keys.rotate(t, order[1], time.Now())

		var jwks struct {
			Keys []map[string]any `json:"keys"`
		}
		_ = json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks)
		kty := map[string]string{AlgRS256: "RSA", AlgES256: "EC"}
		if len(jwks.Keys) != 2 || jwks.Keys[0]["kty"] != kty[order[0]] || jwks.Keys[1]["kty"] != kty[order[1]] {
			t.Fatalf("%v: jwks = %v", order, jwks.Keys)
		}
		if body := f.get("/.well-known/openid-configuration").Body.String(); !strings.Contains(body, `"id_token_signing_alg_values_supported":["RS256","ES256"]`) {
			t.Fatalf("%v: discovery = %s", order, body)
		}
		f.parseIDToken(t, before.IDToken)
		after := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid"), oidcTestVerifier))
		if jwtHeader(t, after.IDToken)["alg"] != order[1] {
			t.Fatalf("%v: a new token is not signed with the new key", order)
		}
		f.parseIDToken(t, after.IDToken)

		// Past the grace window the retired key and its algorithm leave.
		f.srv.now = func() time.Time { return time.Now().Add(25 * time.Hour) }
		_ = json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks)
		if len(jwks.Keys) != 1 || jwks.Keys[0]["kty"] != kty[order[1]] {
			t.Fatalf("%v: jwks after grace = %v", order, jwks.Keys)
		}
		if body := f.get("/.well-known/openid-configuration").Body.String(); !strings.Contains(body, `"id_token_signing_alg_values_supported":["`+order[1]+`"]`) {
			t.Fatalf("%v: discovery after grace = %s", order, body)
		}
	}
}