//
// Copyright 2022 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 cose

import (
	"bytes"
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/go-openapi/strfmt"
	"github.com/go-openapi/swag/conv"
	"github.com/in-toto/in-toto-golang/in_toto"
	gocose "github.com/veraison/go-cose"

	"github.com/sigstore/rekor/pkg/generated/models"
	"github.com/sigstore/rekor/pkg/log"
	"github.com/sigstore/rekor/pkg/pki"
	"github.com/sigstore/rekor/pkg/pki/x509"
	"github.com/sigstore/rekor/pkg/types"
	"github.com/sigstore/rekor/pkg/types/cose"
)

const (
	APIVERSION = "0.0.1"
)

const (
	CurveP256 = "P-256"
)

func init() {
	if err := cose.VersionMap.SetEntryFactory(APIVERSION, NewEntry); err != nil {
		log.Logger.Panic(err)
	}
}

var maxAttestationSize = 100 * 1024

func SetMaxAttestationSize(limit int) {
	maxAttestationSize = limit
}

type V001Entry struct {
	CoseObj      models.CoseV001Schema
	keyObj       pki.PublicKey
	sign1Msg     *gocose.Sign1Message
	envelopeHash []byte
}

func (v V001Entry) APIVersion() string {
	return APIVERSION
}

func NewEntry() types.EntryImpl {
	return &V001Entry{}
}

func (v V001Entry) IndexKeys() ([]string, error) {
	var result []string

	// We add the key, the hash of the overall cose envelope, and the hash of the payload itself as keys.
	if v.CoseObj.PublicKey == nil {
		return nil, errors.New("missing public key")
	}
	keyObj, err := x509.NewPublicKey(bytes.NewReader(*v.CoseObj.PublicKey))
	if err != nil {
		return nil, err
	}

	// 1. Key
	key, err := keyObj.CanonicalValue()
	if err != nil {
		log.Logger.Error(err)
	} else {
		keyHash := sha256.Sum256(key)
		result = append(result, strings.ToLower(hex.EncodeToString(keyHash[:])))
	}
	result = append(result, keyObj.Subjects()...)

	// 2. Overall envelope
	result = append(result, formatKey(v.CoseObj.Message))

	// 3. Payload
	if v.sign1Msg != nil {
		result = append(result, formatKey(v.sign1Msg.Payload))
	} else {
		// If no payload exists (it's unpacked in validate() method)
		// return now, as we will not be able to extract any headers
		return result, nil
	}

	// If payload is an in-toto statement, let's grab the subjects.
	if rawContentType, ok := v.sign1Msg.Headers.Protected[gocose.HeaderLabelContentType]; ok {
		contentType, ok := rawContentType.(string)
		// Integers as defined by CoAP content format are valid too,
		// but in-intoto payload type is not defined there, so only
		// proceed if content type is a string.
		// See list of CoAP content formats here:
		// https://www.iana.org/assignments/core-parameters/core-parameters.xhtml#content-formats
		if ok && contentType == in_toto.PayloadType {
			stmt, err := getIntotoStatement(v.sign1Msg.Payload)
			if err != nil {
				// ContentType header says intoto statement, but
				// parsing failed, continue with a warning.
				log.Logger.Warnf("Failed to parse intoto statement")
			} else {
				for _, sub := range stmt.Subject {
					for alg, digest := range sub.Digest {
						index := alg + ":" + digest
						result = append(result, index)
					}
				}
			}
		}
	}

	return result, nil
}

// indexKeyExtract captures only the fields of an in-toto statement that are
// used to derive index keys; it intentionally uses encoding/json semantics so
// that payloads accepted today continue to parse identically.
type indexKeyExtract struct {
	Subject []struct {
		Digest map[string]string `json:"digest"`
	} `json:"subject"`
}

func getIntotoStatement(b []byte) (*indexKeyExtract, error) {
	var stmt indexKeyExtract
	if err := json.Unmarshal(b, &stmt); err != nil {
		return nil, err
	}

	return &stmt, nil
}

func formatKey(b []byte) string {
	h := sha256.Sum256(b)
	hash := hex.EncodeToString(h[:])
	return strings.ToLower(fmt.Sprintf("%s:%s", models.CoseV001SchemaDataPayloadHashAlgorithmSha256, hash))
}

func (v *V001Entry) Unmarshal(pe models.ProposedEntry) error {
	it, ok := pe.(*models.Cose)
	if !ok {
		return errors.New("cannot unmarshal non Cose v0.0.1 type")
	}
	if it == nil {
		return errors.New("proposed entry cannot be nil")
	}

	var err error
	if err := DecodeEntry(it.Spec, &v.CoseObj); err != nil {
		return err
	}

	// field validation
	if err := v.CoseObj.Validate(strfmt.Default); err != nil {
		return err
	}

	if v.CoseObj.PublicKey == nil {
		return errors.New("missing public key")
	}
	v.keyObj, err = x509.NewPublicKey(bytes.NewReader(*v.CoseObj.PublicKey))
	if err != nil {
		return err
	}

	// Store the envelope hash.
	// The CoseObj.Message is only populated during entry creation.
	// When marshalling from the database (retrieval) the envelope
	// hash must be decoded from the stored hex string.
	// The envelope hash is used to create the attestation key during
	// retrieval of a record.
	if len(v.CoseObj.Message) == 0 {
		if v.CoseObj.Data == nil || v.CoseObj.Data.EnvelopeHash == nil || v.CoseObj.Data.EnvelopeHash.Value == nil {
			return errors.New("envelope hash should have been previously computed")
		}
		b, err := hex.DecodeString(*v.CoseObj.Data.EnvelopeHash.Value)
		if err != nil {
			return err
		}
		v.envelopeHash = b
	} else {
		h := sha256.Sum256(v.CoseObj.Message)
		v.envelopeHash = h[:]
	}

	return v.validate()
}

func (v *V001Entry) Canonicalize(_ context.Context) ([]byte, error) {
	if v.keyObj == nil {
		return nil, errors.New("cannot canonicalize empty key")
	}
	if v.sign1Msg == nil {
		return nil, errors.New("signed message uninitialized")
	}
	if v.sign1Msg.Payload == nil {
		return nil, errors.New("payload empty")
	}

	pk, err := v.keyObj.CanonicalValue()
	if err != nil {
		return nil, err
	}
	pkb := strfmt.Base64(pk)

	h := sha256.Sum256([]byte(v.sign1Msg.Payload))

	canonicalEntry := models.CoseV001Schema{
		PublicKey: &pkb,
		Data: &models.CoseV001SchemaData{
			PayloadHash: &models.CoseV001SchemaDataPayloadHash{
				Algorithm: conv.Pointer(models.CoseV001SchemaDataPayloadHashAlgorithmSha256),
				Value:     conv.Pointer(hex.EncodeToString(h[:])),
			},
			EnvelopeHash: &models.CoseV001SchemaDataEnvelopeHash{
				Algorithm: conv.Pointer(models.CoseV001SchemaDataEnvelopeHashAlgorithmSha256),
				Value:     conv.Pointer(hex.EncodeToString(v.envelopeHash)),
			},
		},
	}

	itObj := models.Cose{}
	itObj.APIVersion = conv.Pointer(APIVERSION)
	itObj.Spec = &canonicalEntry

	return json.Marshal(&itObj)
}

// validate performs cross-field validation for fields in object
func (v *V001Entry) validate() error {
	// This also gets called in the CLI, where we won't have this data
	// or during record retrieval (message is the raw COSE object) which
	// is only stored as an attestation.
	if len(v.CoseObj.Message) == 0 {
		return nil
	}

	alg, pk, err := getPublicKey(v.keyObj)
	if err != nil {
		return err
	}

	bv, err := gocose.NewVerifier(alg, pk)
	if err != nil {
		return err
	}
	sign1Msg := gocose.NewSign1Message()
	if err := sign1Msg.UnmarshalCBOR(v.CoseObj.Message); err != nil {
		return err
	}

	if v.CoseObj.Data == nil {
		return errors.New("cose object data cannot be nil")
	}
	if err := sign1Msg.Verify(v.CoseObj.Data.Aad, bv); err != nil {
		return err
	}

	v.sign1Msg = sign1Msg
	return nil
}

func getPublicKey(pk pki.PublicKey) (gocose.Algorithm, crypto.PublicKey, error) {
	invAlg := gocose.Algorithm(0)
	x5pk, ok := pk.(*x509.PublicKey)

	if !ok || x5pk == nil {
		return invAlg, nil, errors.New("invalid public key type")
	}

	cryptoPub := x5pk.CryptoPubKey()

	var alg gocose.Algorithm
	switch t := cryptoPub.(type) {
	case *rsa.PublicKey:
		alg = gocose.AlgorithmPS256
	case *ecdsa.PublicKey:
		alg = gocose.AlgorithmES256
		if t.Params().Name != CurveP256 {
			return invAlg, nil, fmt.Errorf("unsupported elliptic curve %s",
				t.Params().Name)
		}
	default:
		return invAlg, nil, fmt.Errorf("unsupported algorithm type %T", t)
	}

	return alg, cryptoPub, nil
}

// AttestationKey returns the digest of the COSE envelope that was uploaded,
// to be used to lookup the attestation from storage.
func (v *V001Entry) AttestationKey() string {
	return fmt.Sprintf("%s:%s",
		models.CoseV001SchemaDataEnvelopeHashAlgorithmSha256,
		hex.EncodeToString(v.envelopeHash))
}

// DecodeEntry performs direct decode into the provided output pointer
// without mutating the receiver on error.
func DecodeEntry(input any, output *models.CoseV001Schema) error {
	if output == nil {
		return fmt.Errorf("nil output *models.CoseV001Schema")
	}
	var m models.CoseV001Schema
	// Typed switch including map fast path
	switch data := input.(type) {
	case map[string]any:
		mm := data
		if msg, ok := mm["message"].(string); ok && msg != "" {
			outb := make([]byte, base64.StdEncoding.DecodedLen(len(msg)))
			n, err := base64.StdEncoding.Decode(outb, []byte(msg))
			if err != nil {
				return fmt.Errorf("failed parsing base64 data for message: %w", err)
			}
			m.Message = strfmt.Base64(outb[:n])
		}
		if pk, ok := mm["publicKey"].(string); ok && pk != "" {
			outb := make([]byte, base64.StdEncoding.DecodedLen(len(pk)))
			n, err := base64.StdEncoding.Decode(outb, []byte(pk))
			if err != nil {
				return fmt.Errorf("failed parsing base64 data for publicKey: %w", err)
			}
			b := strfmt.Base64(outb[:n])
			m.PublicKey = &b
		}
		if d, ok := mm["data"].(map[string]any); ok {
			m.Data = &models.CoseV001SchemaData{}
			if ph, ok := d["payloadHash"].(map[string]any); ok {
				m.Data.PayloadHash = &models.CoseV001SchemaDataPayloadHash{}
				if alg, ok := ph["algorithm"].(string); ok {
					m.Data.PayloadHash.Algorithm = &alg
				}
				if val, ok := ph["value"].(string); ok {
					m.Data.PayloadHash.Value = &val
				}
			}
			if eh, ok := d["envelopeHash"].(map[string]any); ok {
				m.Data.EnvelopeHash = &models.CoseV001SchemaDataEnvelopeHash{}
				if alg, ok := eh["algorithm"].(string); ok {
					m.Data.EnvelopeHash.Algorithm = &alg
				}
				if val, ok := eh["value"].(string); ok {
					m.Data.EnvelopeHash.Value = &val
				}
			}
			if aad, ok := d["aad"].(string); ok && aad != "" {
				outb := make([]byte, base64.StdEncoding.DecodedLen(len(aad)))
				n, err := base64.StdEncoding.Decode(outb, []byte(aad))
				if err != nil {
					return fmt.Errorf("failed parsing base64 data for aad: %w", err)
				}
				m.Data.Aad = strfmt.Base64(outb[:n])
			}
		}
		*output = m
		return nil
	case *models.CoseV001Schema:
		if data == nil {
			return fmt.Errorf("nil *models.CoseV001Schema")
		}
		*output = *data
		return nil
	case models.CoseV001Schema:
		*output = data
		return nil
	default:
		return fmt.Errorf("unsupported input type %T for DecodeEntry", input)
	}
}

// AttestationKeyValue returns both the key and value to be persisted
// into attestation storage
func (v *V001Entry) AttestationKeyValue() (string, []byte) {
	storageSize := len(v.CoseObj.Message)
	if storageSize > maxAttestationSize {
		log.Logger.Infof("Skipping attestation storage, size %d is greater than max %d", storageSize, maxAttestationSize)
		return "", nil
	}

	return v.AttestationKey(), v.CoseObj.Message
}

func (v V001Entry) CreateFromArtifactProperties(_ context.Context, props types.ArtifactProperties) (models.ProposedEntry, error) {
	returnVal := models.Cose{}
	var err error
	messageBytes := props.ArtifactBytes
	if messageBytes == nil {
		if props.ArtifactPath == nil {
			return nil, errors.New("path to artifact file must be specified")
		}
		if props.ArtifactPath.IsAbs() {
			return nil, errors.New("cose envelopes cannot be fetched over HTTP(S)")
		}
		messageBytes, err = os.ReadFile(filepath.Clean(props.ArtifactPath.Path))
		if err != nil {
			return nil, err
		}
	}
	publicKeyBytes := props.PublicKeyBytes
	if len(publicKeyBytes) == 0 {
		if len(props.PublicKeyPaths) != 1 {
			return nil, errors.New("only one public key must be provided to verify signature")
		}
		if props.PublicKeyPaths[0] == nil {
			return nil, errors.New("public key path must be provided")
		}
		keyBytes, err := os.ReadFile(filepath.Clean(props.PublicKeyPaths[0].Path))
		if err != nil {
			return nil, fmt.Errorf("error reading public key file: %w", err)
		}
		publicKeyBytes = append(publicKeyBytes, keyBytes)
	} else if len(publicKeyBytes) != 1 {
		return nil, errors.New("only one public key must be provided")
	}

	kb := strfmt.Base64(publicKeyBytes[0])
	mb := strfmt.Base64(messageBytes)

	re := V001Entry{
		CoseObj: models.CoseV001Schema{
			Data: &models.CoseV001SchemaData{
				Aad: props.AdditionalAuthenticatedData,
			},
			PublicKey: &kb,
			Message:   mb,
		},
	}

	returnVal.Spec = re.CoseObj
	returnVal.APIVersion = conv.Pointer(re.APIVersion())

	return &returnVal, nil
}

func (v V001Entry) Verifiers() ([]pki.PublicKey, error) {
	if v.CoseObj.PublicKey == nil {
		return nil, errors.New("cose v0.0.1 entry not initialized")
	}
	key, err := x509.NewPublicKey(bytes.NewReader(*v.CoseObj.PublicKey))
	if err != nil {
		return nil, err
	}
	return []pki.PublicKey{key}, nil
}

func (v V001Entry) ArtifactHash() (string, error) {
	if v.CoseObj.Data == nil || v.CoseObj.Data.PayloadHash == nil || v.CoseObj.Data.PayloadHash.Value == nil || v.CoseObj.Data.PayloadHash.Algorithm == nil {
		return "", errors.New("cose v0.0.1 entry not initialized")
	}
	return strings.ToLower(fmt.Sprintf("%s:%s", *v.CoseObj.Data.PayloadHash.Algorithm, *v.CoseObj.Data.PayloadHash.Value)), nil
}

func (v V001Entry) Insertable() (bool, error) {
	if len(v.CoseObj.Message) == 0 {
		return false, errors.New("missing COSE Sign1 message")
	}
	if v.CoseObj.PublicKey == nil || len(*v.CoseObj.PublicKey) == 0 {
		return false, errors.New("missing public key")
	}
	if v.CoseObj.Data == nil {
		return false, errors.New("missing COSE data property")
	}
	if len(v.envelopeHash) == 0 {
		return false, errors.New("envelope hash has not been computed")
	}
	if v.keyObj == nil {
		return false, errors.New("public key has not been parsed")
	}
	if v.sign1Msg == nil {
		return false, errors.New("signature has not been validated")
	}

	return true, nil
}
