/*
Copyright (c) Edgeless Systems GmbH

SPDX-License-Identifier: AGPL-3.0-only
*/

package snp

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"testing"

	tpmclient "github.com/google/go-tpm-tools/client"
	"github.com/google/go-tpm-tools/simulator"
	"github.com/google/go-tpm/tpm2"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGetSNPAttestation(t *testing.T) {
	testCases := map[string]struct {
		report    string
		vcek      string
		certChain string
		apiError  error
		wantErr   bool
	}{
		"success": {
			report:    "[encoded data, 4096 characters]",
			vcek:      "someVCEKCert",
			certChain: "ASKARK",
			wantErr:   false,
		},
		"report too short": {
			report:  "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
			wantErr: true,
		},
		"vcek api fails": {
			report:   "[encoded data, 4096 characters]",
			apiError: errors.New(""),
			wantErr:  true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			assert := assert.New(t)
			require := require.New(t)

			tpm, err := simulator.Get()
			require.NoError(err)
			defer tpm.Close()

			imdsClient := stubImdsClient{
				vcek:      tc.vcek,
				certChain: tc.certChain,
				apiError:  tc.apiError,
			}
			reportDecoded, err := hex.DecodeString(tc.report)
			assert.NoError(err)
			snpAttestationReport := stubTpmReport{
				tpmContent: reportDecoded,
				err:        nil,
			}

			attestationJSON, err := getInstanceInfo(&snpAttestationReport, imdsClient)(tpm)
			if tc.wantErr {
				assert.Error(err)
				return
			}
			assert.NoError(err)

			var instanceInfo azureInstanceInfo
			err = json.Unmarshal(attestationJSON, &instanceInfo)
			assert.NoError(err)

			if tc.wantErr {
				assert.NotEqualValues(tc.report[0x20:], instanceInfo.AttestationReport)
			} else {
				assert.EqualValues(reportDecoded[0x20:0x4a0+0x20], instanceInfo.AttestationReport)
				assert.Equal(tc.vcek, string(instanceInfo.Vcek))
				assert.Equal(tc.certChain, string(instanceInfo.CertChain))
			}
		})
	}
}

// TestGetHCLAttestationKey is a basic smoke test that only checks if getAkPub can be run error free.
// Testing anything else will only verify that the simulator works as expected, since getAkPub
// only retrieves the attestation key from the TPM.
func TestGetHCLAttestationKey(t *testing.T) {
	require := require.New(t)
	assert := assert.New(t)

	tpm, err := simulator.Get()
	require.NoError(err)
	defer tpm.Close()

	// we should receive an error if no key was saved at index `tpmAkIdx`
	_, err = getAttestationKey(tpm)
	assert.Error(err)

	// create a key at the index
	tpmAk, err := tpmclient.NewCachedKey(tpm, tpm2.HandleOwner, tpm2.Public{
		Type:       tpm2.AlgRSA,
		NameAlg:    tpm2.AlgSHA256,
		Attributes: tpm2.FlagFixedTPM | tpm2.FlagFixedParent | tpm2.FlagSensitiveDataOrigin | tpm2.FlagUserWithAuth | tpm2.FlagNoDA | tpm2.FlagRestricted | tpm2.FlagSign,
		RSAParameters: &tpm2.RSAParams{
			Sign: &tpm2.SigScheme{
				Alg:  tpm2.AlgRSASSA,
				Hash: tpm2.AlgSHA256,
			},
			KeyBits: 2048,
		},
	}, tpmAkIdx)
	require.NoError(err)
	defer tpmAk.Close()

	// we should now be able to retrieve the key
	_, err = getAttestationKey(tpm)
	assert.NoError(err)
}

type stubImdsClient struct {
	vcek      string
	certChain string
	apiError  error
}

func (c stubImdsClient) getVcek(ctx context.Context) (vcekResponse, error) {
	if c.apiError != nil {
		return vcekResponse{}, c.apiError
	}

	return vcekResponse{
		VcekCert:         c.vcek,
		CertificateChain: c.certChain,
	}, nil
}

type stubTpmReport struct {
	tpmContent []byte
	err        error
}

func (s *stubTpmReport) get(tpm io.ReadWriteCloser) ([]byte, error) {
	if s.err != nil {
		return nil, s.err
	}
	return s.tpmContent, nil
}
