/*
   Copyright Docker attest authors

   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.
*/

package attest

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/Masterminds/semver/v3"
	"github.com/distribution/reference"
	"github.com/docker/attest/attestation"
	"github.com/docker/attest/internal/test"
	"github.com/docker/attest/mapping"
	"github.com/docker/attest/oci"
	"github.com/docker/attest/policy"
	"github.com/docker/attest/tlog"
	"github.com/docker/attest/tuf"
	"github.com/docker/attest/version"
	intoto "github.com/in-toto/in-toto-golang/in_toto"
	"github.com/secure-systems-lab/go-securesystemslib/dsse"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sigs.k8s.io/yaml"
)

var (
	ExampleAttestation = filepath.Join("test", "testdata", "example_attestation.json")
	LocalKeysPolicy    = filepath.Join("test", "testdata", "local-policy-real")
	LocalParamPolicy   = filepath.Join("test", "testdata", "local-policy-param")
	ExpiresPolicy      = filepath.Join("test", "testdata", "expires")
)

const (
	LinuxAMD64          = "linux/amd64"
	TestVerifierVersion = "9.9.9"
)

type MockVersionFetcher struct{}

func (m *MockVersionFetcher) Get() (*semver.Version, error) {
	return semver.NewVersion(TestVerifierVersion)
}

func TestVerifyAttestations(t *testing.T) {
	ex, err := os.ReadFile(ExampleAttestation)
	assert.NoError(t, err)

	env := new(attestation.EnvelopeReference)
	err = json.Unmarshal(ex, env)
	assert.NoError(t, err)
	resolver := &attestation.MockResolver{
		Envs: []*attestation.EnvelopeReference{env},
	}

	testCases := []struct {
		name                  string
		policyEvaluationError error
		expectedError         error
	}{
		{"policy ok", nil, nil},
		{"policy error", fmt.Errorf("policy error"), fmt.Errorf("policy evaluation failed: policy error")},
	}
	ctx := context.Background()
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			mockPE := policy.MockPolicyEvaluator{
				EvaluateFunc: func(_ context.Context, _ attestation.Resolver, _ *policy.Policy, _ *policy.Input) (*policy.Result, error) {
					return policy.AllowedResult(), tc.policyEvaluationError
				},
			}
			verifier, err := NewImageVerifier(ctx, &policy.Options{})
			require.NoError(t, err)
			_, err = verifier.verifyAttestations(ctx, resolver, &mockPE, &policy.Policy{ResolvedName: ""})
			if tc.expectedError != nil {
				if assert.Error(t, err) {
					assert.Equal(t, tc.expectedError.Error(), err.Error())
				}
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestVSA(t *testing.T) {
	ctx, signer := test.Setup(t)
	// setup an image with signed attestations
	outputLayout := test.CreateTempDir(t, "", TestTempDir)

	opts := &attestation.SigningOptions{
		TransparencyLog: tlog.GetMockTL(),
	}
	attIdx, err := oci.IndexFromPath(test.UnsignedTestIndex())
	assert.NoError(t, err)
	signedManifests, err := SignStatements(ctx, attIdx.Index, signer, opts)
	require.NoError(t, err)
	signedIndex := attIdx.Index
	signedIndex, err = attestation.UpdateIndexImages(signedIndex, signedManifests)
	require.NoError(t, err)

	// output signed attestations
	spec, err := oci.ParseImageSpec(oci.LocalPrefix+outputLayout, oci.WithPlatform(LinuxAMD64))
	require.NoError(t, err)
	err = oci.SaveIndex(ctx, []*oci.ImageSpec{spec}, signedIndex, attIdx.Name)
	assert.NoError(t, err)

	// mocked vsa query should pass
	policyOpts := &policy.Options{
		LocalPolicyDir:   PassPolicyDir,
		AttestationStyle: mapping.AttestationStyleAttached,
		DisableTUF:       true,
	}
	verifier, err := NewImageVerifier(ctx, policyOpts)
	require.NoError(t, err)
	verifier.versionFetcher = &MockVersionFetcher{}
	results, err := verifier.Verify(ctx, spec)
	require.NoError(t, err)
	assert.Equal(t, OutcomeSuccess, results.Outcome)
	assert.Empty(t, results.Violations)

	if assert.NotNil(t, results.Input) {
		assert.Equal(t, test.UnsignedLinuxAMD64ImageDigest, results.Input.Digest)
		assert.NotNil(t, results.Input.Tag)
	}

	assert.Equal(t, intoto.StatementInTotoV01, results.VSA.Type)
	assert.Equal(t, attestation.VSAPredicateType, results.VSA.PredicateType)
	assert.Len(t, results.VSA.Subject, 1)

	require.IsType(t, attestation.VSAPredicate{}, results.VSA.Predicate)
	attestationPredicate, ok := results.VSA.Predicate.(attestation.VSAPredicate)
	require.True(t, ok)

	assert.Equal(t, "PASSED", attestationPredicate.VerificationResult)
	assert.Equal(t, "docker-official-images", attestationPredicate.Verifier.ID)
	assert.Equal(t, []string{"SLSA_BUILD_LEVEL_3"}, attestationPredicate.VerifiedLevels)
	assert.Equal(t, PassPolicyDir+"/policy.rego", attestationPredicate.Policy.DownloadLocation)
	assert.Equal(t, "https://docker.com/official/policy/v0.1", attestationPredicate.Policy.URI)
	// this is the digest of the policy file
	assert.Equal(t, map[string]string{"sha256": "fe1d4973f3521009a3adec206946e12aae935a2aceeb1e01f52b5d4cb9de79a5"}, attestationPredicate.Policy.Digest)
	assert.Greater(t, len(attestationPredicate.InputAttestations), 0)
	for _, input := range attestationPredicate.InputAttestations {
		require.NotEmpty(t, input.Digest)
		digest, ok := input.Digest["sha256"]
		assert.True(t, ok)
		assert.NotEmpty(t, digest)
		assert.Contains(t, []string{"application/vnd.in-toto.provenance+dsse", "application/vnd.in-toto.spdx+dsse"}, input.MediaType)
	}
	assert.Equal(t, TestVerifierVersion, attestationPredicate.Verifier.Version[version.ThisModulePath])
}

func TestVerificationFailure(t *testing.T) {
	ctx, signer := test.Setup(t)
	// setup an image with signed attestations
	outputLayout := test.CreateTempDir(t, "", TestTempDir)

	opts := &attestation.SigningOptions{
		TransparencyLog: tlog.GetMockTL(),
	}
	attIdx, err := oci.IndexFromPath(test.UnsignedTestIndex())
	assert.NoError(t, err)
	signedManifests, err := SignStatements(ctx, attIdx.Index, signer, opts)
	require.NoError(t, err)
	signedIndex := attIdx.Index
	signedIndex, err = attestation.UpdateIndexImages(signedIndex, signedManifests, attestation.WithReplacedLayers(true))
	require.NoError(t, err)

	// output signed attestations
	spec, err := oci.ParseImageSpec(oci.LocalPrefix+outputLayout, oci.WithPlatform(LinuxAMD64))
	require.NoError(t, err)
	err = oci.SaveIndex(ctx, []*oci.ImageSpec{spec}, signedIndex, attIdx.Name)
	assert.NoError(t, err)

	// mocked vsa query should fail
	policyOpts := &policy.Options{
		LocalPolicyDir:   FailPolicyDir,
		AttestationStyle: mapping.AttestationStyleAttached,
		DisableTUF:       true,
	}
	results, err := Verify(ctx, spec, policyOpts)
	require.NoError(t, err)
	assert.Equal(t, OutcomeFailure, results.Outcome)
	assert.Len(t, results.Violations, 1)

	violation := results.Violations[0]
	assert.Equal(t, "missing_attestation", violation.Type)
	assert.Equal(t, "Attestation missing for subject", violation.Description)
	assert.Nil(t, violation.Attestation)

	assert.Equal(t, intoto.StatementInTotoV01, results.VSA.Type)
	assert.Equal(t, attestation.VSAPredicateType, results.VSA.PredicateType)
	assert.Len(t, results.VSA.Subject, 1)

	require.IsType(t, attestation.VSAPredicate{}, results.VSA.Predicate)
	attestationPredicate, ok := results.VSA.Predicate.(attestation.VSAPredicate)
	require.True(t, ok)

	assert.Equal(t, "FAILED", attestationPredicate.VerificationResult)
	assert.Equal(t, "docker-official-images", attestationPredicate.Verifier.ID)
	assert.Equal(t, []string{"SLSA_BUILD_LEVEL_3"}, attestationPredicate.VerifiedLevels)
	assert.Equal(t, FailPolicyDir+"/policy.rego", attestationPredicate.Policy.DownloadLocation)
	assert.Equal(t, "https://docker.com/official/policy/v0.1", attestationPredicate.Policy.URI)
	assert.Equal(t, map[string]string{"sha256": "4345a4f5db3ce02664bd83f8e4aad03bd9a26d4edb334338c762d9648e16bed1"}, attestationPredicate.Policy.Digest)
}

func TestSignVerify(t *testing.T) {
	ctx, signer := test.Setup(t)
	// setup an image with signed attestations
	outputLayout := test.CreateTempDir(t, "", TestTempDir)

	keys, err := GenKeyMetadata(signer)
	require.NoError(t, err)
	config := struct {
		Keys []*attestation.KeyMetadata `json:"keys"`
	}{
		Keys: []*attestation.KeyMetadata{keys},
	}
	keysYaml, err := yaml.Marshal(config)
	require.NoError(t, err)

	testCases := []struct {
		name               string
		signTL             bool
		policyDir          string
		imageName          string
		expectedNonSuccess Outcome
		spitConfig         bool
		param              string
	}{
		{name: "happy path", signTL: true, policyDir: PassNoTLPolicyDir},
		{name: "sign tl, verify no tl", signTL: true, policyDir: PassPolicyDir},
		{name: "no tl", signTL: false, policyDir: PassPolicyDir},
		{name: "mirror", signTL: false, policyDir: PassMirrorPolicyDir, imageName: "mirror.org/library/test-image:test"},
		{name: "mirror no match", signTL: false, policyDir: PassMirrorPolicyDir, imageName: "incorrect.org/library/test-image:test", expectedNonSuccess: OutcomeNoPolicy},
		{name: "verify inputs", signTL: false, policyDir: InputsPolicyDir},
		{name: "mirror with verification", signTL: false, policyDir: LocalKeysPolicy, imageName: "mirror.org/library/test-image:test", spitConfig: true},
		{name: "policy with input params", spitConfig: true, signTL: false, policyDir: LocalParamPolicy, param: "bar"},
		{name: "policy without expected param", spitConfig: true, signTL: false, policyDir: LocalParamPolicy, param: "baz", expectedNonSuccess: OutcomeFailure},
	}

	attIdx, err := oci.IndexFromPath(test.UnsignedTestIndex())
	assert.NoError(t, err)
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			opts := &attestation.SigningOptions{}
			if tc.signTL {
				opts.TransparencyLog = tlog.GetMockTL()
			}
			if tc.spitConfig {
				// write keysYaml to config.yaml in LocalKeysPolicy.
				err = os.WriteFile(filepath.Join(tc.policyDir, "config.yaml"), keysYaml, 0o600)
				require.NoError(t, err)
			}

			signedManifests, err := SignStatements(ctx, attIdx.Index, signer, opts)
			require.NoError(t, err)
			signedIndex := attIdx.Index
			signedIndex, err = attestation.UpdateIndexImages(signedIndex, signedManifests, attestation.WithReplacedLayers(true))
			require.NoError(t, err)

			imageName := tc.imageName
			if imageName == "" {
				imageName = attIdx.Name
			}
			// output signed attestations
			spec, err := oci.ParseImageSpec(oci.LocalPrefix+outputLayout, oci.WithPlatform(LinuxAMD64))
			require.NoError(t, err)
			err = oci.SaveIndex(ctx, []*oci.ImageSpec{spec}, signedIndex, imageName)
			require.NoError(t, err)

			policyOpts := &policy.Options{
				LocalPolicyDir: tc.policyDir,
				DisableTUF:     true,
				Debug:          true,
			}
			if tc.signTL {
				getTL := func(_ context.Context, _ *attestation.VerifyOptions) (tlog.TransparencyLog, error) {
					return tlog.GetMockTL(), nil
				}
				verifier, err := attestation.NewVerfier(attestation.WithLogVerifierFactory(getTL))
				require.NoError(t, err)
				policyOpts.AttestationVerifier = verifier
			}
			if tc.param != "" {
				policyOpts.Parameters = policy.Parameters{"foo": tc.param}
			}
			results, err := Verify(ctx, spec, policyOpts)
			require.NoError(t, err)
			if tc.expectedNonSuccess != "" {
				assert.Equal(t, tc.expectedNonSuccess, results.Outcome)
				return
			}
			if results.Outcome == OutcomeFailure {
				t.Logf("Violations: %v", results.Violations)
			}
			assert.Equal(t, OutcomeSuccess, results.Outcome)
			platform, err := oci.ParsePlatform(LinuxAMD64)
			require.NoError(t, err)

			ref, err := reference.ParseNormalizedNamed(attIdx.Name)
			require.NoError(t, err)
			expectedPURL, _, err := oci.RefToPURL(ref, platform)
			require.NoError(t, err)
			assert.Equal(t, expectedPURL, results.Input.PURL)
		})
	}
}

func TestDefaultOptions(t *testing.T) {
	testCases := []struct {
		name             string
		tufOpts          *tuf.ClientOptions
		localTargetsDir  string
		attestationStyle mapping.AttestationStyle
		referrersRepo    string
		expectedError    string
		disableTuf       bool
		localPolicyDir   string
	}{
		{name: "empty"},
		{name: "tufClient provided", tufOpts: &tuf.ClientOptions{MetadataSource: "a", TargetsSource: "b"}},
		{name: "localTargetsDir provided", localTargetsDir: test.CreateTempDir(t, "", TestTempDir)},
		{name: "attestationStyle provided", attestationStyle: mapping.AttestationStyleAttached},
		{name: "referrersRepo provided", referrersRepo: "referrers"},
		{name: "referrersRepo provided with attached", referrersRepo: "referrers", attestationStyle: mapping.AttestationStyleAttached, expectedError: "referrers repo specified but attestation source not set to referrers"},
		{name: "tuf disabled and no local-policy-dir", disableTuf: true, expectedError: "local policy dir must be set if not using TUF"},
		{name: "tuf disabled but options set", disableTuf: true, tufOpts: &tuf.ClientOptions{MetadataSource: "a", TargetsSource: "b"}, localPolicyDir: "foo", expectedError: "TUF client options set but TUF disabled"},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			defaultTargets, err := defaultLocalTargetsDir()
			require.NoError(t, err)

			opts := &policy.Options{
				TUFClientOptions: tc.tufOpts,
				LocalTargetsDir:  tc.localTargetsDir,
				AttestationStyle: tc.attestationStyle,
				ReferrersRepo:    tc.referrersRepo,
				DisableTUF:       tc.disableTuf,
				LocalPolicyDir:   tc.localPolicyDir,
			}

			err = populateDefaultOptions(opts)
			if tc.expectedError != "" {
				require.Error(t, err)
				assert.Equal(t, tc.expectedError, err.Error())
				return
			}

			require.NoError(t, err)

			if tc.localTargetsDir != "" {
				assert.Equal(t, tc.localTargetsDir, opts.LocalTargetsDir)
			} else {
				assert.Equal(t, defaultTargets, opts.LocalTargetsDir)
			}

			if tc.attestationStyle != "" {
				assert.Equal(t, tc.attestationStyle, opts.AttestationStyle)
			} else {
				assert.Equal(t, mapping.AttestationStyleReferrers, opts.AttestationStyle)
			}

			if tc.tufOpts != nil {
				assert.Equal(t, tc.tufOpts, opts.TUFClientOptions)
			} else {
				assert.NotNil(t, opts.TUFClientOptions)
			}

			if tc.referrersRepo != "" {
				assert.Equal(t, tc.referrersRepo, opts.ReferrersRepo)
			} else {
				assert.Empty(t, opts.ReferrersRepo)
			}
		})
	}
}

// LoadKeyMetadata loads the key metadata for the given signer verifier.
func GenKeyMetadata(sv dsse.SignerVerifier) (*attestation.KeyMetadata, error) {
	pub := sv.Public()
	pem, err := test.PublicKeyToPEM(pub)
	if err != nil {
		return nil, fmt.Errorf("failed to convert public key to PEM: %w", err)
	}
	id, err := sv.KeyID()
	if err != nil {
		return nil, err
	}

	return &attestation.KeyMetadata{
		ID:            id,
		Status:        "active",
		SigningFormat: "dssev1",
		From:          time.Now(),
		PEM:           pem,
	}, nil
}