// Copyright 2020 The Monogon Project Authors.
//
// SPDX-License-Identifier: Apache-2.0
//
// 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 tpm

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	tpm2tools "github.com/google/go-tpm-tools/client"
	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpmutil"
	"github.com/pkg/errors"
	"golang.org/x/crypto/nacl/secretbox"
	"golang.org/x/sys/unix"

	"source.monogon.dev/metropolis/pkg/logtree"
	"source.monogon.dev/metropolis/pkg/sysfs"
	tpmpb "source.monogon.dev/metropolis/pkg/tpm/proto"
)

var (
	// SecureBootPCRs are all PCRs that measure the current Secure Boot
	// configuration.  This is what we want if we rely on secure boot to verify
	// boot integrity. The firmware hashes the secure boot policy and custom
	// keys into the PCR.
	//
	// This requires an extra step that provisions the custom keys.
	//
	// Some background: https://mjg59.dreamwidth.org/48897.html?thread=1847297
	// (the initramfs issue mentioned in the article has been solved by
	// integrating it into the kernel binary, and we don't have a shim
	// bootloader)
	//
	// PCR7 alone is not sufficient - it needs to be combined with firmware
	// measurements.
	SecureBootPCRs = []int{7}

	// FirmwarePCRs are alle PCRs that contain the firmware measurements. See:
	//   https://trustedcomputinggroup.org/wp-content/uploads/TCG_EFI_Platform_1_22_Final_-v15.pdf
	FirmwarePCRs = []int{
		0, // platform firmware
		2, // option ROM code
		3, // option ROM configuration and data
	}

	// FullSystemPCRs are all PCRs that contain any measurements up to the
	// currently running EFI payload.
	FullSystemPCRs = []int{
		0, // platform firmware
		1, // host platform configuration
		2, // option ROM code
		3, // option ROM configuration and data
		4, // EFI payload
	}

	// Using FullSystemPCRs is the most secure, but also the most brittle
	// option since updating the EFI binary, updating the platform firmware,
	// changing platform settings or updating the binary would invalidate the
	// sealed data. It's annoying (but possible) to predict values for PCR4,
	// and even more annoying for the firmware PCR (comparison to known values
	// on similar hardware is the only thing that comes to mind).
	//
	// See also: https://github.com/mxre/sealkey (generates PCR4 from EFI
	// image, BSD license)
	//
	// Using only SecureBootPCRs is the easiest and still reasonably secure, if
	// we assume that the platform knows how to take care of itself (i.e. Intel
	// Boot Guard), and that secure boot is implemented properly. It is,
	// however, a much larger amount of code we need to trust.
	//
	// We do not care about PCR 5 (GPT partition table) since modifying it is
	// harmless. All of the boot options and cmdline are hardcoded in the
	// kernel image, and we use no bootloader, so there's no PCR for bootloader
	// configuration or kernel cmdline.
)

var (
	numSRTMPCRs = 16
	srtmPCRs    = tpm2.PCRSelection{Hash: tpm2.AlgSHA256, PCRs: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}}
	// TCG Trusted Platform Module Library Level 00 Revision 0.99 Table 6
	tpmGeneratedValue = uint32(0xff544347)
)

var (
	// ErrNotExists is returned when no TPMs are available in the system
	ErrNotExists = errors.New("no TPMs found")
	// ErrNotInitialized is returned when this package was not initialized
	// successfully
	ErrNotInitialized = errors.New("no TPM was initialized")
)

// Singleton since the TPM is too
var tpm *TPM

// We're serializing all TPM operations since it has a limited number of
// handles and recovering if it runs out is difficult to implement correctly.
// Might also be marginally more secure.
var lock sync.Mutex

// TPM represents a high-level interface to a connected TPM 2.0
type TPM struct {
	logger logtree.LeveledLogger
	device io.ReadWriteCloser

	// We keep the AK loaded since it's used fairly often and deriving it is
	// expensive
	akHandleCache tpmutil.Handle
	akPublicKey   crypto.PublicKey
}

// Initialize finds and opens the TPM (if any). If there is no TPM available it
// returns ErrNotExists
func Initialize(logger logtree.LeveledLogger) error {
	lock.Lock()
	defer lock.Unlock()
	tpmDir, err := os.Open("/sys/class/tpm")
	if err != nil {
		return errors.Wrap(err, "failed to open sysfs TPM class")
	}
	defer tpmDir.Close()

	tpms, err := tpmDir.Readdirnames(2)
	if err != nil {
		return errors.Wrap(err, "failed to read TPM device class")
	}

	if len(tpms) == 0 {
		return ErrNotExists
	}
	if len(tpms) > 1 {
		// If this is changed GetMeasurementLog() needs to be updated too
		logger.Warningf("Found more than one TPM, using the first one")
	}
	tpmName := tpms[0]
	ueventData, err := sysfs.ReadUevents(filepath.Join("/sys/class/tpm", tpmName, "uevent"))
	majorDev, err := strconv.Atoi(ueventData["MAJOR"])
	if err != nil {
		return fmt.Errorf("failed to convert uevent: %w", err)
	}
	minorDev, err := strconv.Atoi(ueventData["MINOR"])
	if err != nil {
		return fmt.Errorf("failed to convert uevent: %w", err)
	}
	if err := unix.Mknod("/dev/tpm", 0600|unix.S_IFCHR, int(unix.Mkdev(uint32(majorDev), uint32(minorDev)))); err != nil {
		return errors.Wrap(err, "failed to create TPM device node")
	}
	device, err := tpm2.OpenTPM("/dev/tpm")
	if err != nil {
		return errors.Wrap(err, "failed to open TPM")
	}
	tpm = &TPM{
		device: device,
		logger: logger,
	}
	return nil
}

// IsInitialized returns true if Initialize was called an at least one
// TPM 2.0 was found and initialized. Otherwise it returns false.
func IsInitialized() bool {
	lock.Lock()
	defer lock.Unlock()
	return !(tpm == nil)
}

// GenerateSafeKey uses two sources of randomness (Kernel & TPM) to generate
// the key
func GenerateSafeKey(size uint16) ([]byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, ErrNotInitialized
	}
	encryptionKeyHost := make([]byte, size)
	if _, err := io.ReadFull(rand.Reader, encryptionKeyHost); err != nil {
		return []byte{}, errors.Wrap(err, "failed to generate host portion of new key")
	}
	var encryptionKeyTPM []byte
	for i := 48; i > 0; i-- {
		tpmKeyPart, err := tpm2.GetRandom(tpm.device, size-uint16(len(encryptionKeyTPM)))
		if err != nil {
			return []byte{}, errors.Wrap(err, "failed to generate TPM portion of new key")
		}
		encryptionKeyTPM = append(encryptionKeyTPM, tpmKeyPart...)
		if len(encryptionKeyTPM) >= int(size) {
			break
		}
	}

	if len(encryptionKeyTPM) != int(size) {
		return []byte{}, fmt.Errorf("got incorrect amount of TPM randomess: %v, requested %v", len(encryptionKeyTPM), size)
	}

	encryptionKey := make([]byte, size)
	for i := uint16(0); i < size; i++ {
		encryptionKey[i] = encryptionKeyHost[i] ^ encryptionKeyTPM[i]
	}
	return encryptionKey, nil
}

// Seal seals sensitive data and only allows access if the current platform
// configuration in matches the one the data was sealed on.
func Seal(data []byte, pcrs []int) ([]byte, error) {
	// Generate a key and use secretbox to encrypt and authenticate the actual
	// payload as go-tpm2 uses a raw seal operation limiting payload size to
	// 128 bytes which is insufficient.
	boxKey, err := GenerateSafeKey(32)
	if err != nil {
		return []byte{}, fmt.Errorf("failed to generate boxKey: %w", err)
	}
	lock.Lock()
	defer lock.Unlock()
	srk, err := tpm2tools.StorageRootKeyRSA(tpm.device)
	if err != nil {
		return []byte{}, errors.Wrap(err, "failed to load TPM SRK")
	}
	defer srk.Close()
	var boxKeyArr [32]byte
	copy(boxKeyArr[:], boxKey)
	// Nonce is not used as we're generating a new boxKey for every operation,
	// therefore we can just leave it all-zero.
	var unusedNonce [24]byte
	encryptedData := secretbox.Seal(nil, data, &unusedNonce, &boxKeyArr)
	sealedKey, err := srk.Seal(boxKey, tpm2tools.SealOpts{Current: tpm2.PCRSelection{Hash: tpm2.AlgSHA256, PCRs: pcrs}})
	if err != nil {
		return []byte{}, fmt.Errorf("failed to seal boxKey: %w", err)
	}
	sealedBytes := tpmpb.ExtendedSealedBytes{
		SealedKey:        sealedKey,
		EncryptedPayload: encryptedData,
	}
	rawSealedBytes, err := proto.Marshal(&sealedBytes)
	if err != nil {
		return []byte{}, errors.Wrapf(err, "failed to marshal sealed data")
	}
	return rawSealedBytes, nil
}

// Unseal unseals sensitive data if the current platform configuration allows
// and sealing constraints allow it.
func Unseal(data []byte) ([]byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, ErrNotInitialized
	}
	srk, err := tpm2tools.StorageRootKeyRSA(tpm.device)
	if err != nil {
		return []byte{}, errors.Wrap(err, "failed to load TPM SRK")
	}
	defer srk.Close()

	var sealedBytes tpmpb.ExtendedSealedBytes
	if err := proto.Unmarshal(data, &sealedBytes); err != nil {
		return []byte{}, errors.Wrap(err, "failed to unmarshal sealed data")
	}
	if sealedBytes.SealedKey == nil {
		return []byte{}, fmt.Errorf("sealed data structure is invalid: no sealed key")
	}
	// Logging this for auditing purposes
	pcrList := []string{}
	for _, pcr := range sealedBytes.SealedKey.Pcrs {
		pcrList = append(pcrList, string(pcr))
	}
	tpm.logger.Infof("Attempting to unseal key protected with PCRs %s", strings.Join(pcrList, ","))
	unsealedKey, err := srk.Unseal(sealedBytes.SealedKey, tpm2tools.UnsealOpts{})
	if err != nil {
		return []byte{}, errors.Wrap(err, "failed to unseal key")
	}
	var key [32]byte
	if len(unsealedKey) != len(key) {
		return []byte{}, fmt.Errorf("unsealed key has wrong length: expected %v bytes, got %v", len(key), len(unsealedKey))
	}
	copy(key[:], unsealedKey)
	var unusedNonce [24]byte
	payload, ok := secretbox.Open(nil, sealedBytes.EncryptedPayload, &unusedNonce, &key)
	if !ok {
		return []byte{}, errors.New("payload box cannot be opened")
	}
	return payload, nil
}

// Standard AK template for RSA2048 non-duplicatable restricted signing for
// attestation
var akTemplate = tpm2.Public{
	Type:       tpm2.AlgRSA,
	NameAlg:    tpm2.AlgSHA256,
	Attributes: tpm2.FlagSignerDefault,
	RSAParameters: &tpm2.RSAParams{
		Sign: &tpm2.SigScheme{
			Alg:  tpm2.AlgRSASSA,
			Hash: tpm2.AlgSHA256,
		},
		KeyBits: 2048,
	},
}

func loadAK() error {
	var err error
	// Rationale: The AK is an EK-equivalent key and used only for attestation.
	// Using a non-primary key here would require us to store the wrapped
	// version somewhere, which is inconvenient.  This being a primary key in
	// the Endorsement hierarchy means that it can always be recreated and can
	// never be "destroyed". Under our security model this is of no concern
	// since we identify a node by its IK (Identity Key) which we can destroy.
	tpm.akHandleCache, tpm.akPublicKey, err = tpm2.CreatePrimary(tpm.device, tpm2.HandleEndorsement,
		tpm2.PCRSelection{}, "", "", akTemplate)
	return err
}

// Process documented in TCG EK Credential Profile 2.2.1
func loadEK() (tpmutil.Handle, crypto.PublicKey, error) {
	// The EK is a primary key which is supposed to be certified by the
	// manufacturer of the TPM.  Its public attributes are standardized in TCG
	// EK Credential Profile 2.0 Table 1. These need to match exactly or we
	// aren't getting the key the manufacturere signed. tpm2tools contains such
	// a template already, so we're using that instead of redoing it ourselves.
	// This ignores the more complicated ways EKs can be specified, the
	// additional stuff you can do is just absolutely crazy (see 2.2.1.2
	// onward)
	return tpm2.CreatePrimary(tpm.device, tpm2.HandleEndorsement,
		tpm2.PCRSelection{}, "", "", tpm2tools.DefaultEKTemplateRSA())
}

// GetAKPublic gets the TPM2T_PUBLIC of the AK key
func GetAKPublic() ([]byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, ErrNotInitialized
	}
	if tpm.akHandleCache == tpmutil.Handle(0) {
		if err := loadAK(); err != nil {
			return []byte{}, fmt.Errorf("failed to load AK primary key: %w", err)
		}
	}
	public, _, _, err := tpm2.ReadPublic(tpm.device, tpm.akHandleCache)
	if err != nil {
		return []byte{}, err
	}
	return public.Encode()
}

// TCG TPM v2.0 Provisioning Guidance v1.0 7.8 Table 2 and TCG EK Credential
// Profile v2.1 2.2.1.4 de-facto Standard for Windows These are both
// non-normative and reference Windows 10 documentation that's no longer
// available :( But in practice this is what people are using, so if it's
// normative or not doesn't really matter
const ekCertHandle = 0x01c00002

// GetEKPublic gets the public key and (if available) Certificate of the EK
func GetEKPublic() ([]byte, []byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, []byte{}, ErrNotInitialized
	}
	ekHandle, publicRaw, err := loadEK()
	if err != nil {
		return []byte{}, []byte{}, fmt.Errorf("failed to load EK primary key: %w", err)
	}
	defer tpm2.FlushContext(tpm.device, ekHandle)
	// Don't question the use of HandleOwner, that's the Standard™
	ekCertRaw, err := tpm2.NVReadEx(tpm.device, ekCertHandle, tpm2.HandleOwner, "", 0)
	if err != nil {
		return []byte{}, []byte{}, err
	}

	publicKey, err := x509.MarshalPKIXPublicKey(publicRaw)
	if err != nil {
		return []byte{}, []byte{}, err
	}

	return publicKey, ekCertRaw, nil
}

// MakeAKChallenge generates a challenge for TPM residency and attributes of
// the AK
func MakeAKChallenge(ekPubKey, akPub []byte, nonce []byte) ([]byte, []byte, error) {
	ekPubKeyData, err := x509.ParsePKIXPublicKey(ekPubKey)
	if err != nil {
		return []byte{}, []byte{}, fmt.Errorf("failed to decode EK pubkey: %w", err)
	}
	akPubData, err := tpm2.DecodePublic(akPub)
	if err != nil {
		return []byte{}, []byte{}, fmt.Errorf("failed to decode AK public part: %w", err)
	}
	// Make sure we're attesting the right attributes (in particular Restricted)
	if !akPubData.MatchesTemplate(akTemplate) {
		return []byte{}, []byte{}, errors.New("the key being challenged is not a valid AK")
	}
	akName, err := akPubData.Name()
	if err != nil {
		return []byte{}, []byte{}, fmt.Errorf("failed to derive AK name: %w", err)
	}
	return generateRSA(akName.Digest, ekPubKeyData.(*rsa.PublicKey), 16, nonce, rand.Reader)
}

// SolveAKChallenge solves a challenge for TPM residency of the AK
func SolveAKChallenge(credBlob, secretChallenge []byte) ([]byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, ErrNotInitialized
	}
	if tpm.akHandleCache == tpmutil.Handle(0) {
		if err := loadAK(); err != nil {
			return []byte{}, fmt.Errorf("failed to load AK primary key: %w", err)
		}
	}

	ekHandle, _, err := loadEK()
	if err != nil {
		return []byte{}, fmt.Errorf("failed to load EK: %w", err)
	}
	defer tpm2.FlushContext(tpm.device, ekHandle)

	// This is necessary since the EK requires an endorsement handle policy in
	// its session.  For us this is stupid because we keep all hierarchies open
	// anyways since a) we cannot safely store secrets on the OS side
	// pre-global unlock and b) it makes no sense in this security model since
	// an uncompromised host OS will not let an untrusted entity attest as
	// itself and a compromised OS can either not pass PCR policy checks or the
	// game's already over (you successfully runtime-exploited a production
	// Metropolis node).
	endorsementSession, _, err := tpm2.StartAuthSession(
		tpm.device,
		tpm2.HandleNull,
		tpm2.HandleNull,
		make([]byte, 16),
		nil,
		tpm2.SessionPolicy,
		tpm2.AlgNull,
		tpm2.AlgSHA256)
	if err != nil {
		panic(err)
	}
	defer tpm2.FlushContext(tpm.device, endorsementSession)

	_, _, err = tpm2.PolicySecret(tpm.device, tpm2.HandleEndorsement, tpm2.AuthCommand{Session: tpm2.HandlePasswordSession, Attributes: tpm2.AttrContinueSession}, endorsementSession, nil, nil, nil, 0)
	if err != nil {
		return []byte{}, fmt.Errorf("failed to make a policy secret session: %w", err)
	}

	for {
		solution, err := tpm2.ActivateCredentialUsingAuth(tpm.device, []tpm2.AuthCommand{
			// Use standard no-password authenatication
			{Session: tpm2.HandlePasswordSession, Attributes: tpm2.AttrContinueSession},
			// Use a full policy session for the EK
			{Session: endorsementSession, Attributes: tpm2.AttrContinueSession},
		}, tpm.akHandleCache, ekHandle, credBlob, secretChallenge)
		if warn, ok := err.(tpm2.Warning); ok && warn.Code == tpm2.RCRetry {
			time.Sleep(100 * time.Millisecond)
			continue
		}
		return solution, err
	}
}

// FlushTransientHandles flushes all sessions and non-persistent handles
func FlushTransientHandles() error {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return ErrNotInitialized
	}
	flushHandleTypes := []tpm2.HandleType{tpm2.HandleTypeTransient, tpm2.HandleTypeLoadedSession, tpm2.HandleTypeSavedSession}
	for _, handleType := range flushHandleTypes {
		handles, err := tpm2tools.Handles(tpm.device, handleType)
		if err != nil {
			return err
		}
		for _, handle := range handles {
			if err := tpm2.FlushContext(tpm.device, handle); err != nil {
				return err
			}
		}
	}
	return nil
}

// AttestPlatform performs a PCR quote using the AK and returns the quote and
// its signature
func AttestPlatform(nonce []byte) ([]byte, []byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return []byte{}, []byte{}, ErrNotInitialized
	}
	if tpm.akHandleCache == tpmutil.Handle(0) {
		if err := loadAK(); err != nil {
			return []byte{}, []byte{}, fmt.Errorf("failed to load AK primary key: %w", err)
		}
	}
	// We only care about SHA256 since SHA1 is weak. This is supported on at
	// least GCE and Intel / AMD fTPM, which is good enough for now. Alg is
	// null because that would just hash the nonce, which is dumb.
	quote, signature, err := tpm2.Quote(tpm.device, tpm.akHandleCache, "", "", nonce, srtmPCRs,
		tpm2.AlgNull)
	if err != nil {
		return []byte{}, []byte{}, fmt.Errorf("failed to quote PCRs: %w", err)
	}
	return quote, signature.RSA.Signature, err
}

// VerifyAttestPlatform verifies a given attestation. You can rely on all data
// coming back as being from the TPM on which the AK is bound to.
func VerifyAttestPlatform(nonce, akPub, quote, signature []byte) (*tpm2.AttestationData, error) {
	hash := crypto.SHA256.New()
	hash.Write(quote)

	akPubData, err := tpm2.DecodePublic(akPub)
	if err != nil {
		return nil, fmt.Errorf("invalid AK: %w", err)
	}
	akPublicKey, err := akPubData.Key()
	if err != nil {
		return nil, fmt.Errorf("invalid AK: %w", err)
	}
	akRSAKey, ok := akPublicKey.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("invalid AK: invalid key type")
	}

	if err := rsa.VerifyPKCS1v15(akRSAKey, crypto.SHA256, hash.Sum(nil), signature); err != nil {
		return nil, err
	}

	quoteData, err := tpm2.DecodeAttestationData(quote)
	if err != nil {
		return nil, err
	}
	// quoteData.Magic works together with the TPM's Restricted key attribute.
	// If this attribute is set (which it needs to be for the AK to be
	// considered valid) the TPM will not sign external data having this prefix
	// with such a key. Only data that originates inside the TPM like quotes
	// and key certifications can have this prefix and sill be signed by a
	// restricted key. This check is thus vital, otherwise somebody can just
	// feed the TPM an arbitrary attestation to sign with its AK and this
	// function will happily accept the forged attestation.
	if quoteData.Magic != tpmGeneratedValue {
		return nil, errors.New("invalid TPM quote: data marker for internal data not set - forged attestation")
	}
	if quoteData.Type != tpm2.TagAttestQuote {
		return nil, errors.New("invalid TPM qoute: not a TPM quote")
	}
	if !bytes.Equal(quoteData.ExtraData, nonce) {
		return nil, errors.New("invalid TPM quote: wrong nonce")
	}

	return quoteData, nil
}

// GetPCRs returns all SRTM PCRs in-order
func GetPCRs() ([][]byte, error) {
	lock.Lock()
	defer lock.Unlock()
	if tpm == nil {
		return [][]byte{}, ErrNotInitialized
	}
	pcrs := make([][]byte, numSRTMPCRs)

	// The TPM can (and most do) return partial results. Let's just retry as
	// many times as we have PCRs since each read should return at least one
	// PCR.
readLoop:
	for i := 0; i < numSRTMPCRs; i++ {
		sel := tpm2.PCRSelection{Hash: tpm2.AlgSHA256}
		for pcrN := 0; pcrN < numSRTMPCRs; pcrN++ {
			if len(pcrs[pcrN]) == 0 {
				sel.PCRs = append(sel.PCRs, pcrN)
			}
		}

		readPCRs, err := tpm2.ReadPCRs(tpm.device, sel)
		if err != nil {
			return nil, fmt.Errorf("failed to read PCRs: %w", err)
		}

		for pcrN, pcr := range readPCRs {
			pcrs[pcrN] = pcr
		}
		for _, pcr := range pcrs {
			// If at least one PCR is still not read, continue
			if len(pcr) == 0 {
				continue readLoop
			}
		}
		break
	}

	return pcrs, nil
}

// GetMeasurmentLog returns the binary log of all data hashed into PCRs. The
// result can be parsed by eventlog.  As this library currently doesn't support
// extending PCRs it just returns the log as supplied by the EFI interface.
func GetMeasurementLog() ([]byte, error) {
	return os.ReadFile("/sys/kernel/security/tpm0/binary_bios_measurements")
}