// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Package tpm2 provides TPM2.0 related functionality helpers.
package tpm2

import (
	"bytes"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"

	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpm2/transport"

	"github.com/siderolabs/talos/internal/pkg/tpm"
	"github.com/siderolabs/talos/pkg/machinery/constants"
)

// Unseal unseals a sealed blob using the TPM
//
//nolint:gocyclo,cyclop
func Unseal(sealed SealedResponse) ([]byte, error) {
	t, err := tpm.Open()
	if err != nil {
		return nil, err
	}
	defer t.Close() //nolint:errcheck

	// fail early if PCR banks are not present or filled with all zeroes or 0xff
	if err = validatePCRBanks(t, sealed.PCRs); err != nil {
		return nil, err
	}

	tpmPub, err := tpm2.Unmarshal[tpm2.TPM2BPublic](sealed.SealedBlobPublic)
	if err != nil {
		return nil, err
	}

	tpmPriv, err := tpm2.Unmarshal[tpm2.TPM2BPrivate](sealed.SealedBlobPrivate)
	if err != nil {
		return nil, err
	}

	srk, err := tpm2.Unmarshal[tpm2.TPM2BName](sealed.KeyName)
	if err != nil {
		return nil, err
	}

	// we need to create a primary since we don't persist the SRK
	primary := tpm2.CreatePrimary{
		PrimaryHandle: tpm2.TPMRHOwner,
		InPublic:      tpm2.New2B(tpm2.ECCSRKTemplate),
	}

	createPrimaryResponse, err := primary.Execute(t)
	if err != nil {
		return nil, err
	}

	defer func() {
		flush := tpm2.FlushContext{
			FlushHandle: createPrimaryResponse.ObjectHandle,
		}

		_, flushErr := flush.Execute(t)
		if flushErr != nil {
			err = flushErr
		}
	}()

	outPub, err := createPrimaryResponse.OutPublic.Contents()
	if err != nil {
		return nil, err
	}

	if !bytes.Equal(createPrimaryResponse.Name.Buffer, srk.Buffer) {
		// this means the srk name does not match, possibly due to a different TPM or tpm was reset
		// could also mean the disk was used on a different machine
		return nil, fmt.Errorf("srk name does not match, expected %x, got %x", srk.Buffer, createPrimaryResponse.Name.Buffer)
	}

	load := tpm2.Load{
		ParentHandle: tpm2.NamedHandle{
			Handle: createPrimaryResponse.ObjectHandle,
			Name:   createPrimaryResponse.Name,
		},
		InPrivate: *tpmPriv,
		InPublic:  *tpmPub,
	}

	loadResponse, err := load.Execute(t)
	if err != nil {
		return nil, err
	}

	policySess, policyCloseFunc, err := tpm2.PolicySession(
		t,
		tpm2.TPMAlgSHA256,
		20,
		tpm2.Salted(createPrimaryResponse.ObjectHandle, *outPub),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create policy session: %w", err)
	}

	defer policyCloseFunc() //nolint:errcheck

	pubKey, err := ParsePCRSigningPubKey(constants.PCRPublicKey)
	if err != nil {
		return nil, err
	}

	loadExternal := tpm2.LoadExternal{
		Hierarchy: tpm2.TPMRHOwner,
		InPublic:  tpm2.New2B(RSAPubKeyTemplate(pubKey.N.BitLen(), pubKey.E, pubKey.N.Bytes())),
	}

	loadExternalResponse, err := loadExternal.Execute(t)
	if err != nil {
		return nil, fmt.Errorf("failed to load external key: %w", err)
	}

	defer func() {
		flush := tpm2.FlushContext{
			FlushHandle: loadExternalResponse.ObjectHandle,
		}

		_, flushErr := flush.Execute(t)
		if flushErr != nil {
			err = flushErr
		}
	}()

	pcrSelector, err := CreateSelector([]int{constants.UKIPCR})
	if err != nil {
		return nil, err
	}

	policyDigest, err := PolicyPCRDigest(t, policySess.Handle(), tpm2.TPMLPCRSelection{
		PCRSelections: []tpm2.TPMSPCRSelection{
			{
				Hash:      tpm2.TPMAlgSHA256,
				PCRSelect: pcrSelector,
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve policy digest: %w", err)
	}

	sigJSON, err := ParsePCRSignature()
	if err != nil {
		return nil, err
	}

	pubKeyFingerprint := sha256.Sum256(x509.MarshalPKCS1PublicKey(pubKey))

	var signature string
	// TODO: maybe we should use the highest supported algorithm of the TPM
	// fallback to the next one if the signature is not found
	for _, bank := range sigJSON.SHA256 {
		digest, decodeErr := hex.DecodeString(bank.Pol)
		if decodeErr != nil {
			return nil, decodeErr
		}

		if bytes.Equal(digest, policyDigest.Buffer) {
			signature = bank.Sig

			if hex.EncodeToString(pubKeyFingerprint[:]) != bank.PKFP {
				return nil, errors.New("certificate fingerprint does not match")
			}

			break
		}
	}

	if signature == "" {
		return nil, errors.New("no signatures matching PCR SHA256 digest found in signature JSON")
	}

	signatureDecoded, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return nil, err
	}

	// Verify will only verify the RSA part of the RSA+SHA256 signature,
	// hence we need to do the SHA256 part ourselves
	policyDigestHash := sha256.Sum256(policyDigest.Buffer)

	verifySignature := tpm2.VerifySignature{
		KeyHandle: loadExternalResponse.ObjectHandle,
		Digest: tpm2.TPM2BDigest{
			Buffer: policyDigestHash[:],
		},
		Signature: tpm2.TPMTSignature{
			SigAlg: tpm2.TPMAlgRSASSA,
			Signature: tpm2.NewTPMUSignature(tpm2.TPMAlgRSASSA, &tpm2.TPMSSignatureRSA{
				Hash: tpm2.TPMAlgSHA256,
				Sig: tpm2.TPM2BPublicKeyRSA{
					Buffer: signatureDecoded,
				},
			}),
		},
	}

	verifySignatureResponse, err := verifySignature.Execute(t)
	if err != nil {
		return nil, fmt.Errorf("failed to verify signature: %w", err)
	}

	policyAuthorize := tpm2.PolicyAuthorize{
		PolicySession:  policySess.Handle(),
		ApprovedPolicy: *policyDigest,
		KeySign:        loadExternalResponse.Name,
		CheckTicket:    verifySignatureResponse.Validation,
	}

	if _, err = policyAuthorize.Execute(t); err != nil {
		return nil, fmt.Errorf("failed to execute policy authorize: %w", err)
	}

	// this handles the case when talos is upgraded from pre Talos 1.12 and we had PCRs field
	// set to just PCR 11 but locked to PCR 7
	// in this case the EncryptionVersion is empty and the PubKeyPCRs field is also empty
	// since both EncryptionVersion and PubKeyPCRs are introduced for Talos 1.12+ only
	if sealed.EncryptionVersion == "" && len(sealed.PubKeyPCRs) == 0 {
		if len(sealed.PCRs) == 1 && sealed.PCRs[0] == constants.UKIPCR {
			if err := validatePCRPolicyDigest(t, policySess.Handle(), []int{constants.SecureBootStatePCR}, sealed.PolicyDigest); err != nil {
				return nil, fmt.Errorf("failed to validate PCR policy digest for PCRs %v: %w", []int{constants.SecureBootStatePCR}, err)
			}
		}
	}

	// Talos 1.12+ sets the EncryptionVersion and PubKeyPCRs to PCR 11 and any other PCR used to PCRs field
	if sealed.EncryptionVersion != "" && len(sealed.PCRs) > 0 {
		if err := validatePCRPolicyDigest(t, policySess.Handle(), sealed.PCRs, sealed.PolicyDigest); err != nil {
			return nil, fmt.Errorf("failed to validate PCR policy digest for PCRs %v: %w", sealed.PCRs, err)
		}
	}

	unsealOp := tpm2.Unseal{
		ItemHandle: tpm2.AuthHandle{
			Handle: loadResponse.ObjectHandle,
			Name:   loadResponse.Name,
			Auth:   policySess,
		},
	}

	unsealResponse, err := unsealOp.Execute(t, tpm2.HMAC(
		tpm2.TPMAlgSHA256,
		20,
		tpm2.Salted(createPrimaryResponse.ObjectHandle, *outPub),
		tpm2.AESEncryption(128, tpm2.EncryptOut),
		tpm2.Bound(loadResponse.ObjectHandle, loadResponse.Name, nil),
	))
	if err != nil {
		return nil, fmt.Errorf("failed to unseal op: %w", err)
	}

	return unsealResponse.OutData.Buffer, nil
}

func validatePCRPolicyDigest(t transport.TPM, handle tpm2.TPMHandle, pcrs []int, digest []byte) error {
	pcrSelector, err := CreateSelector(pcrs)
	if err != nil {
		return fmt.Errorf("failed to create PCR selector for PCRs %v: %w", pcrs, err)
	}

	pcrPolicyDigest, err := PolicyPCRDigest(t, handle, tpm2.TPMLPCRSelection{
		PCRSelections: []tpm2.TPMSPCRSelection{
			{
				Hash:      tpm2.TPMAlgSHA256,
				PCRSelect: pcrSelector,
			},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to calculate policy PCR digest: %w", err)
	}

	if !bytes.Equal(pcrPolicyDigest.Buffer, digest) {
		return fmt.Errorf("sealing policy digest does not match, expected %x, got %x", digest, pcrPolicyDigest.Buffer)
	}

	return nil
}
