// Copyright 2023 Canonical Ltd.
// Licensed under the LGPLv3 with static-linking exception.
// See LICENCE file for details.

package policyutil

import (
	"bytes"
	"crypto"
	"errors"
	"fmt"
	"hash/fnv"
	"io"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/canonical/go-tpm2"
	"github.com/canonical/go-tpm2/mu"
)

const (
	pathForbiddenChars = "{}*<>[]/"

	// We use command codes to identify element types. In the case
	// where we need a custom command code for a special element,
	// we set the vendor bit (0x20000000) and also set one of the
	// reserved bits (0xdfff0000)
	commandPolicyBranchNode tpm2.CommandCode = 0x20010171
	commandPolicyPCRValues  tpm2.CommandCode = 0x2002017F
)

var (
	// ErrMissingDigest is returned from [Policy.Execute] when a TPM2_PolicyCpHash or
	// TPM2_PolicyNameHash assertion is missing a digest for the selected session algorithm.
	ErrMissingDigest = errors.New("missing digest for session algorithm")
)

type (
	taskFn       func() error
	authMapKey   uint32
	ticketMapKey uint32
)

func mapKey(vals ...interface{}) uint32 {
	h := fnv.New32()
	mu.MustMarshalToWriter(h, vals...)
	return h.Sum32()
}

func makeAuthMapKey(authName tpm2.Name, policyRef tpm2.Nonce) authMapKey {
	return authMapKey(mapKey(authName, policyRef))
}

func makeTicketMapKey(ticket *PolicyTicket) ticketMapKey {
	return ticketMapKey(mapKey(ticket.AuthName, ticket.PolicyRef, ticket.CpHash))
}

// PolicyTicket corresponds to a ticket generated from a TPM2_PolicySigned or TPM2_PolicySecret
// assertion and is returned by [Policy.Execute]. Generated tickets can be supplied to
// [Policy.Execute] in the future in order to satisfy these assertions as long as they haven't
// expired.
type PolicyTicket struct {
	AuthName  tpm2.Name    // The name of the auth object associated with the corresponding assertion
	PolicyRef tpm2.Nonce   // The policy ref of the corresponding assertion
	CpHash    tpm2.Digest  // The cpHash supplied to the assertion that generated this ticket
	Timeout   tpm2.Timeout // The timeout returned by the assertion that generated this ticket

	// Ticket is the actual ticket returned by the TPM for the assertion that generated this ticket.
	// The Tag field indicates whether this was generated by TPM2_PolicySigned or TPM2_PolicySecret.
	Ticket *tpm2.TkAuth
}

// PolicyError is returned from [Policy.Execute] and other methods when an error
// is encountered during some processing of a policy. It provides an indication of
// where an error occurred.
type PolicyError struct {
	Path string // the path of the branch at which the error occurred

	task string
	err  error
}

type policyDelimiterError interface {
	error
	isPolicyDelimiterError()
}

// makePolicyError returns a PolicyError in the following way:
//   - If the supplied error already contains a PolicyError from the current policy,
//     it is unwrapped to the first error within the current policy and returned.
//   - If the supplied error does not contain a PolicyError from the current policy
//     the error is wrapped with a new PolicyError.
//
// A policy boundary is indicated by the presence of a policyDelimiterError.
func makePolicyError(err error, path policyBranchPath, task string) *PolicyError {
	pErr := &PolicyError{Path: string(path), task: task, err: err}

	var delim policyDelimiterError
	errors.As(err, &delim)

	var pErrNext *PolicyError
	for errors.As(errors.Unwrap(pErr), &pErrNext) {
		if delim != nil {
			var delim2 policyDelimiterError
			if !errors.As(pErrNext, &delim2) || delim2 != delim {
				break
			}
		}
		pErr = pErrNext
	}

	return pErr
}

func (e *PolicyError) Error() string {
	branch := "root branch"
	if len(e.Path) > 0 {
		branch = "branch '" + e.Path + "'"
	}
	return fmt.Sprintf("cannot run '%s' task in %s: %v", e.task, branch, e.err)
}

func (e *PolicyError) Unwrap() error {
	return e.err
}

// PolicyNVError is returned from [Policy.Execute] and other methods when an error
// is encountered when executing a TPM2_PolicyNV assertion. If there was an error
// authorizing use of the NV index with a policy session, this will wrap a
// *[ResourceAuthorizeError].
type PolicyNVError struct {
	Index tpm2.Handle // The NV index handle
	Name  tpm2.Name   // The NV index name

	err error
}

func (e *PolicyNVError) Error() string {
	return fmt.Sprintf("cannot complete assertion with NV index %v (name: %#x): %v", e.Index, e.Name, e.err)
}

func (e *PolicyNVError) Unwrap() error {
	return e.err
}

// PolicyAuthorizationError is returned from [Policy.Execute] if:
//   - the policy uses TPM2_PolicySecret and the associated resource could not be authorized. When
//     this occurs because there was an error loading the associated resource, this will wrap a
//     *[ResourceLoadError]. If there was an error authorizing use of the resource with a policy
//     session, this will wrap a *[ResourceAuthorizeError].
//   - the policy uses TPM2_PolicySigned and no or an invalid signed authorization was supplied.
//   - the policy uses TPM2_PolicyAuthorize and no or an invalid authorized policy was supplied.
type PolicyAuthorizationError struct {
	AuthName  tpm2.Name
	PolicyRef tpm2.Nonce
	err       error
}

func (e *PolicyAuthorizationError) Error() string {
	return fmt.Sprintf("cannot complete authorization with authName=%#x, policyRef=%#x: %v", e.AuthName, e.PolicyRef, e.err)
}

func (e *PolicyAuthorizationError) Unwrap() error {
	return e.err
}

// ResourceLoadError is returned from [Policy.Execute] if the policy uses TPM2_PolicySecret
// and the associated resource could not be loaded. If loading the resource required
// authorization with a policy session and that failed, this will wrap another *[PolicyError].
type ResourceLoadError struct {
	Name tpm2.Name
	err  error
}

func (e *ResourceLoadError) Error() string {
	return fmt.Sprintf("cannot load resource with name %#x: %v", e.Name, e.err)
}

func (e *ResourceLoadError) Unwrap() error {
	return e.err
}

func (*ResourceLoadError) isPolicyDelimiterError() {}

// ResourceAuthorizeError is returned from [Policy.Execute] if an error is encountered
// when trying to authorize a resource required by a policy. This should be wrappped in
// either a *[PolicyNVError] or *[PolicyAuthorizationError] which indicates the assertion
// that the error occurred for. This may wrap another *[PolicyError].
type ResourceAuthorizeError struct {
	Name tpm2.Name
	err  error
}

func (e *ResourceAuthorizeError) Error() string {
	return fmt.Sprintf("cannot authorize resource with name %#x: %v", e.Name, e.err)
}

func (e *ResourceAuthorizeError) Unwrap() error {
	return e.err
}

func (*ResourceAuthorizeError) isPolicyDelimiterError() {}

type policyBranchName string

func (n policyBranchName) isValid() bool {
	if !utf8.ValidString(string(n)) {
		return false
	}
	if strings.ContainsAny(string(n), pathForbiddenChars) {
		return false
	}
	return true
}

func (n policyBranchName) Marshal(w io.Writer) error {
	if !n.isValid() {
		return errors.New("invalid name")
	}
	_, err := mu.MarshalToWriter(w, []byte(n))
	return err
}

func (n *policyBranchName) Unmarshal(r io.Reader) error {
	var b []byte
	if _, err := mu.UnmarshalFromReader(r, &b); err != nil {
		return err
	}
	name := policyBranchName(b)
	if !name.isValid() {
		return errors.New("invalid name")
	}
	*n = name
	return nil
}

type policyBranchPath string

func (p policyBranchPath) PopNextComponent() (next string, remaining policyBranchPath) {
	remaining = p
	for len(remaining) > 0 {
		s := strings.SplitN(string(remaining), "/", 2)
		remaining = ""
		if len(s) == 2 {
			remaining = policyBranchPath(s[1])
		}
		component := s[0]
		if len(component) > 0 {
			return component, remaining
		}
	}

	return "", ""
}

func (p policyBranchPath) Concat(path string) policyBranchPath {
	var pathElements []string
	if p != "" {
		pathElements = append(pathElements, string(p))
	}
	if path != "" {
		pathElements = append(pathElements, path)
	}
	return policyBranchPath(strings.Join(pathElements, "/"))
}

type authorizedPolicy struct {
	policyBranch
	authorization *PolicyAuthorization
}

type authorizationCommandParams struct {
	commandCode tpm2.CommandCode
	handles     []NamedHandle
	params      []interface{}
}

// policySession corresponds to a policy session. This is a limited version
// of PolicyExecuteSession that's used in all code paths in Policy.
type policySession interface {
	Name() tpm2.Name
	HashAlg() tpm2.HashAlgorithmId

	PolicySigned(authKey tpm2.ResourceContext, includeNonceTPM bool, cpHashA tpm2.Digest, policyRef tpm2.Nonce, expiration int32, auth *tpm2.Signature) (tpm2.Timeout, *tpm2.TkAuth, error)
	PolicySecret(authObject tpm2.ResourceContext, cpHashA tpm2.Digest, policyRef tpm2.Nonce, expiration int32, authObjectAuthSession tpm2.SessionContext) (tpm2.Timeout, *tpm2.TkAuth, error)
	PolicyTicket(timeout tpm2.Timeout, cpHashA tpm2.Digest, policyRef tpm2.Nonce, authName tpm2.Name, ticket *tpm2.TkAuth) error
	PolicyOR(pHashList tpm2.DigestList) error
	PolicyPCR(pcrDigest tpm2.Digest, pcrs tpm2.PCRSelectionList) error
	PolicyNV(auth, index tpm2.ResourceContext, operandB tpm2.Operand, offset uint16, operation tpm2.ArithmeticOp, authAuthSession tpm2.SessionContext) error
	PolicyCounterTimer(operandB tpm2.Operand, offset uint16, operation tpm2.ArithmeticOp) error
	PolicyCommandCode(code tpm2.CommandCode) error
	PolicyCpHash(cpHashA tpm2.Digest) error
	PolicyNameHash(nameHash tpm2.Digest) error
	PolicyDuplicationSelect(objectName, newParentName tpm2.Name, includeObject bool) error
	PolicyAuthorize(approvedPolicy tpm2.Digest, policyRef tpm2.Nonce, keySign tpm2.Name, verified *tpm2.TkVerified) error
	PolicyAuthValue() error
	PolicyPassword() error
	PolicyGetDigest() (tpm2.Digest, error)
	PolicyNvWritten(writtenSet bool) error
}

type policyTickets interface {
	ticket(authName tpm2.Name, policyRef tpm2.Nonce) *PolicyTicket
	addTicket(ticket *PolicyTicket)
	invalidTicket(ticket *PolicyTicket)
}

type policyResources interface {
	loadedResource(name tpm2.Name) (ResourceContext, error)
	authorizedPolicies(keySign tpm2.Name, policyRef tpm2.Nonce) ([]*Policy, error)
	signedAuthorization(authKey tpm2.Name, policyRef tpm2.Nonce) (*PolicySignedAuthorization, error)
}

type policyRunner interface {
	session() policySession
	tickets() policyTickets
	resources() policyResources

	authResourceName() tpm2.Name
	loadExternal(public *tpm2.Public) (ResourceContext, error)
	authorize(auth ResourceContext, askForPolicy bool, params *authorizationCommandParams, prefer tpm2.SessionType) (SessionContext, error)
	runBranch(branches policyBranches) (selected int, err error)
	runAuthorizedPolicy(keySign *tpm2.Public, policyRef tpm2.Nonce, policies []*authorizedPolicy) (approvedPolicy tpm2.Digest, checkTicket *tpm2.TkVerified, err error)
	notifyPolicyPCRDigest() error
}

type taggedHash struct {
	HashAlg tpm2.HashAlgorithmId
	Digest  tpm2.Digest
}

func (h taggedHash) Marshal(w io.Writer) error {
	ta := tpm2.MakeTaggedHash(h.HashAlg, h.Digest)
	_, err := mu.MarshalToWriter(w, ta)
	return err
}

func (h *taggedHash) Unmarshal(r io.Reader) error {
	var ta tpm2.TaggedHash
	if _, err := mu.UnmarshalFromReader(r, &ta); err != nil {
		return err
	}

	if ta.HashAlg != tpm2.HashAlgorithmNull && !ta.HashAlg.IsValid() {
		return errors.New("invalid digest algorithm")
	}

	*h = taggedHash{
		HashAlg: ta.HashAlg,
		Digest:  ta.Digest()}
	return nil
}

type taggedHashList []taggedHash

type policyNVElement struct {
	NvIndex   *tpm2.NVPublic
	OperandB  tpm2.Operand
	Offset    uint16
	Operation tpm2.ArithmeticOp
}

func (*policyNVElement) name() string { return "TPM2_PolicyNV assertion" }

func (e *policyNVElement) run(runner policyRunner) (err error) {
	nvIndex, err := tpm2.NewNVIndexResourceContextFromPub(e.NvIndex)
	if err != nil {
		return fmt.Errorf("cannot create nvIndex context: %w", err)
	}

	var auth ResourceContext = newResourceContext(nvIndex, nil)
	askForPolicy := true
	switch {
	case e.NvIndex.Attrs&tpm2.AttrNVPolicyRead != 0:
		// use NV index for auth
	case e.NvIndex.Attrs&tpm2.AttrNVAuthRead != 0:
		// use NV index for auth
	case e.NvIndex.Attrs&tpm2.AttrNVOwnerRead != 0:
		auth, err = runner.resources().loadedResource(tpm2.MakeHandleName(tpm2.HandleOwner))
		askForPolicy = false
	case e.NvIndex.Attrs&tpm2.AttrNVPPRead != 0:
		auth, err = runner.resources().loadedResource(tpm2.MakeHandleName(tpm2.HandlePlatform))
		askForPolicy = false
	default:
		return errors.New("invalid nvIndex read auth mode")
	}
	if err != nil {
		return &PolicyNVError{
			Index: nvIndex.Handle(),
			Name:  nvIndex.Name(),
			err:   fmt.Errorf("cannot create auth context: %w", err),
		}
	}

	params := &authorizationCommandParams{
		commandCode: tpm2.CommandPolicyNV,
		handles:     []NamedHandle{auth.Resource(), nvIndex, runner.session().Name()},
		params:      []any{e.OperandB, e.Offset, e.Operation},
	}

	authSession, err := runner.authorize(auth, askForPolicy, params, tpm2.SessionTypePolicy)
	if err != nil {
		return &PolicyNVError{
			Index: nvIndex.Handle(),
			Name:  nvIndex.Name(),
			err:   &ResourceAuthorizeError{Name: nvIndex.Name(), err: err},
		}
	}
	defer authSession.Flush()

	if err := runner.session().PolicyNV(auth.Resource(), nvIndex, e.OperandB, e.Offset, e.Operation, authSession.Session()); err != nil {
		return &PolicyNVError{Index: nvIndex.Handle(), Name: nvIndex.Name(), err: err}
	}

	return nil
}

type policySecretElement struct {
	AuthObjectName tpm2.Name
	PolicyRef      tpm2.Nonce
	CpHashA        tpm2.Digest
	Expiration     int32
}

func (*policySecretElement) name() string { return "TPM2_PolicySecret assertion" }

func (e *policySecretElement) run(runner policyRunner) (err error) {
	if ticket := runner.tickets().ticket(e.AuthObjectName, e.PolicyRef); ticket != nil {
		err := runner.session().PolicyTicket(ticket.Timeout, ticket.CpHash, ticket.PolicyRef, ticket.AuthName, ticket.Ticket)
		switch {
		case tpm2.IsTPMParameterError(err, tpm2.ErrorExpired, tpm2.CommandPolicyTicket, 1):
			// The ticket has expired - ignore this and fall through to PolicySecret
			runner.tickets().invalidTicket(ticket)
		case tpm2.IsTPMParameterError(err, tpm2.ErrorTicket, tpm2.CommandPolicyTicket, 5):
			// The ticket is invalid - ignore this and fall through to PolicySecret
			runner.tickets().invalidTicket(ticket)
		case err != nil:
			return &PolicyAuthorizationError{AuthName: e.AuthObjectName, PolicyRef: e.PolicyRef, err: err}
		default:
			// The ticket was accepted
			return nil
		}
	}

	authObject, err := runner.resources().loadedResource(e.AuthObjectName)
	if err != nil {
		return &PolicyAuthorizationError{
			AuthName:  e.AuthObjectName,
			PolicyRef: e.PolicyRef,
			err:       &ResourceLoadError{Name: e.AuthObjectName, err: err},
		}
	}
	defer authObject.Flush()

	params := &authorizationCommandParams{
		commandCode: tpm2.CommandPolicySecret,
		handles:     []NamedHandle{authObject.Resource(), runner.session().Name()},
		params:      []any{e.CpHashA, e.PolicyRef, e.Expiration},
	}

	authSession, err := runner.authorize(authObject, false, params, tpm2.SessionTypeHMAC)
	if err != nil {
		return &PolicyAuthorizationError{
			AuthName:  e.AuthObjectName,
			PolicyRef: e.PolicyRef,
			err:       &ResourceAuthorizeError{Name: e.AuthObjectName, err: err},
		}
	}
	defer authSession.Flush()

	timeout, ticket, err := runner.session().PolicySecret(authObject.Resource(), e.CpHashA, e.PolicyRef, e.Expiration, authSession.Session())
	if err != nil {
		return &PolicyAuthorizationError{AuthName: e.AuthObjectName, PolicyRef: e.PolicyRef, err: err}
	}

	runner.tickets().addTicket(&PolicyTicket{
		AuthName:  e.AuthObjectName,
		PolicyRef: e.PolicyRef,
		CpHash:    nil,
		Timeout:   timeout,
		Ticket:    ticket})
	return nil
}

type policySignedElement struct {
	AuthKey   *tpm2.Public
	PolicyRef tpm2.Nonce
	Unused1   tpm2.Digest
	Unused2   int32
}

func (*policySignedElement) name() string { return "TPM2_PolicySigned assertion" }

func (e *policySignedElement) run(runner policyRunner) error {
	authKeyName := e.AuthKey.Name()
	if !authKeyName.IsValid() {
		return errors.New("invalid auth key name")
	}

	if ticket := runner.tickets().ticket(authKeyName, e.PolicyRef); ticket != nil {
		err := runner.session().PolicyTicket(ticket.Timeout, ticket.CpHash, ticket.PolicyRef, ticket.AuthName, ticket.Ticket)
		switch {
		case tpm2.IsTPMParameterError(err, tpm2.ErrorExpired, tpm2.CommandPolicyTicket, 1):
			// The ticket has expired - ignore this and fall through to PolicySigned
			runner.tickets().invalidTicket(ticket)
		case tpm2.IsTPMParameterError(err, tpm2.ErrorTicket, tpm2.CommandPolicyTicket, 5):
			// The ticket is invalid - ignore this and fall through to PolicySigned
			runner.tickets().invalidTicket(ticket)
		case err != nil:
			return &PolicyAuthorizationError{AuthName: authKeyName, PolicyRef: e.PolicyRef, err: err}
		default:
			// The ticket was accepted
			return nil
		}
	}

	auth, err := runner.resources().signedAuthorization(authKeyName, e.PolicyRef)
	if err != nil {
		return &PolicyAuthorizationError{
			AuthName:  authKeyName,
			PolicyRef: e.PolicyRef,
			err:       fmt.Errorf("cannot obtain signed authorization: %w", err),
		}
	}

	authKey, err := runner.loadExternal(e.AuthKey)
	if err != nil {
		return fmt.Errorf("cannot create authKey context: %w", err)
	}
	defer authKey.Flush()

	includeNonceTPM := false
	if len(auth.NonceTPM) > 0 {
		includeNonceTPM = true
	}

	timeout, ticket, err := runner.session().PolicySigned(authKey.Resource(), includeNonceTPM, auth.CpHash, e.PolicyRef, auth.Expiration, auth.PolicyAuthorization.Signature)
	if err != nil {
		return &PolicyAuthorizationError{AuthName: authKeyName, PolicyRef: e.PolicyRef, err: err}
	}

	runner.tickets().addTicket(&PolicyTicket{
		AuthName:  authKeyName,
		PolicyRef: e.PolicyRef,
		CpHash:    auth.CpHash,
		Timeout:   timeout,
		Ticket:    ticket})
	return nil
}

type policyAuthorizations []PolicyAuthorization

type policyAuthorizeElement struct {
	PolicyRef tpm2.Nonce
	KeySign   *tpm2.Public
}

func (*policyAuthorizeElement) name() string { return "authorized policy" }

func (e *policyAuthorizeElement) run(runner policyRunner) error {
	keySignName := e.KeySign.Name()
	if !keySignName.IsValid() {
		return errors.New("invalid keySign")
	}

	policies, err := runner.resources().authorizedPolicies(keySignName, e.PolicyRef)
	if err != nil {
		return &PolicyAuthorizationError{AuthName: keySignName, PolicyRef: e.PolicyRef, err: err}
	}

	// Filter out policies that aren't computed for the current session algorithm or
	// don't have a matching authorization, although we shouldn't really have any
	// without a matching authorization.
	var candidatePolicies []*authorizedPolicy
	for _, policy := range policies {
		digest, err := policy.Digest(runner.session().HashAlg())
		if err == ErrMissingDigest {
			// no suitable digest
			continue
		}
		if err != nil {
			return err
		}

		// Find the signed authorization
		var policyAuth *PolicyAuthorization
		for _, auth := range policy.policy.PolicyAuthorizations {
			if auth.Signature == nil {
				continue
			}
			if !bytes.Equal(auth.AuthKey.Name(), keySignName) {
				continue
			}
			if !bytes.Equal(auth.PolicyRef, e.PolicyRef) {
				continue
			}
			if ok, _ := auth.Verify(digest); !ok {
				continue
			}
			policyAuth = &auth
			break
		}
		if policyAuth == nil {
			// no matching authorization - this shouldn't really happen.
			continue
		}

		candidatePolicies = append(candidatePolicies, &authorizedPolicy{
			policyBranch: policyBranch{
				Name:          policyBranchName(fmt.Sprintf("%x", digest)),
				Policy:        policy.policy.Policy,
				PolicyDigests: taggedHashList{{HashAlg: runner.session().HashAlg(), Digest: digest}},
			},
			authorization: policyAuth,
		})
	}

	approvedPolicy, checkTicket, err := runner.runAuthorizedPolicy(e.KeySign, e.PolicyRef, candidatePolicies)
	if err != nil {
		return &PolicyAuthorizationError{AuthName: keySignName, PolicyRef: e.PolicyRef, err: err}
	}

	if err := runner.session().PolicyAuthorize(approvedPolicy, e.PolicyRef, keySignName, checkTicket); err != nil {
		return &PolicyAuthorizationError{AuthName: keySignName, PolicyRef: e.PolicyRef, err: err}
	}
	return nil
}

type policyAuthValueElement struct{}

func (*policyAuthValueElement) name() string { return "TPM2_PolicyAuthValue assertion" }

func (*policyAuthValueElement) run(runner policyRunner) error {
	return runner.session().PolicyAuthValue()
}

type policyCommandCodeElement struct {
	CommandCode tpm2.CommandCode
}

func (*policyCommandCodeElement) name() string { return "TPM2_PolicyCommandCode assertion" }

func (e *policyCommandCodeElement) run(runner policyRunner) error {
	return runner.session().PolicyCommandCode(e.CommandCode)
}

type policyCounterTimerElement struct {
	OperandB  tpm2.Operand
	Offset    uint16
	Operation tpm2.ArithmeticOp
}

func (*policyCounterTimerElement) name() string { return "TPM2_PolicyCounterTimer assertion" }

func (e *policyCounterTimerElement) run(runner policyRunner) error {
	return runner.session().PolicyCounterTimer(e.OperandB, e.Offset, e.Operation)
}

type cpHashParams struct {
	CommandCode tpm2.CommandCode
	Handles     []tpm2.Name
	CpBytes     []byte
}

type policyCpHashElement struct {
	Digest tpm2.Digest
}

func (*policyCpHashElement) name() string { return "TPM2_PolicyCpHash assertion" }

func (e *policyCpHashElement) run(runner policyRunner) error {
	return runner.session().PolicyCpHash(e.Digest)
}

type policyNameHashElement struct {
	Digest tpm2.Digest
}

func (*policyNameHashElement) name() string { return "TPM2_PolicyNameHash assertion" }

func (e *policyNameHashElement) run(runner policyRunner) error {
	return runner.session().PolicyNameHash(e.Digest)
}

type policyBranch struct {
	Name          policyBranchName
	PolicyDigests taggedHashList
	Policy        policyElements

	index int `tpm2:"ignore"` // index of this branch in a branch node
}

func (b *policyBranch) name() string {
	switch len(b.Name) {
	case 0:
		return "{" + strconv.Itoa(b.index) + "}"
	default:
		return string(b.Name)
	}
}

func (b *policyBranch) nameMatches(pattern string) (bool, error) {
	names := []string{"{" + strconv.Itoa(b.index) + "}"}
	if len(b.Name) > 0 {
		names = append(names, string(b.Name))
	}
	for _, name := range names {
		switch match, err := filepath.Match(pattern, name); {
		case err != nil:
			return false, err
		case match:
			return true, nil
		}
	}
	return false, nil
}

type policyBranches []*policyBranch

func (b policyBranches) filterBranches(pattern string) ([]int, error) {
	var indices []int
	for i, br := range b {
		switch match, err := br.nameMatches(pattern); {
		case err != nil:
			return nil, fmt.Errorf("cannot match: %w", err)
		case match:
			indices = append(indices, i)
		}
	}
	return indices, nil
}

func (b policyBranches) Marshal(w io.Writer) error {
	_, err := mu.MarshalToWriter(w, []*policyBranch(b))
	return err
}

func (b *policyBranches) Unmarshal(r io.Reader) error {
	var branches []*policyBranch
	if _, err := mu.UnmarshalFromReader(r, &branches); err != nil {
		return err
	}

	for i, branch := range branches {
		branch.index = i
	}

	*b = policyBranches(branches)
	return nil
}

type policyORElement struct {
	HashList tpm2.DigestList
}

func (*policyORElement) name() string { return "TPM2_PolicyOR assertion" }

func (e *policyORElement) run(runner policyRunner) error {
	return runner.session().PolicyOR(e.HashList)
}

type policyBranchNodeElement struct {
	Branches policyBranches
}

func (*policyBranchNodeElement) name() string { return "branch node" }

func (e *policyBranchNodeElement) run(runner policyRunner) error {
	if len(e.Branches) == 0 {
		return errors.New("no branches")
	}

	selected, err := runner.runBranch(e.Branches)
	if err != nil {
		return err
	}

	// Obtain the branch digests
	var digests tpm2.DigestList
	for _, branch := range e.Branches {
		found := false
		for _, digest := range branch.PolicyDigests {
			if digest.HashAlg != runner.session().HashAlg() {
				continue
			}

			digests = append(digests, digest.Digest)
			found = true
			break
		}
		if !found {
			return ErrMissingDigest
		}
	}

	tree, err := newPolicyOrTree(runner.session().HashAlg(), digests)
	if err != nil {
		return fmt.Errorf("cannot compute PolicyOR tree: %w", err)
	}

	skipIntermediates := false
	if selected < 0 {
		selected = 0
		skipIntermediates = true
	}
	pHashLists := tree.selectBranch(selected)

	if skipIntermediates {
		return runner.session().PolicyOR(pHashLists[len(pHashLists)-1])
	} else {
		for _, pHashList := range pHashLists {
			if err := runner.session().PolicyOR(pHashList); err != nil {
				return err
			}
		}
	}
	return nil
}

type pcrValue struct {
	PCR    tpm2.Handle
	Digest taggedHash
}

type pcrValueList []pcrValue

type policyPCRValuesElement struct {
	PCRs pcrValueList
}

func (*policyPCRValuesElement) name() string { return "TPM2_PolicyPCR values assertion" }

func (e *policyPCRValuesElement) run(runner policyRunner) error {
	values, err := e.pcrValues()
	if err != nil {
		return err
	}
	pcrs, pcrDigest, err := ComputePCRDigestFromAllValues(runner.session().HashAlg(), values)
	if err != nil {
		return fmt.Errorf("cannot compute PCR digest: %w", err)
	}
	return runner.session().PolicyPCR(pcrDigest, pcrs)
}

func (e *policyPCRValuesElement) pcrValues() (tpm2.PCRValues, error) {
	values := make(tpm2.PCRValues)
	for i, value := range e.PCRs {
		if value.PCR.Type() != tpm2.HandleTypePCR {
			return nil, fmt.Errorf("invalid PCR handle at index %d", i)
		}
		if err := values.SetValue(value.Digest.HashAlg, int(value.PCR), value.Digest.Digest); err != nil {
			return nil, fmt.Errorf("invalid PCR value at index %d: %w", i, err)
		}
	}
	return values, nil
}

type policyPCRElement struct {
	PCRDigest tpm2.Digest
	PCRs      tpm2.PCRSelectionList
}

func (*policyPCRElement) name() string { return "TPM2_PolicyPCR assertion" }

func (e *policyPCRElement) run(runner policyRunner) error {
	if err := runner.notifyPolicyPCRDigest(); err != nil {
		return err
	}
	return runner.session().PolicyPCR(e.PCRDigest, e.PCRs)
}

type policyDuplicationSelectElement struct {
	Object        tpm2.Name
	NewParent     tpm2.Name
	IncludeObject bool
}

func (*policyDuplicationSelectElement) name() string { return "TPM2_PolicyDuplicationSelect assertion" }

func (e *policyDuplicationSelectElement) run(runner policyRunner) error {
	object := e.Object
	if len(object) == 0 && !e.IncludeObject {
		object = runner.authResourceName()
	}
	return runner.session().PolicyDuplicationSelect(object, e.NewParent, e.IncludeObject)
}

type policyPasswordElement struct{}

func (*policyPasswordElement) name() string { return "TPM2_PolicyPassword assertion" }

func (*policyPasswordElement) run(runner policyRunner) error {
	return runner.session().PolicyPassword()
}

type policyNvWrittenElement struct {
	WrittenSet bool
}

func (*policyNvWrittenElement) name() string { return "TPM2_PolicyNvWritten assertion" }

func (e *policyNvWrittenElement) run(runner policyRunner) error {
	return runner.session().PolicyNvWritten(e.WrittenSet)
}

type policyElementDetails struct {
	NV                *policyNVElement
	Secret            *policySecretElement
	Signed            *policySignedElement
	Authorize         *policyAuthorizeElement
	AuthValue         *policyAuthValueElement
	CommandCode       *policyCommandCodeElement
	CounterTimer      *policyCounterTimerElement
	CpHash            *policyCpHashElement
	NameHash          *policyNameHashElement
	OR                *policyORElement
	PCR               *policyPCRElement
	DuplicationSelect *policyDuplicationSelectElement
	Password          *policyPasswordElement
	NvWritten         *policyNvWrittenElement

	BranchNode *policyBranchNodeElement
	PCRValues  *policyPCRValuesElement
}

func (d *policyElementDetails) Select(selector reflect.Value) interface{} {
	switch selector.Interface().(tpm2.CommandCode) {
	case tpm2.CommandPolicyNV:
		return &d.NV
	case tpm2.CommandPolicySecret:
		return &d.Secret
	case tpm2.CommandPolicySigned:
		return &d.Signed
	case tpm2.CommandPolicyAuthorize:
		return &d.Authorize
	case tpm2.CommandPolicyAuthValue:
		return &d.AuthValue
	case tpm2.CommandPolicyCommandCode:
		return &d.CommandCode
	case tpm2.CommandPolicyCounterTimer:
		return &d.CounterTimer
	case tpm2.CommandPolicyCpHash:
		return &d.CpHash
	case tpm2.CommandPolicyNameHash:
		return &d.NameHash
	case tpm2.CommandPolicyOR:
		return &d.OR
	case tpm2.CommandPolicyPCR:
		return &d.PCR
	case tpm2.CommandPolicyDuplicationSelect:
		return &d.DuplicationSelect
	case tpm2.CommandPolicyPassword:
		return &d.Password
	case tpm2.CommandPolicyNvWritten:
		return &d.NvWritten
	case commandPolicyBranchNode:
		return &d.BranchNode
	case commandPolicyPCRValues:
		return &d.PCRValues
	default:
		return nil
	}
}

type policyElementRunner interface {
	name() string
	run(runner policyRunner) error
}

type policyElement struct {
	Type    tpm2.CommandCode
	Details *policyElementDetails
}

func (e *policyElement) runner() policyElementRunner {
	switch e.Type {
	case tpm2.CommandPolicyNV:
		return e.Details.NV
	case tpm2.CommandPolicySecret:
		return e.Details.Secret
	case tpm2.CommandPolicySigned:
		return e.Details.Signed
	case tpm2.CommandPolicyAuthorize:
		return e.Details.Authorize
	case tpm2.CommandPolicyAuthValue:
		return e.Details.AuthValue
	case tpm2.CommandPolicyCommandCode:
		return e.Details.CommandCode
	case tpm2.CommandPolicyCounterTimer:
		return e.Details.CounterTimer
	case tpm2.CommandPolicyCpHash:
		return e.Details.CpHash
	case tpm2.CommandPolicyNameHash:
		return e.Details.NameHash
	case tpm2.CommandPolicyOR:
		return e.Details.OR
	case tpm2.CommandPolicyPCR:
		return e.Details.PCR
	case tpm2.CommandPolicyDuplicationSelect:
		return e.Details.DuplicationSelect
	case tpm2.CommandPolicyPassword:
		return e.Details.Password
	case tpm2.CommandPolicyNvWritten:
		return e.Details.NvWritten
	case commandPolicyBranchNode:
		return e.Details.BranchNode
	case commandPolicyPCRValues:
		return e.Details.PCRValues
	default:
		panic("invalid type")
	}
}

type policyElements []*policyElement

type policy struct {
	PolicyDigests        taggedHashList
	PolicyAuthorizations policyAuthorizations
	Policy               policyElements
}

type nullTickets struct{}

func (*nullTickets) ticket(authName tpm2.Name, policyRef tpm2.Nonce) *PolicyTicket {
	return nil
}

func (*nullTickets) addTicket(ticket *PolicyTicket)     {}
func (*nullTickets) invalidTicket(ticket *PolicyTicket) {}

type mockPolicyResources struct {
	authorized AuthorizedPolicies
}

func newMockPolicyResources(authorizedPolicies AuthorizedPolicies) *mockPolicyResources {
	return &mockPolicyResources{
		authorized: authorizedPolicies,
	}
}

func (*mockPolicyResources) loadedResource(name tpm2.Name) (ResourceContext, error) {
	// the handle is not relevant here
	return newResourceContext(tpm2.NewResourceContext(0x80000000, name), nil), nil
}

func (r *mockPolicyResources) policy(name tpm2.Name) (*Policy, error) {
	return nil, nil
}

func (r *mockPolicyResources) authorizedPolicies(keySign tpm2.Name, policyRef tpm2.Nonce) ([]*Policy, error) {
	if r.authorized == nil {
		return nil, nil
	}
	return r.authorized.AuthorizedPolicies(keySign, policyRef)
}

func (*mockPolicyResources) signedAuthorization(authKey tpm2.Name, policyRef tpm2.Nonce) (*PolicySignedAuthorization, error) {
	return new(PolicySignedAuthorization), nil
}

// Policy corresponds to an authorization policy. It can be serialized with
// [github.com/canonical/go-tpm2/mu].
type Policy struct {
	policy policy
}

// Marshal implements [mu.CustomMarshaller.Marshal].
func (p Policy) Marshal(w io.Writer) error {
	_, err := mu.MarshalToWriter(w, uint32(0), p.policy)
	return err
}

// Unmarshal implements [mu.CustomMarshaller.Unarshal].
func (p *Policy) Unmarshal(r io.Reader) error {
	var version uint32
	_, err := mu.UnmarshalFromReader(r, &version, &p.policy)
	if err != nil {
		return err
	}
	if version != 0 {
		return errors.New("invalid version")
	}
	return nil
}

// Digest returns the digest for this policy for the specified algorithm, if it
// has been computed. If it hasn't been computed, ErrMissingDigest is returned.
func (p *Policy) Digest(alg tpm2.HashAlgorithmId) (tpm2.Digest, error) {
	if !alg.IsValid() {
		return nil, errors.New("invalid algorithm")
	}

	for _, digest := range p.policy.PolicyDigests {
		if digest.HashAlg == alg {
			return digest.Digest, nil
		}
	}

	return nil, ErrMissingDigest
}

// Authorize signs this policy with the supplied signer so that it can be used as an
// authorized policy for a TPM2_PolicyAuthorize assertion with the supplied authKey and
// policyRef. Calling this updates the policy, so it should be persisted afterwards.
// This signs every digest that the policy has been computed for.
//
// TPM2_PolicyAuthorize expects the digest algorithm of the signature to match the name
// algorithm of the public key, so the name algorithm of authKey must match the algorithm
// supplied through the opts argument.
//
// This expects the policy to contain a digest for the selected algorithm already.
func (p *Policy) Authorize(rand io.Reader, authKey *tpm2.Public, policyRef tpm2.Nonce, signer crypto.Signer, opts crypto.SignerOpts) (err error) {
	authName := authKey.Name()
	authAlg := authName.Algorithm()
	if opts.HashFunc() != authAlg.GetHash() {
		return errors.New("mismatched authKey name and opts")
	}
	if !authAlg.Available() {
		return errors.New("auth algorithm is unavailable")
	}

	var authorizations policyAuthorizations
	for _, auth := range p.policy.PolicyAuthorizations {
		if bytes.Equal(auth.AuthKey.Name(), authName) && bytes.Equal(auth.PolicyRef, policyRef) {
			continue
		}
		authorizations = append(authorizations, auth)
	}

	for _, approvedPolicy := range p.policy.PolicyDigests {
		auth, err := SignPolicyAuthorization(rand, approvedPolicy.Digest, authKey, policyRef, signer, opts)
		if err != nil {
			return fmt.Errorf("cannot sign authorization for digest %v: %w", approvedPolicy.HashAlg, err)
		}
		authorizations = append(authorizations, *auth)
	}

	p.policy.PolicyAuthorizations = authorizations
	return nil
}
