//
// 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/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"os"

	"github.com/secure-systems-lab/go-securesystemslib/dsse"
	"github.com/sigstore/rekor/pkg/generated/client"
	"github.com/sigstore/sigstore/pkg/cryptoutils"
	"github.com/sigstore/sigstore/pkg/signature"

	"github.com/sigstore/cosign/v2/cmd/cosign/cli/options"
	"github.com/sigstore/cosign/v2/cmd/cosign/cli/rekor"
	"github.com/sigstore/cosign/v2/cmd/cosign/cli/verify"
	"github.com/sigstore/cosign/v2/pkg/cosign"
	"github.com/sigstore/cosign/v2/pkg/cosign/bundle"
	"github.com/sigstore/cosign/v2/pkg/cosign/pivkey"
	"github.com/sigstore/cosign/v2/pkg/cosign/pkcs11key"
	sigs "github.com/sigstore/cosign/v2/pkg/signature"
)

type CreateCmd struct {
	Artifact             string
	AttestationPath      string
	BundlePath           string
	CertificatePath      string
	IgnoreTlog           bool
	KeyRef               string
	Out                  string
	RekorURL             string
	RFC3161TimestampPath string
	SignaturePath        string
	Sk                   bool
	Slot                 string
}

func (c *CreateCmd) Exec(ctx context.Context) (err error) {
	if c.Artifact == "" {
		return fmt.Errorf("must supply --artifact")
	}

	// We require some signature
	if options.NOf(c.BundlePath, c.SignaturePath) == 0 {
		return fmt.Errorf("must at least supply signature via --bundle or --signature")
	}

	var cert *x509.Certificate
	var envelope dsse.Envelope
	var rekorClient *client.Rekor
	var sigBytes, signedTimestamp []byte
	var sigVerifier signature.Verifier

	if c.BundlePath != "" {
		b, err := cosign.FetchLocalSignedPayloadFromPath(c.BundlePath)
		if err != nil {
			return err
		}

		if b.Cert != "" {
			certPEM, err := base64.StdEncoding.DecodeString(b.Cert)
			if err != nil {
				return err
			}
			certs, err := cryptoutils.UnmarshalCertificatesFromPEM(certPEM)
			if err != nil {
				return err
			}
			if len(certs) == 0 {
				return fmt.Errorf("no certs found in bundle")
			}
			cert = certs[0]
		}

		if b.Base64Signature != "" {
			// Could be a DSSE envelope or plain signature
			signature, err := base64.StdEncoding.DecodeString(b.Base64Signature)
			if err != nil {
				return err
			}

			// See if DSSE JSON unmashalling succeeds
			err = json.Unmarshal(signature, &envelope)
			if err != nil {
				// Guess that it is a plain signature
				sigBytes = signature
			}
		}
	}

	if c.SignaturePath != "" {
		signatureB64, err := os.ReadFile(c.SignaturePath)
		if err != nil {
			return err
		}

		sigBytes, err = base64.StdEncoding.DecodeString(string(signatureB64))
		if err != nil {
			return err
		}
	}

	if c.RFC3161TimestampPath != "" {
		timestampBytes, err := os.ReadFile(c.RFC3161TimestampPath)
		if err != nil {
			return err
		}

		var rfc3161Timestamp bundle.RFC3161Timestamp
		err = json.Unmarshal(timestampBytes, &rfc3161Timestamp)
		if err != nil {
			return err
		}

		signedTimestamp = rfc3161Timestamp.SignedRFC3161Timestamp
	}

	if c.CertificatePath != "" {
		certBytes, err := os.ReadFile(c.CertificatePath)
		if err != nil {
			return err
		}

		certDecoded, err := base64.StdEncoding.DecodeString(string(certBytes))
		if err != nil {
			return err
		}

		block, _ := pem.Decode(certDecoded)
		if block == nil {
			return fmt.Errorf("unable to decode provided certificate")
		}

		cert, err = x509.ParseCertificate(block.Bytes)
		if err != nil {
			return err
		}
	}

	if c.AttestationPath != "" {
		attestationBytes, err := os.ReadFile(c.AttestationPath)
		if err != nil {
			return err
		}

		err = json.Unmarshal(attestationBytes, &envelope)
		if err != nil {
			return err
		}
	}

	if c.KeyRef != "" {
		sigVerifier, err = sigs.PublicKeyFromKeyRef(ctx, c.KeyRef)
		if err != nil {
			return fmt.Errorf("loading public key: %w", err)
		}
		pkcs11Key, ok := sigVerifier.(*pkcs11key.Key)
		if ok {
			defer pkcs11Key.Close()
		}
	} else if c.Sk {
		sk, err := pivkey.GetKeyWithSlot(c.Slot)
		if err != nil {
			return fmt.Errorf("opening piv token: %w", err)
		}
		defer sk.Close()
		sigVerifier, err = sk.Verifier()
		if err != nil {
			return fmt.Errorf("loading public key from token: %w", err)
		}
	}

	if c.RekorURL != "" {
		rekorClient, err = rekor.NewClient(c.RekorURL)
		if err != nil {
			return err
		}
	}

	bundle, err := verify.AssembleNewBundle(ctx, sigBytes, signedTimestamp, &envelope, c.Artifact, cert, c.IgnoreTlog, sigVerifier, nil, rekorClient)
	if err != nil {
		return err
	}

	bundleBytes, err := bundle.MarshalJSON()
	if err != nil {
		return err
	}

	if c.Out != "" {
		err = os.WriteFile(c.Out, bundleBytes, 0600)
		if err != nil {
			return err
		}
	} else {
		fmt.Println(string(bundleBytes))
	}

	return nil
}
