//
// Copyright 2024 The Sigstore 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 bundle

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"os"
	"path/filepath"
	"testing"

	"github.com/sigstore/cosign/v3/internal/test"
	"github.com/sigstore/cosign/v3/pkg/cosign"
	"github.com/sigstore/cosign/v3/pkg/cosign/bundle"
	sgBundle "github.com/sigstore/sigstore-go/pkg/bundle"
	"github.com/sigstore/sigstore/pkg/cryptoutils"
)

func TestCreateCmd(t *testing.T) {
	ctx := context.Background()

	artifact := "hello world"
	digest := sha256.Sum256([]byte(artifact))

	td := t.TempDir()
	artifactPath := filepath.Join(td, "artifact")
	err := os.WriteFile(artifactPath, []byte(artifact), 0600)
	checkErr(t, err)

	// Test signing with a key
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	checkErr(t, err)
	sigBytes, err := privateKey.Sign(rand.Reader, digest[:], crypto.SHA256)
	checkErr(t, err)

	signature := base64.StdEncoding.EncodeToString(sigBytes)
	sigPath := filepath.Join(td, "sig")
	err = os.WriteFile(sigPath, []byte(signature), 0600)
	checkErr(t, err)

	publicKeyPath := filepath.Join(td, "key.pub")
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	checkErr(t, err)
	pemBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	}
	err = os.WriteFile(publicKeyPath, pem.EncodeToMemory(pemBlock), 0600)
	checkErr(t, err)

	outPath := filepath.Join(td, "bundle.sigstore.json")

	bundleCreate := CreateCmd{
		Artifact:      artifactPath,
		KeyRef:        publicKeyPath,
		IgnoreTlog:    true,
		Out:           outPath,
		SignaturePath: sigPath,
	}

	err = bundleCreate.Exec(ctx)
	checkErr(t, err)

	b, err := sgBundle.LoadJSONFromPath(outPath)
	checkErr(t, err)

	if b.VerificationMaterial == nil {
		t.Fatal("bundle does not have verification material")
	}

	if b.VerificationMaterial.GetPublicKey() == nil {
		t.Fatal("bundle verification material does not have public key")
	}

	if b.GetMessageSignature() == nil {
		t.Fatal("bundle does not have message signature")
	}

	// Test using an identity certificate in an old bundle format
	rootCert, rootKey, _ := test.GenerateRootCa()
	leafCert, privKey, _ := test.GenerateLeafCert("subject", "oidc-issuer", rootCert, rootKey)

	sigBytes, err = privKey.Sign(rand.Reader, digest[:], crypto.SHA256)
	checkErr(t, err)

	signedPayload := cosign.LocalSignedPayload{}
	signedPayload.Base64Signature = base64.StdEncoding.EncodeToString(sigBytes)

	certBytes, err := cryptoutils.MarshalCertificateToPEM(leafCert)
	checkErr(t, err)

	signedPayload.Cert = base64.StdEncoding.EncodeToString(certBytes)
	bundleContents, err := json.Marshal(signedPayload)
	checkErr(t, err)

	bundlePath := filepath.Join(td, "old-bundle.json")
	err = os.WriteFile(bundlePath, bundleContents, 0600)
	checkErr(t, err)

	bundleCreate = CreateCmd{
		Artifact:   artifactPath,
		BundlePath: bundlePath,
		IgnoreTlog: true,
		Out:        outPath,
	}

	err = bundleCreate.Exec(ctx)
	checkErr(t, err)

	b, err = sgBundle.LoadJSONFromPath(outPath)
	checkErr(t, err)

	if b.VerificationMaterial == nil {
		t.Fatal("bundle does not have verification material")
	}

	if b.VerificationMaterial.GetCertificate() == nil {
		t.Fatal("bundle verification material does not have certificate")
	}

	if b.GetMessageSignature() == nil {
		t.Fatal("bundle does not have message signature")
	}
}

func checkErr(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func TestCreateCmd_FailOnIgnoreTlogWithSET(t *testing.T) {
	ctx := context.Background()

	artifact := "hello world"
	digest := sha256.Sum256([]byte(artifact))

	td := t.TempDir()
	artifactPath := filepath.Join(td, "artifact")
	err := os.WriteFile(artifactPath, []byte(artifact), 0600)
	checkErr(t, err)

	rootCert, rootKey, _ := test.GenerateRootCa()
	leafCert, privKey, _ := test.GenerateLeafCert("subject", "oidc-issuer", rootCert, rootKey)

	sigBytes, err := privKey.Sign(rand.Reader, digest[:], crypto.SHA256)
	checkErr(t, err)

	// Test using an old bundle with a SET
	signedPayloadWithSET := cosign.LocalSignedPayload{}
	signedPayloadWithSET.Base64Signature = base64.StdEncoding.EncodeToString(sigBytes)

	certBytes, err := cryptoutils.MarshalCertificateToPEM(leafCert)
	checkErr(t, err)

	signedPayloadWithSET.Cert = base64.StdEncoding.EncodeToString(certBytes)
	signedPayloadWithSET.Bundle = &bundle.RekorBundle{
		SignedEntryTimestamp: []byte("set"),
	}

	bundleContentsWithSET, err := json.Marshal(signedPayloadWithSET)
	checkErr(t, err)
	bundlePathWithSET := filepath.Join(td, "old-bundle-with-set.json")
	err = os.WriteFile(bundlePathWithSET, bundleContentsWithSET, 0600)
	checkErr(t, err)

	// Test using an old bundle without a SET
	signedPayloadWithoutSET := cosign.LocalSignedPayload{}
	signedPayloadWithoutSET.Base64Signature = signedPayloadWithSET.Base64Signature
	signedPayloadWithoutSET.Cert = signedPayloadWithSET.Cert

	bundleContentsWithoutSET, err := json.Marshal(signedPayloadWithoutSET)
	checkErr(t, err)
	bundlePathWithoutSET := filepath.Join(td, "old-bundle-without-set.json")
	err = os.WriteFile(bundlePathWithoutSET, bundleContentsWithoutSET, 0600)
	checkErr(t, err)

	tests := []struct {
		name        string
		bundlePath  string
		ignoreTlog  bool
		expectError bool
		errStr      string
	}{
		{
			name:        "Fail when bundle has SET and IgnoreTlog is true",
			bundlePath:  bundlePathWithSET,
			ignoreTlog:  true,
			expectError: true,
			errStr:      "cannot ignore transparency log when the provided bundle contains a Signed Entry Timestamp",
		},
		{
			name:        "Pass when bundle has no SET and IgnoreTlog is true",
			bundlePath:  bundlePathWithoutSET,
			ignoreTlog:  true,
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			outPath := filepath.Join(td, "out-"+filepath.Base(tt.bundlePath))
			bundleCreate := CreateCmd{
				Artifact:   artifactPath,
				BundlePath: tt.bundlePath,
				IgnoreTlog: tt.ignoreTlog,
				Out:        outPath,
			}

			err := bundleCreate.Exec(ctx)
			if tt.expectError {
				if err == nil {
					t.Fatal("expected error, got nil")
				}
				if tt.errStr != "" && err.Error() != tt.errStr {
					t.Fatalf("expected error %q, got %q", tt.errStr, err.Error())
				}
			} else {
				checkErr(t, err)
			}
		})
	}
}