// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"fmt"
	"math/big"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jws"
	"github.com/openpubkey/openpubkey/util"
	"github.com/stretchr/testify/require"
)

func TestSignVerifyJWT(t *testing.T) {
	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	oidcPubKey := &oidcPrivKey.PublicKey

	idToken, err := createOIDCToken(oidcPrivKey, "test")
	require.NoError(t, err)

	signerVerifier, err := NewSignerVerifier(oidcPubKey, 256)
	require.NoError(t, err)

	gqToken, err := signerVerifier.SignJWT(idToken)
	require.NoError(t, err)

	ok := signerVerifier.VerifyJWT(gqToken)
	require.True(t, ok, "signature verification failed")
}

func TestGQ256SignJWT(t *testing.T) {
	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	idToken, err := createOIDCToken(oidcPrivKey, "test")
	require.NoError(t, err)

	gqToken1, err := GQ256SignJWT(&oidcPrivKey.PublicKey, idToken)
	require.NoError(t, err)
	ok, err := GQ256VerifyJWT(&oidcPrivKey.PublicKey, gqToken1)
	require.NoError(t, err)
	require.True(t, ok)

	// Test that we throw the correct error if the wrong RSA public key is sent to SignJWT
	oidcPrivKeyWrong, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	require.NotEqual(t, oidcPrivKey, oidcPrivKeyWrong)

	gqToken2, err := GQ256SignJWT(&oidcPrivKeyWrong.PublicKey, idToken)
	require.ErrorContains(t, err, "incorrect public key supplied when GQ signing jwt")
	require.ErrorContains(t, err, "could not verify message using any of the signatures or keys")
	require.Nil(t, gqToken2)

	// Test specifying with extra claims
	expKey1 := "key1" // Expected claim keys, values
	expValue1 := "value1"
	expKey2 := "key2"
	expValue2 := "value2"

	gqTokenExtraClaims, err := GQ256SignJWT(&oidcPrivKey.PublicKey, idToken,
		WithExtraClaim(expKey1, expValue1),
		WithExtraClaim(expKey2, expValue2))
	require.NoError(t, err)

	retValue1, ok, err := getClaimInProtected(expKey1, gqTokenExtraClaims)
	require.NoError(t, err)
	require.True(t, ok)
	require.Equal(t, expValue1, retValue1)

	retValue2, ok, err := getClaimInProtected(expKey2, gqTokenExtraClaims)
	require.NoError(t, err)
	require.True(t, ok)
	require.Equal(t, expValue2, retValue2)

	// Test that we don't find a claim we didn't add
	retClaimValue, ok, err := getClaimInProtected("noSuchKey", gqTokenExtraClaims)
	require.NoError(t, err)
	require.False(t, ok, "we didn't add this claim, yet somehow it was in the protected header")
	require.Nil(t, retClaimValue)

	// Test that we throw the correct error if reserved header is used
	gqTokenReservedClaim, err := GQ256SignJWT(&oidcPrivKey.PublicKey, idToken,
		WithExtraClaim("alg", "ES256"))
	require.Error(t, err, "use of reserved claim name, alg, should throw an error")
	require.EqualError(t, err,
		"error creating GQ signature: use of reserved header name, alg, in additional headers",
		"incorrect error throw")
	require.Nil(t, gqTokenReservedClaim)
}

// TestGQ256SignJWT_PS256RejectsPS256 verifies that GQ256SignJWT rejects PS256 tokens.
// GQ signing is incompatible with PS256: PSS padding uses a randomized salt, so
// the encoded identity cannot be reconstructed deterministically for GQ1 verification.
func TestGQ256SignJWT_RejectsPS256(t *testing.T) {
	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	idToken, err := createOIDCTokenPS256(oidcPrivKey, "test")
	require.NoError(t, err)

	_, err = GQ256SignJWT(&oidcPrivKey.PublicKey, idToken)
	require.ErrorContains(t, err, "GQ signing requires an RS256-signed JWT")
}

func TestVerifyModifiedIdPayload(t *testing.T) {
	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	oidcPubKey := &oidcPrivKey.PublicKey

	idToken, err := createOIDCToken(oidcPrivKey, "test")
	require.NoError(t, err)

	// modify the ID Token payload to detect IdP signature invalidity via GQ verify
	modifiedToken, err := modifyTokenPayload(idToken, "fail")
	require.NoError(t, err)

	_, err = jws.Verify(modifiedToken, jws.WithKey(jwa.RS256(), oidcPubKey))
	require.Error(t, err, "ID token signature should fail for modified token")

	signerVerifier, err := NewSignerVerifier(oidcPubKey, 256)
	require.NoError(t, err)
	gqToken, err := signerVerifier.SignJWT(modifiedToken)
	require.NoError(t, err)

	ok := signerVerifier.VerifyJWT(gqToken)
	require.False(t, ok, "GQ signature verification passed for invalid payload")
}

func TestVerifyModifiedGqPayload(t *testing.T) {
	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	oidcPubKey := &oidcPrivKey.PublicKey

	idToken, err := createOIDCToken(oidcPrivKey, "test")
	require.NoError(t, err)

	signerVerifier, err := NewSignerVerifier(oidcPubKey, 256)
	require.NoError(t, err)
	gqToken, err := signerVerifier.SignJWT(idToken)
	require.NoError(t, err)
	// modify the ID Token payload to detect GQ signature invalidity
	modifiedToken, err := modifyTokenPayload(gqToken, "fail")
	require.NoError(t, err)

	ok := signerVerifier.VerifyJWT(modifiedToken)
	require.False(t, ok, "GQ signature verification passed for invalid payload")

}

func TestVerifyModifiedGqProtectedHeader(t *testing.T) {
	// This tests exists as regression test for the bug fixed in https://github.com/openpubkey/openpubkey/pull/379

	oidcPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	oidcPubKey := &oidcPrivKey.PublicKey

	tests := []struct {
		name          string
		claimToModify string
		origValue     string
		newValue      string
	}{
		{"modify JKT", "jkt", "original-thumbprint", "attacker-thumbprint"},
		{"modify CIC", "cic", "original-commitment", "attacker-commitment"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			idToken, err := createOIDCToken(oidcPrivKey, "test")
			require.NoError(t, err)

			signerVerifier, err := NewSignerVerifier(oidcPubKey, 256)
			require.NoError(t, err)
			gqToken, err := signerVerifier.SignJWT(idToken, WithExtraClaim(tt.claimToModify, tt.origValue))
			require.NoError(t, err)

			modifiedToken, err := modifyTokenProtectedHeaderClaim(gqToken, tt.claimToModify, tt.newValue)
			require.NoError(t, err)
			require.NotEqual(t, string(gqToken), modifiedToken, "expected token to be modified but token unchanged")

			ok := signerVerifier.VerifyJWT(modifiedToken)
			// If ok is true, it means that signature verification is not working correctly
			//  as altering the message should break the signature.
			require.False(t, ok, "verify passed for tampered %s", tt.claimToModify)
		})
	}

	// To ensure that the above test doesn't accidentally pass due to encoding breaking
	// the signature and not the tampering itself, we also test a find and replace
	// modification that does not require JSON parsing or re-encoding the header
	idToken, err := createOIDCToken(oidcPrivKey, "test")
	require.NoError(t, err)

	signerVerifier, err := NewSignerVerifier(oidcPubKey, 256)
	require.NoError(t, err)

	// We choose an 8 character string so that we can do a clean find and replace (no padding):
	// 8 bits*6 characters = 48 bits
	// 48/6 = 8 base64 characters
	gqSignedValue := "123456"
	gqSignedValueB64 := []byte("MTIzNDU2") // base64 of "123456"
	modificationB64 := []byte("YWJjZGVm")  // base64 of "abcdef"
	gqToken, err := signerVerifier.SignJWT(idToken, WithExtraClaim("claim", gqSignedValue))
	require.NoError(t, err)
	gqHeadersB64, payload, signature, err := jws.SplitCompact(gqToken)
	require.NoError(t, err)

	// Replace base64("123456") with base64("abcdef")
	modifiedGqHeaders := bytes.Replace(gqHeadersB64, gqSignedValueB64, modificationB64, 1)

	// Ensure headers were actually modified, and modification carries into the GQ signature, otherwise the test is invalid
	require.NotEqual(t, string(gqHeadersB64), string(modifiedGqHeaders), "modified token should not equal original token")
	modifiedGqToken := util.JoinJWTSegments(modifiedGqHeaders, payload, signature)
	require.NotEqual(t, string(gqToken), modifiedGqToken, "modified token should not equal original token")

	ok := signerVerifier.VerifyJWT(modifiedGqToken)
	require.False(t, ok, "verify passed for tampered token")
}

func TestNewSignerVerifierRejectsInvalidExponents(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	// 0, 1, 2 are rejected by E < 3; 4 and 6 are rejected by E%2 == 0.
	for _, e := range []int{0, 1, 2, 4, 6} {
		t.Run(fmt.Sprintf("E=%d", e), func(t *testing.T) {
			key.E = e
			sv, err := NewSignerVerifier(&key.PublicKey, 256)
			require.Error(t, err, "NewSignerVerifier(E=%d) should have returned an error", e)
			require.ErrorContains(t, err, "unsupported RSA public key exponent")
			require.Nil(t, sv)
		})
	}
}

func TestTComputationForSoundness(t *testing.T) {
	// t = ceil(256 / floor(log2(E))), computed via integer arithmetic.
	tests := []struct {
		name     string
		exponent int
		wantT    int
	}{
		{"E=65537 standard", 65537, 16}, // ceil(256 / 16) = 16
		{"E=257", 257, 32},              // ceil(256 / 8) = 32
		{"E=17", 17, 64},                // ceil(256 / 4) = 64
		{"E=5", 5, 128},                 // ceil(256 / 2) = 128
		{"E=3 small exponent", 3, 256},  // ceil(256 / 1) = 256
	}
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			key.E = tt.exponent
			sv, err := NewSignerVerifier(&key.PublicKey, 256)
			require.NoError(t, err)
			svImpl, ok := sv.(*signerVerifier)
			require.True(t, ok, "expected *signerVerifier concrete type, got %T", sv)
			require.Equal(t, tt.wantT, svImpl.t, "NewSignerVerifier(E=%d, securityParam=256).t", tt.exponent)
		})
	}
}

func TestSignVerifyJWTWithExponentE3(t *testing.T) {
	privKey := generateTestRSAKeyWithExponent(t, 2048, 3)

	idToken, err := createOIDCToken(privKey, "test")
	require.NoError(t, err)

	sv, err := NewSignerVerifier(&privKey.PublicKey, 256)
	require.NoError(t, err)

	gqToken, err := sv.SignJWT(idToken)
	require.NoError(t, err)

	ok := sv.VerifyJWT(gqToken)
	require.True(t, ok, "GQ sign+verify with E=3 failed")
}

// generateTestRSAKeyWithExponent generates an RSA key with the given public exponent.
// rsa.GenerateKey always uses 65537, so this constructs the key manually.
// e must be odd and >= 3; even or small values will never satisfy gcd(e, phi)=1 and loop forever.
// Intended for generating test keys for unittests.	DO NOT call this function if security is desired!
func generateTestRSAKeyWithExponent(t *testing.T, bits, e int) *rsa.PrivateKey {
	t.Helper()
	if e < 3 || e%2 == 0 {
		t.Fatalf("generateRSAKeyWithExponent: e=%d is invalid; must be odd and >= 3", e)
	}
	// DO NOT use the code below in settings where security is desired.
	// This code is for generating test data, not secure RSA keys
	bigE := big.NewInt(int64(e))
	one := big.NewInt(1)
	for {
		p, err := rand.Prime(rand.Reader, bits/2)
		require.NoError(t, err)
		q, err := rand.Prime(rand.Reader, bits/2)
		require.NoError(t, err)

		pMinus1 := new(big.Int).Sub(p, one)
		qMinus1 := new(big.Int).Sub(q, one)
		phi := new(big.Int).Mul(pMinus1, qMinus1)

		gcd := new(big.Int).GCD(nil, nil, bigE, phi)
		if gcd.Cmp(one) != 0 {
			continue
		}

		n := new(big.Int).Mul(p, q)
		d := new(big.Int).ModInverse(bigE, phi)

		key := &rsa.PrivateKey{
			PublicKey: rsa.PublicKey{N: n, E: e},
			D:         d,
			Primes:    []*big.Int{p, q},
		}
		key.Precompute()
		return key
	}
}

func modifyTokenPayload(token []byte, audience string) ([]byte, error) {
	headers, _, signature, err := jws.SplitCompact(token)
	if err != nil {
		return nil, err
	}
	newPayload := map[string]any{
		"sub": "1",
		"iss": "test",
		"aud": audience,
		"iat": time.Now().Unix(),
	}
	modifiedPayload, err := json.Marshal(newPayload)
	if err != nil {
		return nil, err
	}
	newToken := util.JoinJWTSegments(headers, util.Base64EncodeForJWT(modifiedPayload), signature)
	return newToken, nil
}

func modifyTokenProtectedHeaderClaim(token []byte, claimKey string, newValue string) ([]byte, error) {
	headersB64, payload, signature, err := jws.SplitCompact(token)
	if err != nil {
		return nil, err
	}

	headersJSON, err := util.Base64DecodeForJWT(headersB64)
	if err != nil {
		return nil, err
	}

	var headers map[string]any
	if err := json.Unmarshal(headersJSON, &headers); err != nil {
		return nil, err
	}
	headers[claimKey] = newValue

	modifiedHeadersJSON, err := json.Marshal(headers)
	if err != nil {
		return nil, err
	}

	newToken := util.JoinJWTSegments(util.Base64EncodeForJWT(modifiedHeadersJSON), payload, signature)
	return newToken, nil
}

func createOIDCTokenPS256(oidcPrivKey *rsa.PrivateKey, audience string) ([]byte, error) {
	alg := jwa.PS256() // RSASSA-PSS using SHA-256

	oidcHeader := jws.NewHeaders()
	err := oidcHeader.Set(jws.AlgorithmKey, alg)
	if err != nil {
		return nil, err
	}
	err = oidcHeader.Set(jws.TypeKey, "JWT")
	if err != nil {
		return nil, err
	}

	oidcPayload := map[string]any{
		"sub": "1",
		"iss": "test",
		"aud": audience,
		"iat": time.Now().Unix(),
	}
	payloadBytes, err := json.Marshal(oidcPayload)
	if err != nil {
		return nil, err
	}

	return jws.Sign(
		payloadBytes,
		jws.WithKey(
			alg,
			oidcPrivKey,
			jws.WithProtectedHeaders(oidcHeader),
		),
	)
}

func createOIDCToken(oidcPrivKey *rsa.PrivateKey, audience string) ([]byte, error) {
	alg := jwa.RS256() // RSASSA-PKCS-v1.5 using SHA-256

	oidcHeader := jws.NewHeaders()
	err := oidcHeader.Set(jws.AlgorithmKey, alg)
	if err != nil {
		return nil, err
	}
	err = oidcHeader.Set(jws.TypeKey, "JWT")
	if err != nil {
		return nil, err
	}

	oidcPayload := map[string]any{
		"sub": "1",
		"iss": "test",
		"aud": audience,
		"iat": time.Now().Unix(),
	}
	payloadBytes, err := json.Marshal(oidcPayload)
	if err != nil {
		return nil, err
	}

	return jws.Sign(
		payloadBytes,
		jws.WithKey(
			alg,
			oidcPrivKey,
			jws.WithProtectedHeaders(oidcHeader),
		),
	)
}

func getClaimInProtected(claimKey string, token []byte) (any, bool, error) {
	headersB64, _, _, err := jws.SplitCompact(token)
	if err != nil {
		return nil, false, err
	}

	headersJson, err := util.Base64DecodeForJWT(headersB64)
	if err != nil {
		return nil, false, err
	}

	headers := jws.NewHeaders()
	err = json.Unmarshal(headersJson, &headers)
	if err != nil {
		return nil, false, err
	}

	var claimValue string
	err = headers.Get(claimKey, &claimValue)
	if err != nil {
		// swallow error and communicate via bool instead
		return nil, false, nil //nolint:nilerr // error is intentionally discarded
	}
	return claimValue, true, nil
}
