package protocol

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/binary"
	"math/big"
	"testing"

	"github.com/google/go-tpm/legacy/tpm2"
	"github.com/stretchr/testify/assert"

	"github.com/go-webauthn/webauthn/protocol/webauthncbor"
	"github.com/go-webauthn/webauthn/protocol/webauthncose"
)

func TestTPMAttestationVerificationSuccess(t *testing.T) {
	for i := range testAttestationTPMResponses {
		t.Run("TPM Positive tests", func(t *testing.T) {
			pcc := attestationTestUnpackResponse(t, testAttestationTPMResponses[i])
			clientDataHash := sha256.Sum256(pcc.Raw.AttestationResponse.ClientDataJSON)

			attestationType, _, err := verifyTPMFormat(pcc.Response.AttestationObject, clientDataHash[:])
			if err != nil {
				t.Fatalf("Not valid: %+v", err)
			}
			assert.Equal(t, "attca", attestationType)
		})
	}
}

var testAttestationTPMResponses = []string{
	// TPM attestation with ECC P256.
	`{
		"id": "hsS2ywFz_LWf9-lC35vC9uJTVD3ZCVdweZvESUbjXnQ",
		"rawId": "hsS2ywFz_LWf9-lC35vC9uJTVD3ZCVdweZvESUbjXnQ",
		"type": "public-key",
		"response": {
				"attestationObject": " ",
				"clientDataJSON": "eyJ0eXBlIjoid2ViYXV0aG4uY3JlYXRlIiwiY2hhbGxlbmdlIjoidXpuOXUwVHgtTEJkdEdnRVJzYmtIUkJqaVV0NWkycnZtMkJCVFpyV3FFbyIsIm9yaWdpbiI6Imh0dHBzOi8vd2ViYXV0aG4uaW8iLCJjcm9zc09yaWdpbiI6ZmFsc2V9"
		}
	}`,
	// TPM attestation with RSA SHA1.
	`{
		"rawId": "UJDoUJoGiDQF_EEZ3G_z9Lfq16_KFaXtMTjwTUrrRlc",
		"id": "UJDoUJoGiDQF_EEZ3G_z9Lfq16_KFaXtMTjwTUrrRlc",
		"response": {
			"clientDataJSON": "eyJvcmlnaW4iOiJodHRwczovL2xvY2FsaG9zdDo0NDMyOSIsImNoYWxsZW5nZSI6IjlKeVVmSmtnOFBxb0tadUQ3Rkh6T0U5ZGJ5Y3VsQzl1ckdUcEdxQm5Fd25oS21uaTRyR1JYeG0zLVpCSEs4eDZyaUpRcUlwQzhxRWEtVDBxSUZUS1RRIiwidHlwZSI6IndlYmF1dGhuLmNyZWF0ZSJ9",
			"attestationObject": " "
		},
		"type": "public-key"
	}`,
	// TPM attestation with RSA SHA256.
	`{
		"rawId": "h9XMhkVePN1Prq9Ks_VfwIsVZvt-jmSRTEnevTc-KB8",
		"id": "h9XMhkVePN1Prq9Ks_VfwIsVZvt-jmSRTEnevTc-KB8",
		"response": {
			"clientDataJSON": "eyJvcmlnaW4iOiJodHRwczovL2xvY2FsaG9zdDo0NDMyOSIsImNoYWxsZW5nZSI6ImdIckFrNHBOZTJWbEIwSExlS2NsSTJQNlFFYTgzUHVHZWlqVEhNdHBiaFk5S2x5YnlobHdGX1Z6UmU3eWhhYlhhZ1d1WTZya0RXZnZ2aE5xZ2gybzdBIiwidHlwZSI6IndlYmF1dGhuLmNyZWF0ZSJ9",
			"attestationObject": " "
		},
		"type": "public-key"
	}`,
}

func TestTPMAttestationVerificationFailAttStatement(t *testing.T) {
	tests := []struct {
		name    string
		att     AttestationObject
		wantErr string
	}{
		{
			"TPM Negative Test AttStatement Missing Ver",
			AttestationObject{},
			"Error retrieving ver value",
		},
		{
			"TPM Negative Test AttStatement Ver not 2.0",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "foo.bar"}},
			"WebAuthn only supports TPM 2.0 currently",
		},
		{
			"TPM Negative Test AttStatement Alg not present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0"}},
			"Error retrieving alg value",
		},
		{
			"TPM Negative Test AttStatement x5c not present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0)}},
			ErrNotImplemented.Details,
		},
		{
			"TPM Negative Test AttStatement ecdaaKeyId present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0), "x5c": []interface{}{}, "ecdaaKeyId": []byte{}}},
			ErrNotImplemented.Details,
		},
		{
			"TPM Negative Test AttStatement sig not present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0), "x5c": []interface{}{}}},
			"Error retrieving sig value",
		},
		{
			"TPM Negative Test AttStatement certInfo not present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0), "x5c": []interface{}{}, "sig": []byte{}}},
			"Error retrieving certInfo value",
		},
		{
			"TPM Negative Test AttStatement pubArea not present",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0), "x5c": []interface{}{}, "sig": []byte{}, "certInfo": []byte{}}},
			"Error retrieving pubArea value",
		},
		{
			"TPM Negative Test pubArea not TPMT_PUBLIC",
			AttestationObject{AttStatement: defaultAttStatement},
			"Unable to decode TPMT_PUBLIC in attestation statement",
		},
		{
			"TPM Negative Test Unsupported Public Key Type",
			AttestationObject{AttStatement: map[string]interface{}{"ver": "2.0", "alg": int64(0), "x5c": []interface{}{}, "sig": []byte{}, "certInfo": []byte{}, "pubArea": makeDefaultRSAPublicBytes()}, AuthData: AuthenticatorData{AttData: AttestedCredentialData{CredentialPublicKey: []byte{}}}},
			"Unsupported Public Key Type",
		},
	}
	for _, tt := range tests {
		attestationType, _, err := verifyTPMFormat(tt.att, nil)
		if tt.wantErr != "" {
			assert.Contains(t, err.Error(), tt.wantErr)
		} else {
			assert.Equal(t, "attca", attestationType)
		}
	}
}

func makeDefaultRSAPublicBytes() []byte {
	r, _ := defaultRSAPublic.Encode()
	return r
}

var (
	defaultRSAPublic = tpm2.Public{
		Type:       tpm2.AlgRSA,
		NameAlg:    tpm2.AlgSHA256,
		Attributes: tpm2.FlagSignerDefault,
		RSAParameters: &tpm2.RSAParams{
			Sign: &tpm2.SigScheme{
				Alg:  tpm2.AlgRSASSA,
				Hash: tpm2.AlgSHA256,
			},
			KeyBits: 2048,
		},
	}

	defaultECCPublic = tpm2.Public{
		Type:       tpm2.AlgECC,
		NameAlg:    tpm2.AlgSHA256,
		Attributes: tpm2.FlagSignerDefault,
		ECCParameters: &tpm2.ECCParams{
			Sign: &tpm2.SigScheme{
				Alg:  tpm2.AlgECDSA,
				Hash: tpm2.AlgSHA256,
			},
			CurveID: tpm2.CurveNISTP256,
		},
	}
)

var defaultAttStatement = map[string]interface{}{"ver": "2.0", "alg": int64(-257), "x5c": []interface{}{}, "sig": []byte{}, "certInfo": []byte{}, "pubArea": []byte{}}

type CredentialPublicKey struct {
	KeyType   int64  `cbor:"1,keyasint" json:"kty"`
	Algorithm int64  `cbor:"3,keyasint" json:"alg"`
	Curve     int64  `cbor:"-1,keyasint,omitempty" json:"crv"`
	XCoord    []byte `cbor:"-2,keyasint,omitempty" json:"x"`
	YCoord    []byte `cbor:"-3,keyasint,omitempty" json:"y"`
	Modulus   []byte `cbor:"-1,keyasint,omitempty" json:"n"`
	Exponent  []byte `cbor:"-2,keyasint,omitempty" json:"e"`
}

func corruptBytes(in []byte) []byte {
	out := make([]byte, len(in))

	copy(out, in)

	out[len(in)-1] ^= 0xff

	return out
}

func uint32ToBytes(i uint32) []byte {
	t, o := make([]byte, 4), make([]byte, 3)

	binary.LittleEndian.PutUint32(t, i)

	copy(o, t)

	return o
}

func getTPMAttestionKeys() ([]byte, []byte, []byte, rsa.PrivateKey, ecdsa.PrivateKey, error) {
	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, nil, rsa.PrivateKey{}, ecdsa.PrivateKey{}, err
	}

	r := webauthncose.RSAPublicKeyData{
		PublicKeyData: webauthncose.PublicKeyData{
			KeyType:   int64(webauthncose.RSAKey),
			Algorithm: int64(webauthncose.AlgRS256),
		},
		Modulus:  rsaKey.N.Bytes(),
		Exponent: uint32ToBytes(uint32(rsaKey.E)),
	}

	eccKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, nil, rsa.PrivateKey{}, ecdsa.PrivateKey{}, err
	}

	c := CredentialPublicKey{
		KeyType:   int64(webauthncose.EllipticKey),
		Algorithm: int64(webauthncose.AlgES256),
		Curve:     int64(webauthncose.P256),
		XCoord:    eccKey.X.Bytes(),
		YCoord:    eccKey.Y.Bytes(),
	}

	e := webauthncose.EC2PublicKeyData{
		PublicKeyData: webauthncose.PublicKeyData{
			KeyType:   c.KeyType,
			Algorithm: c.Algorithm,
		},
		Curve:  c.Curve,
		XCoord: c.XCoord,
		YCoord: c.YCoord,
	}

	okpKey, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, nil, rsa.PrivateKey{}, ecdsa.PrivateKey{}, err
	}

	o := webauthncose.OKPPublicKeyData{
		PublicKeyData: webauthncose.PublicKeyData{
			KeyType:   int64(webauthncose.OctetKey),
			Algorithm: int64(webauthncose.AlgEdDSA),
		},
		Curve:  int64(webauthncose.Ed25519),
		XCoord: okpKey,
	}

	epk, err := webauthncbor.Marshal(e)
	if err != nil {
		return nil, nil, nil, rsa.PrivateKey{}, ecdsa.PrivateKey{}, err
	}

	rpk, err := webauthncbor.Marshal(r)
	if err != nil {
		return nil, nil, nil, rsa.PrivateKey{}, ecdsa.PrivateKey{}, err
	}

	opk, err := webauthncbor.Marshal(o)

	return epk, rpk, opk, *rsaKey, *eccKey, err
}

func TestTPMAttestationVerificationFailPubArea(t *testing.T) {
	epk, rpk, opk, rsaKey, eccKey, err := getTPMAttestionKeys()
	if err != nil {
		t.Fatal(err)
	}

	tests := []struct {
		name      string
		keyType   webauthncose.COSEKeyType
		rsaParams tpm2.RSAParams
		eccParams tpm2.ECCParams
		cpk       []byte
		wantErr   string
	}{
		{
			"TPM Negative Test pubArea curve mismatch",
			webauthncose.EllipticKey,
			tpm2.RSAParams{},
			tpm2.ECCParams{CurveID: tpm2.CurveNISTP224, Point: tpm2.ECPoint{XRaw: eccKey.X.Bytes(), YRaw: eccKey.Y.Bytes()}},
			epk,
			"Mismatch between ECCParameters in pubArea and credentialPublicKey",
		},
		{
			"TPM Negative Test pubArea X mismatch",
			webauthncose.EllipticKey,
			tpm2.RSAParams{},
			tpm2.ECCParams{CurveID: tpm2.CurveNISTP256, Point: tpm2.ECPoint{XRaw: corruptBytes(eccKey.X.Bytes()), YRaw: eccKey.Y.Bytes()}},
			epk,
			"Mismatch between ECCParameters in pubArea and credentialPublicKey",
		},
		{
			"TPM Negative Test pubArea Y mismatch",
			webauthncose.EllipticKey,
			tpm2.RSAParams{},
			tpm2.ECCParams{CurveID: tpm2.CurveNISTP256, Point: tpm2.ECPoint{XRaw: eccKey.X.Bytes(), YRaw: corruptBytes(eccKey.Y.Bytes())}},
			epk,
			"Mismatch between ECCParameters in pubArea and credentialPublicKey",
		},
		{
			"TPM Negative Test pubArea N mismatch",
			webauthncose.RSAKey,
			tpm2.RSAParams{ModulusRaw: corruptBytes(rsaKey.N.Bytes()), ExponentRaw: uint32(rsaKey.E)},
			tpm2.ECCParams{},
			rpk,
			"Mismatch between RSAParameters in pubArea and credentialPublicKey",
		},
		{
			"TPM Negative Test pubArea E mismatch",
			webauthncose.RSAKey,
			tpm2.RSAParams{ModulusRaw: rsaKey.N.Bytes(), ExponentRaw: uint32(rsaKey.E + 1)},
			tpm2.ECCParams{},
			rpk,
			"Mismatch between RSAParameters in pubArea and credentialPublicKey",
		},
		{
			"TPM Negative Test pubArea unsupported key type",
			webauthncose.OctetKey,
			tpm2.RSAParams{},
			tpm2.ECCParams{},
			opk,
			"Unsupported Public Key Type",
		},
	}
	for _, tt := range tests {
		attStmt := make(map[string]interface{}, len(defaultAttStatement))
		for id, v := range defaultAttStatement {
			attStmt[id] = v
		}

		public := tpm2.Public{}

		switch tt.keyType {
		case webauthncose.EllipticKey:
			public = defaultECCPublic
			public.ECCParameters.CurveID = tt.eccParams.CurveID
			public.ECCParameters.Point.XRaw = tt.eccParams.Point.XRaw
			public.ECCParameters.Point.YRaw = tt.eccParams.Point.YRaw
		case webauthncose.RSAKey:
			public = defaultRSAPublic
			public.RSAParameters.ExponentRaw = tt.rsaParams.ExponentRaw
			public.RSAParameters.ModulusRaw = tt.rsaParams.ModulusRaw
		case webauthncose.OctetKey:
			public = defaultECCPublic
		}

		attStmt["pubArea"], _ = public.Encode()
		att := AttestationObject{
			AttStatement: attStmt,
			AuthData: AuthenticatorData{
				AttData: AttestedCredentialData{
					CredentialPublicKey: tt.cpk,
				},
			},
		}

		attestationType, _, err := verifyTPMFormat(att, nil)
		if tt.wantErr != "" {
			assert.Contains(t, err.Error(), tt.wantErr)
		} else {
			assert.Equal(t, "attca", attestationType)
		}
	}
}

func TestTPMAttestationVerificationFailCertInfo(t *testing.T) {
	attStmt := make(map[string]interface{}, len(defaultAttStatement))

	for id, v := range defaultAttStatement {
		attStmt[id] = v
	}

	rsaKey, _ := rsa.GenerateKey(rand.Reader, 2048)

	r := webauthncose.RSAPublicKeyData{
		PublicKeyData: webauthncose.PublicKeyData{
			KeyType:   int64(webauthncose.RSAKey),
			Algorithm: int64(webauthncose.AlgRS256),
		},
		Modulus:  rsaKey.N.Bytes(),
		Exponent: uint32ToBytes(uint32(rsaKey.E)),
	}

	public := defaultRSAPublic
	public.RSAParameters.ExponentRaw = uint32(rsaKey.E)
	public.RSAParameters.ModulusRaw = rsaKey.N.Bytes()
	attStmt["pubArea"], _ = public.Encode()
	rpk, _ := webauthncbor.Marshal(r)
	att := AttestationObject{
		AttStatement: attStmt,
		AuthData: AuthenticatorData{
			AttData: AttestedCredentialData{
				CredentialPublicKey: rpk,
			},
		},
	}

	f := webauthncose.HasherFromCOSEAlg(webauthncose.AlgRS256)
	h := f()
	h.Write(att.RawAuthData)
	extraData := h.Sum(nil)

	tests := []struct {
		name     string
		certInfo tpm2.AttestationData
		wantErr  string
	}{
		{
			"TPM Negative Test CertInfo invalid certInfo",
			tpm2.AttestationData{},
			"decoding Magic/Type: EOF",
		},
		{
			"TPM Negative Test CertInfo magic value not 0xff544347",
			tpm2.AttestationData{Magic: 42, Type: tpm2.TagAttestCreation, AttestedCreationInfo: &tpm2.CreationInfo{}},
			"incorrect magic value: 2a",
		},
		{
			"TPM Negative Test CertInfo type not TagAttestCertify",
			tpm2.AttestationData{Magic: 0xff544347, Type: tpm2.TagAttestCreation, AttestedCreationInfo: &tpm2.CreationInfo{}},
			"Type is not set to TPM_ST_ATTEST_CERTIFY",
		},
		{
			"TPM Negative Test CertInfo type not TagAttestCertify",
			tpm2.AttestationData{Magic: 0xff544347, Type: tpm2.TagAttestCertify, AttestedCertifyInfo: &tpm2.CertifyInfo{}},
			"ExtraData is not set to hash of attToBeSigned",
		},
		{
			"TPM Negative Test CertInfo Name/pubArea mismatch",
			tpm2.AttestationData{Magic: 0xff544347, Type: tpm2.TagAttestCertify, AttestedCertifyInfo: &tpm2.CertifyInfo{Name: tpm2.Name{Digest: nil}}, ExtraData: extraData},
			"Name doesn't have a Digest, can't compare to Public",
		},
		{
			"TPM Negative Test CertInfo Name/pubArea mismatch",
			tpm2.AttestationData{Magic: 0xff544347, Type: tpm2.TagAttestCertify, AttestedCertifyInfo: &tpm2.CertifyInfo{Name: tpm2.Name{Digest: &tpm2.HashValue{Alg: tpm2.AlgSHA256, Value: extraData}}}, ExtraData: h.Sum(nil)},
			"Hash value mismatch attested and pubArea",
		},
	}
	for _, tt := range tests {
		certInfo, err := tt.certInfo.Encode()
		if tt.certInfo.Magic != 0 && err != nil {
			t.Fatal(err)
		}

		att.AttStatement["certInfo"] = certInfo
		attestationType, _, err := verifyTPMFormat(att, nil)

		if tt.wantErr != "" {
			assert.Contains(t, err.Error(), tt.wantErr)
		} else {
			assert.Equal(t, "attca", attestationType)
		}
	}
}

var (
	x5cTemplate = x509.Certificate{
		SerialNumber: big.NewInt(0),
		Version:      3,
		Subject: pkix.Name{
			Country:            []string{"US"},
			Organization:       []string{"Duo Labs"},
			OrganizationalUnit: []string{"Authenticator Attestation"},
			CommonName:         "WebAuthn.io",
		},
		Extensions: []pkix.Extension{
			{
				Id:       asn1.ObjectIdentifier([]int{1, 3, 6, 1, 4, 1, 45724, 1, 1, 4}),
				Critical: false,
				Value: []byte{
					0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
					0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
				},
			},
		},
		IsCA: false,
	}
)

func TestTPMAttestationVerificationFailX5c(t *testing.T) {
	attStmt := make(map[string]interface{}, len(defaultAttStatement))

	for id, v := range defaultAttStatement {
		attStmt[id] = v
	}

	rsaKey, _ := rsa.GenerateKey(rand.Reader, 2048)

	r := webauthncose.RSAPublicKeyData{
		PublicKeyData: webauthncose.PublicKeyData{
			KeyType:   int64(webauthncose.RSAKey),
			Algorithm: int64(webauthncose.AlgRS256),
		},
		Modulus:  rsaKey.N.Bytes(),
		Exponent: uint32ToBytes(uint32(rsaKey.E)),
	}

	public := defaultRSAPublic
	public.RSAParameters.ExponentRaw = uint32(rsaKey.E)
	public.RSAParameters.ModulusRaw = rsaKey.N.Bytes()
	pubBytes, _ := public.Encode()
	attStmt["pubArea"] = pubBytes
	rpk, _ := webauthncbor.Marshal(r)
	att := AttestationObject{
		AttStatement: attStmt,
		AuthData: AuthenticatorData{
			AttData: AttestedCredentialData{
				CredentialPublicKey: rpk,
			},
		},
	}

	f := webauthncose.HasherFromCOSEAlg(webauthncose.AlgRS256)
	h := f()
	h.Write(att.RawAuthData)
	extraData := h.Sum(nil)

	h.Reset()
	h.Write(pubBytes)
	pubName := h.Sum(nil)

	certInfo := tpm2.AttestationData{
		Magic: 0xff544347,
		Type:  tpm2.TagAttestCertify,
		AttestedCertifyInfo: &tpm2.CertifyInfo{
			Name: tpm2.Name{
				Digest: &tpm2.HashValue{
					Alg:   tpm2.AlgSHA256,
					Value: pubName,
				},
			},
		},
		ExtraData: extraData,
	}
	attStmt["certInfo"], _ = certInfo.Encode()

	makeX5c := func(b []byte) []interface{} {
		q := make([]interface{}, 1)
		q[0] = b

		return q
	}

	tests := []struct {
		name    string
		x5c     []interface{}
		wantErr string
	}{
		{
			"TPM Negative Test x5c empty",
			make([]interface{}, 1),
			"Error getting certificate from x5c cert chain",
		},
		{
			"TPM Negative Test x5c can't parse",
			makeX5c(make([]byte, 1)),
			"Error parsing certificate from ASN.1",
		},
	}

	for _, tt := range tests {
		att.AttStatement["x5c"] = tt.x5c
		attestationType, _, err := verifyTPMFormat(att, nil)

		if tt.wantErr != "" {
			assert.Contains(t, err.Error(), tt.wantErr)
		} else {
			assert.Equal(t, "attca", attestationType)
		}
	}
}
