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

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/docker/attest/internal/util"
	"github.com/docker/attest/signerverifier"
	"github.com/docker/attest/tuf"
	"github.com/docker/attest/useragent"
	"github.com/go-openapi/runtime"
	"github.com/go-openapi/strfmt"
	"github.com/secure-systems-lab/go-securesystemslib/dsse"
	"github.com/sigstore/cosign/v2/pkg/cosign"
	rclient "github.com/sigstore/rekor/pkg/client"
	"github.com/sigstore/rekor/pkg/generated/models"
	"github.com/sigstore/rekor/pkg/types"
	hashedrekord_v001 "github.com/sigstore/rekor/pkg/types/hashedrekord/v0.0.1"

	stuf "github.com/sigstore/sigstore/pkg/tuf"

	_ "embed"
)

const RekorTLExtKind = "Rekor"

// ensure it has all the necessary methods.
var _ TransparencyLog = (*Rekor)(nil)

const defaultPublicKeysDir = "rekor"

type Rekor struct {
	publicKeys    *cosign.TrustedTransparencyLogPubKeys
	tufDownloader tuf.Downloader
	publicKeysDir string
}

//go:embed keys/c0d23d6ad406973f9559f3ba2d1ca01f84147d8ffc5b8445c224f98b9591801d.pem
var rekorPublicKey []byte

func WithTUFDownloader(tufDownloader tuf.Downloader) func(*Rekor) {
	return func(r *Rekor) {
		r.tufDownloader = tufDownloader
	}
}

func WithTUFPublicKeysDir(dir string) func(*Rekor) {
	return func(r *Rekor) {
		r.publicKeysDir = dir
	}
}

func NewRekorLog(options ...func(*Rekor)) (*Rekor, error) {
	pk, err := signerverifier.ParsePublicKey(rekorPublicKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing rekor public key: %w", err)
	}
	kid, err := signerverifier.KeyID(pk)
	if err != nil {
		return nil, fmt.Errorf("error getting keyid: %w", err)
	}
	keys := map[string]cosign.TransparencyLogPubKey{
		kid: {
			PubKey: pk,
			Status: stuf.Active,
		},
	}
	rekor := &Rekor{
		publicKeys: &cosign.TrustedTransparencyLogPubKeys{
			Keys: keys,
		},
		publicKeysDir: defaultPublicKeysDir,
	}
	for _, opt := range options {
		opt(rekor)
	}
	return rekor, nil
}

// UploadEntry submits a PK token signature to the transparency log.
func (tl *Rekor) UploadEntry(ctx context.Context, subject string, encPayload, signature []byte, signer dsse.SignerVerifier) (*DockerTLExtension, error) {
	// generate self-signed x509 cert
	pubCert, err := CreateX509Cert(subject, signer)
	if err != nil {
		return nil, fmt.Errorf("Error creating x509 cert: %w", err)
	}

	// generate hash of payload
	hasher := sha256.New()
	hasher.Write(encPayload)

	// upload entry
	rekorClient, err := rclient.GetRekorClient(DefaultRekorURL, rclient.WithUserAgent(useragent.Get(ctx)))
	if err != nil {
		return nil, fmt.Errorf("Error creating rekor client: %w", err)
	}
	entry, err := cosign.TLogUpload(ctx, rekorClient, signature, hasher, pubCert)
	if err != nil {
		return nil, fmt.Errorf("Error uploading tlog: %w", err)
	}

	return &DockerTLExtension{
		Kind: RekorTLExtKind,
		Data: entry, // transparency log entry metadata
	}, nil
}

// VerifyEntry verifies a transparency log entry.
func (tl *Rekor) VerifyEntry(ctx context.Context, ext *DockerTLExtension, encPayload, publicKey []byte) (time.Time, error) {
	zeroTime := time.Time{}
	// because the Data field has been unmarsalled into a map[string]interface{} we need to marshal it back to bytes
	// for the unmarshaler to work correctly
	entryBytes, err := json.Marshal(ext.Data)
	if err != nil {
		return time.Time{}, fmt.Errorf("error failed to marshal TL entry: %w", err)
	}

	entry, err := tl.UnmarshalEntry(entryBytes)
	if err != nil {
		return zeroTime, fmt.Errorf("error unmarshaling TL entry: %w", err)
	}

	err = entry.Validate(strfmt.Default)
	if err != nil {
		return zeroTime, fmt.Errorf("TL entry failed validation: %w", err)
	}
	// check if tl.publicKeys containers le.LogId
	_, ok := tl.publicKeys.Keys[*entry.LogID]
	if !ok {
		// otherwise check TUF
		pkTarget, err := tl.tufDownloader.DownloadTarget(filepath.Join(tl.publicKeysDir, fmt.Sprintf("%s.pem", *entry.LogID)), "")
		if err != nil {
			return zeroTime, fmt.Errorf("error downloading rekor public key %s: %w", *entry.LogID, err)
		}
		pk, err := signerverifier.ParsePublicKey(pkTarget.Data)
		if err != nil {
			return zeroTime, fmt.Errorf("error parsing public key: %w", err)
		}
		tl.publicKeys.Keys[*entry.LogID] = cosign.TransparencyLogPubKey{
			PubKey: pk,
			Status: stuf.Active,
		}
	}
	err = cosign.VerifyTLogEntryOffline(ctx, entry, tl.publicKeys)
	if err != nil {
		return zeroTime, fmt.Errorf("TL entry failed verification: %w", err)
	}

	integratedTime := time.Unix(*entry.IntegratedTime, 0)

	err = tl.VerifyEntryPayload(entry, encPayload, publicKey)
	if err != nil {
		return zeroTime, fmt.Errorf("error verifying TL entry payload: %w", err)
	}
	return integratedTime, nil
}

// VerifyEntryPayload checks that the TL entry payload matches envelope payload.
func (tl *Rekor) VerifyEntryPayload(entry *models.LogEntryAnon, payload, publicKey []byte) error {
	tlBody, ok := entry.Body.(string)
	if !ok {
		return fmt.Errorf("expected tl body to be of type string, got %T", entry)
	}
	rekord, err := extractHashedRekord(tlBody)
	if err != nil {
		return fmt.Errorf("error extract HashedRekord from TL entry: %w", err)
	}

	// compare payload hashes
	payloadHash := util.SHA256Hex(payload)
	if rekord.Hash != payloadHash {
		return fmt.Errorf("error payload and tl entry hash mismatch")
	}

	// compare public keys
	cert, err := base64.StdEncoding.Strict().DecodeString(rekord.PublicKey)
	if err != nil {
		return fmt.Errorf("failed to decode public key: %w", err)
	}
	p, _ := pem.Decode(cert)
	result, err := x509.ParseCertificate(p.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse certificate: %w", err)
	}
	if !bytes.Equal(result.RawSubjectPublicKeyInfo, publicKey) {
		return fmt.Errorf("error payload and tl entry public key mismatch")
	}
	return nil
}

func (tl *Rekor) UnmarshalEntry(entry []byte) (*models.LogEntryAnon, error) {
	le := new(models.LogEntryAnon)
	err := le.UnmarshalBinary(entry)
	if err != nil {
		return nil, fmt.Errorf("error failed to unmarshal Rekor entry: %w", err)
	}
	return le, nil
}

func extractHashedRekord(body string) (*Payload, error) {
	sig := new(Payload)
	pe, err := models.UnmarshalProposedEntry(base64.NewDecoder(base64.StdEncoding, strings.NewReader(body)), runtime.JSONConsumer())
	if err != nil {
		return nil, err
	}
	impl, err := types.UnmarshalEntry(pe)
	if err != nil {
		return nil, err
	}
	switch entry := impl.(type) {
	case *hashedrekord_v001.V001Entry:
		sig.Algorithm = *entry.HashedRekordObj.Data.Hash.Algorithm
		sig.Hash = *entry.HashedRekordObj.Data.Hash.Value
		sig.Signature = entry.HashedRekordObj.Signature.Content.String()
		sig.PublicKey = entry.HashedRekordObj.Signature.PublicKey.Content.String()
		return sig, nil
	default:
		return nil, fmt.Errorf("failed to extract haskedrekord, unsupported type: %T", entry)
	}
}