// 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"
	"crypto/sha256"
	"fmt"
	"log"
	"os"
	"strings"

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

	"github.com/siderolabs/talos/internal/pkg/secureboot"
)

// CreateSelector converts PCR  numbers into a bitmask.
func CreateSelector(pcrs []int) ([]byte, error) {
	const sizeOfPCRSelect = 3

	mask := make([]byte, sizeOfPCRSelect)

	for _, n := range pcrs {
		if n >= 8*sizeOfPCRSelect {
			return nil, fmt.Errorf("PCR index %d is out of range (exceeds maximum value %d)", n, 8*sizeOfPCRSelect-1)
		}

		mask[n>>3] |= 1 << (n & 0x7)
	}

	return mask, nil
}

// ReadPCR reads the value of a single PCR.
func ReadPCR(t transport.TPM, pcr int) ([]byte, error) {
	pcrSelector, err := CreateSelector([]int{pcr})
	if err != nil {
		return nil, fmt.Errorf("failed to create PCR selection: %v", err)
	}

	pcrRead := tpm2.PCRRead{
		PCRSelectionIn: tpm2.TPMLPCRSelection{
			PCRSelections: []tpm2.TPMSPCRSelection{
				{
					Hash:      tpm2.TPMAlgSHA256,
					PCRSelect: pcrSelector,
				},
			},
		},
	}

	pcrValue, err := pcrRead.Execute(t)
	if err != nil {
		return nil, fmt.Errorf("failed to read PCR: %v", err)
	}

	return pcrValue.PCRValues.Digests[0].Buffer, nil
}

// PCRExtent hashes the input and extends the PCR with the hash.
func PCRExtent(pcr int, data []byte) error {
	t, err := transport.OpenTPM()
	if err != nil {
		// if the TPM is not available or not a TPM 2.0, we can skip the PCR extension
		if os.IsNotExist(err) || strings.Contains(err.Error(), "device is not a TPM 2.0") {
			log.Printf("TPM device is not available, skipping PCR extension")

			return nil
		}

		return err
	}

	defer t.Close() // nolint: errcheck

	// since we are using SHA256, we can assume that the PCR bank is SHA256
	digest := sha256.Sum256(data)

	pcrHandle := tpm2.PCRExtend{
		PCRHandle: tpm2.AuthHandle{
			Handle: tpm2.TPMHandle(pcr),
			Auth:   tpm2.PasswordAuth(nil),
		},
		Digests: tpm2.TPMLDigestValues{
			Digests: []tpm2.TPMTHA{
				{
					HashAlg: tpm2.TPMAlgSHA256,
					Digest:  digest[:],
				},
			},
		},
	}

	if _, err = pcrHandle.Execute(t); err != nil {
		return err
	}

	return nil
}

// PolicyPCRDigest executes policyPCR and returns the digest.
func PolicyPCRDigest(t transport.TPM, policyHandle tpm2.TPMHandle, pcrSelection tpm2.TPMLPCRSelection) (*tpm2.TPM2BDigest, error) {
	policyPCR := tpm2.PolicyPCR{
		PolicySession: policyHandle,
		Pcrs:          pcrSelection,
	}

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

	policyGetDigest := tpm2.PolicyGetDigest{
		PolicySession: policyHandle,
	}

	policyGetDigestResponse, err := policyGetDigest.Execute(t)
	if err != nil {
		return nil, fmt.Errorf("failed to get policy digest: %v", err)
	}

	return &policyGetDigestResponse.PolicyDigest, nil
}

// nolint:gocyclo
func validatePCRBanks(t transport.TPM) error {
	pcrValue, err := ReadPCR(t, secureboot.UKIPCR)
	if err != nil {
		return fmt.Errorf("failed to read PCR: %v", err)
	}

	if err = validatePCRNotZeroAndNotFilled(pcrValue, secureboot.UKIPCR); err != nil {
		return err
	}

	pcrValue, err = ReadPCR(t, secureboot.SecureBootStatePCR)
	if err != nil {
		return fmt.Errorf("failed to read PCR: %v", err)
	}

	if err = validatePCRNotZeroAndNotFilled(pcrValue, secureboot.SecureBootStatePCR); err != nil {
		return err
	}

	caps := tpm2.GetCapability{
		Capability:    tpm2.TPMCapPCRs,
		Property:      0,
		PropertyCount: 1,
	}

	capsResp, err := caps.Execute(t)
	if err != nil {
		return fmt.Errorf("failed to get PCR capabilities: %v", err)
	}

	assignedPCRs, err := capsResp.CapabilityData.Data.AssignedPCR()
	if err != nil {
		return fmt.Errorf("failed to parse assigned PCRs: %v", err)
	}

	for _, s := range assignedPCRs.PCRSelections {
		h, err := s.Hash.Hash()
		if err != nil {
			return fmt.Errorf("failed to parse hash algorithm: %v", err)
		}

		switch h { //nolint:exhaustive
		case crypto.SHA1:
			continue
		case crypto.SHA256:
			// check if 24 banks are available
			if len(s.PCRSelect) != 24/8 {
				return fmt.Errorf("unexpected number of PCR banks: %d", len(s.PCRSelect))
			}

			// check if all banks are available
			if s.PCRSelect[0] != 0xff || s.PCRSelect[1] != 0xff || s.PCRSelect[2] != 0xff {
				return fmt.Errorf("unexpected PCR banks: %v", s.PCRSelect)
			}
		case crypto.SHA384:
			continue
		case crypto.SHA512:
			continue
		default:
			return fmt.Errorf("unsupported hash algorithm: %s", h.String())
		}
	}

	return nil
}

func validatePCRNotZeroAndNotFilled(pcrValue []byte, pcr int) error {
	if bytes.Equal(pcrValue, bytes.Repeat([]byte{0x00}, sha256.Size)) {
		return fmt.Errorf("PCR bank %d is populated with all zeroes", pcr)
	}

	if bytes.Equal(pcrValue, bytes.Repeat([]byte{0xFF}, sha256.Size)) {
		return fmt.Errorf("PCR bank %d is populated with all 0xFF", pcr)
	}

	return nil
}