package azure

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

	"github.com/google/go-tpm-tools/simulator"
	"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:    "48434c41010000001c070000020000000000000000000000000000000000000002000000020000001f0003000000000001000000000000000000000000000000020000000000000000000000000000000000000001000000020000000000065d010000000000000000000000000000000ccc0895ef2f2c3b8c8568f5a2bb65ff5bf9387a09359742ad41e686cacfd38b00000000000000000000000000000000000000000000000000000000000000005677f1de87289e7ad2c7e99c805d0468b1a9ccd83f0d245afa5242d405da4d5725852f8c6550564870e5f3206dfb1841000000000000000000000000000000000000000000000000000000000000000057e229e0ffe5fa92d0faddff6cae0e61c926fc9ef9afd20a8b8cfcf7129db9338cbe5bf3f6987733a2bf65d06dc38fc100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005f7240b24a1babe2ece844c4f792bcd9844bf6907d14aeea00156310b9538daffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff020000000000065d0000000000000000000000000000000000000000000000009e44aaef02cfca6fddbaca669c6cfd29e1ab8d97ebc939857128acbb13b8740df31436d34e86e5f8ae0cdfeb3a0e185db46decac176cc77d761c22a1b9dcf25b020000000000065d0133010001330100020000000000065d000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000bcb7dc15abff884802e774b39adba8e6ff7efcf05e115c91588e657065151056a320f70c788d0e3619391052922e422b000000000000000000000000000000000000000000000000e8dbf581140443bbc681c50eca8639a76ef6cab34e0780cbca977e2e2a03f8b864fd4e9774b0f8055511567e031e59bf00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005c020000010000000200000001000000480200007b226b657973223a5b7b226b6964223a2248434c416b507562222c226b65795f6f7073223a5b22656e6372797074225d2c226b7479223a22525341222c2265223a2241514142222c226e223a22747946717641414166324746656c6b5737566352684a6e4132597659364c6a427a65554e3276614d5a6e5a74685f74466e574d6b4b35415874757379434e656c337569703356475a7a54617a3558327447566a4772732d4d56486361703951647771555856573367394f515f74456269786378372d78626c554a516b474551666e626253646e5049326c764c7a4f73315a5f30766a65444178765351726d616773366e592d634a4157482d706744564a79487470735553735f5142576b6c617a44736f3557486d6e4d743973394d75696c57586f7830525379586e55656151796859316a753752545363526e5658754e7936377a5f454a6e774d393264727746623841556430534a5f396f687645596c34615a52444543476f3056726a635348552d4a474a6575574335566844425235454f6f4356424267716539653833765f6c4a784933574c65326f7653495a49497a416d625351227d5d2c22766d2d636f6e66696775726174696f6e223a7b22636f6e736f6c652d656e61626c6564223a747275652c2263757272656e742d74696d65223a313636313435353339312c227365637572652d626f6f74223a66616c73652c2274706d2d656e61626c6564223a747275652c22766d556e697175654964223a2242364339384333422d344543372d344441362d424432462d374439384432304437423735227d7d000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
			vcek:      "someVCEKCert",
			certChain: "ASKARK",
			wantErr:   false,
		},
		"report too short": {
			report:  "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
			wantErr: true,
		},
		"vcek api fails": {
			report:   "48434c41010000001c070000020000000000000000000000000000000000000002000000020000001f0003000000000001000000000000000000000000000000020000000000000000000000000000000000000001000000020000000000065d010000000000000000000000000000000ccc0895ef2f2c3b8c8568f5a2bb65ff5bf9387a09359742ad41e686cacfd38b00000000000000000000000000000000000000000000000000000000000000005677f1de87289e7ad2c7e99c805d0468b1a9ccd83f0d245afa5242d405da4d5725852f8c6550564870e5f3206dfb1841000000000000000000000000000000000000000000000000000000000000000057e229e0ffe5fa92d0faddff6cae0e61c926fc9ef9afd20a8b8cfcf7129db9338cbe5bf3f6987733a2bf65d06dc38fc100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005f7240b24a1babe2ece844c4f792bcd9844bf6907d14aeea00156310b9538daffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff020000000000065d0000000000000000000000000000000000000000000000009e44aaef02cfca6fddbaca669c6cfd29e1ab8d97ebc939857128acbb13b8740df31436d34e86e5f8ae0cdfeb3a0e185db46decac176cc77d761c22a1b9dcf25b020000000000065d0133010001330100020000000000065d000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000bcb7dc15abff884802e774b39adba8e6ff7efcf05e115c91588e657065151056a320f70c788d0e3619391052922e422b000000000000000000000000000000000000000000000000e8dbf581140443bbc681c50eca8639a76ef6cab34e0780cbca977e2e2a03f8b864fd4e9774b0f8055511567e031e59bf00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005c020000010000000200000001000000480200007b226b657973223a5b7b226b6964223a2248434c416b507562222c226b65795f6f7073223a5b22656e6372797074225d2c226b7479223a22525341222c2265223a2241514142222c226e223a22747946717641414166324746656c6b5737566352684a6e4132597659364c6a427a65554e3276614d5a6e5a74685f74466e574d6b4b35415874757379434e656c337569703356475a7a54617a3558327447566a4772732d4d56486361703951647771555856573367394f515f74456269786378372d78626c554a516b474551666e626253646e5049326c764c7a4f73315a5f30766a65444178765351726d616773366e592d634a4157482d706744564a79487470735553735f5142576b6c617a44736f3557486d6e4d743973394d75696c57586f7830525379586e55656151796859316a753752545363526e5658754e7936377a5f454a6e774d393264727746623841556430534a5f396f687645596c34615a52444543476f3056726a635348552d4a474a6575574335566844425235454f6f4356424267716539653833765f6c4a784933574c65326f7653495a49497a416d625351227d5d2c22766d2d636f6e66696775726174696f6e223a7b22636f6e736f6c652d656e61626c6564223a747275652c2263757272656e742d74696d65223a313636313435353339312c227365637572652d626f6f74223a66616c73652c2274706d2d656e61626c6564223a747275652c22766d556e697175654964223a2242364339384333422d344543372d344441362d424432462d374439384432304437423735227d7d000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
			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 := getSNPAttestation(&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) {
	assert := assert.New(t)

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

	_, err = getHCLAttestationKey(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
}