package main

import (
	"crypto/rand"
	"crypto/rsa"
	"github.com/golang-jwt/jwt/v5"
	"testing"
	"time"
)

func TestJWTAndOIDC(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	pub := &key.PublicKey
	if checkJWT(token(key, "orders-api", "", now.Add(time.Minute)), pub) != nil {
		t.Fatal("valid jwt")
	}
	if checkJWT(token(key, "other-api", "", now.Add(time.Minute)), pub) == nil {
		t.Fatal("wrong audience")
	}
	h := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{"iss": issuer, "aud": "orders-api", "exp": now.Add(time.Minute).Unix()})
	raw, err := h.SignedString([]byte("PUBLIC TEST KEY, NOT A SECRET"))
	if err != nil {
		t.Fatal(err)
	}
	if checkJWT(raw, pub) == nil {
		t.Fatal("wrong algorithm")
	}
	v := verifier(pub)
	raw = token(key, "web-client", publicNonce, now.Add(time.Minute))
	if err = checkOIDC(v, raw, publicNonce); err != nil {
		t.Fatal(err)
	}
	for name, raw := range map[string]string{"wrong-aud": token(key, "other-client", publicNonce, now.Add(time.Minute)), "expired": token(key, "web-client", publicNonce, now.Add(-time.Second)), "expiry-equal": token(key, "web-client", publicNonce, now)} {
		if checkOIDC(v, raw, publicNonce) == nil {
			t.Fatal(name)
		}
	}
	if checkOIDC(v, raw, "wrong") == nil {
		t.Fatal("nonce mismatch")
	}
}
func TestPolicy(t *testing.T) {
	e, err := enforcer()
	if err != nil {
		t.Fatal(err)
	}
	for _, tt := range []struct {
		user, dom string
		r         Resource
		action    string
		want      bool
	}{{"alice", "team-a", Resource{"doc-1", "team-a"}, "read", true}, {"bob", "team-a", Resource{"doc-1", "team-a"}, "read", false}, {"alice", "team-a", Resource{"doc-1", "team-b"}, "read", false}, {"alice", "team-b", Resource{"doc-1", "team-b"}, "read", false}, {"alice", "team-a", Resource{"doc-1", "team-a"}, "delete", false}} {
		got, err := e.Enforce(tt.user, tt.dom, tt.r, tt.action)
		if err != nil || got != tt.want {
			t.Fatal(tt, got, err)
		}
	}
	if _, err = e.RemoveGroupingPolicy("alice", "reader", "team-a"); err != nil {
		t.Fatal(err)
	}
	if ok, err := e.Enforce("alice", "team-a", Resource{"doc-1", "team-a"}, "read"); err != nil || ok {
		t.Fatal("revocation", ok, err)
	}
}

func TestOIDCApplicationBoundaries(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	v := verifier(&key.PublicKey)
	base := func() jwt.MapClaims {
		return jwt.MapClaims{"iss": issuer, "sub": "alice", "aud": "web-client", "iat": now.Add(-time.Second).Unix(), "exp": now.Add(3 * time.Minute).Unix(), "nonce": publicNonce}
	}
	sign := func(c jwt.MapClaims, k *rsa.PrivateKey) string {
		raw, err := jwt.NewWithClaims(jwt.SigningMethodRS256, c).SignedString(k)
		if err != nil {
			t.Fatal(err)
		}
		return raw
	}
	cases := []struct {
		name   string
		mutate func(jwt.MapClaims)
		want   bool
	}{
		{"wrong-issuer", func(c jwt.MapClaims) { c["iss"] = "https://other.example" }, false},
		{"missing-sub", func(c jwt.MapClaims) { delete(c, "sub") }, false},
		{"missing-iat", func(c jwt.MapClaims) { delete(c, "iat") }, false},
		{"future-iat", func(c jwt.MapClaims) { c["iat"] = now.Add(2 * time.Minute).Unix() }, false},
		{"multi-aud-no-azp", func(c jwt.MapClaims) { c["aud"] = []string{"web-client", "other"} }, false},
		{"multi-aud-wrong-azp", func(c jwt.MapClaims) { c["aud"] = []string{"web-client", "other"}; c["azp"] = "other" }, false},
		{"multi-aud-correct-azp", func(c jwt.MapClaims) { c["aud"] = []string{"web-client", "other"}; c["azp"] = "web-client" }, true},
		{"single-aud-wrong-azp", func(c jwt.MapClaims) { c["azp"] = "other" }, false},
	}
	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			c := base()
			tt.mutate(c)
			err := checkOIDC(v, sign(c, key), publicNonce)
			if (err == nil) != tt.want {
				t.Fatal(err)
			}
		})
	}
	if checkOIDC(v, sign(base(), key), "") == nil {
		t.Fatal("empty expected nonce")
	}
	other, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	if checkOIDC(v, sign(base(), other), publicNonce) == nil {
		t.Fatal("wrong signing key")
	}
}
func TestJWTProfileAndClaims(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"typ", "issuer", "expired", "sub", "jti", "iat", "nbf", "exp"} {
		c := jwt.MapClaims{"iss": issuer, "sub": "alice", "aud": "orders-api", "iat": now.Add(-time.Second).Unix(), "nbf": now.Add(-time.Second).Unix(), "exp": now.Add(time.Minute).Unix(), "jti": "test-id"}
		token := jwt.NewWithClaims(jwt.SigningMethodRS256, c)
		token.Header["typ"] = "demo-access+jwt"
		switch name {
		case "typ":
			token.Header["typ"] = "JWT"
		case "issuer":
			c["iss"] = "https://other.example"
		case "expired":
			c["exp"] = now.Add(-time.Second).Unix()
		default:
			delete(c, name)
		}
		raw, err := token.SignedString(key)
		if err != nil {
			t.Fatal(err)
		}
		if checkJWT(raw, &key.PublicKey) == nil {
			t.Fatal("accepted", name)
		}
	}
	e, err := enforcer()
	if err != nil {
		t.Fatal(err)
	}
	if ok, err := e.Enforce("alice", "team-a", Resource{"doc-2", "team-a"}, "read"); err != nil || ok {
		t.Fatal("wrong object", ok, err)
	}
}