// Copyright 2021 Google Inc.
//
// 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 attest

import (
	"bytes"
	"crypto"
	"crypto/rsa"
	"errors"
	"fmt"

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

// secureCurves represents a set of secure elliptic curves. For now,
// the selection is based on the key size only.
var secureCurves = map[tpm2.EllipticCurve]bool{
	tpm2.CurveNISTP256: true,
	tpm2.CurveNISTP384: true,
	tpm2.CurveNISTP521: true,
	tpm2.CurveBNP256:   true,
	tpm2.CurveBNP638:   true,
	tpm2.CurveSM2P256:  true,
}

// CertificationParameters encapsulates the inputs for certifying an application key.
// Only TPM 2.0 is supported at this point.
type CertificationParameters struct {
	// Public represents the key's canonical encoding (a TPMT_PUBLIC structure).
	// It includes the public key and signing parameters.
	Public []byte
	// CreateData represents the properties of a TPM 2.0 key. It is encoded
	// as a TPMS_CREATION_DATA structure.
	CreateData []byte
	// CreateAttestation represents an assertion as to the details of the key.
	// It is encoded as a TPMS_ATTEST structure.
	CreateAttestation []byte
	// CreateSignature represents a signature of the CreateAttestation structure.
	// It is encoded as a TPMT_SIGNATURE structure.
	CreateSignature []byte
}

// VerifyOpts specifies options for the key certification's verification.
type VerifyOpts struct {
	// Public is the public key used to verify key ceritification.
	Public crypto.PublicKey
	// Hash is the hash function used for signature verification. It can be
	// extracted from the properties of the certifying key.
	Hash crypto.Hash
}

// Verify verifies the TPM2-produced certification parameters checking whether:
// - the key length is secure
// - the attestation parameters matched the attested key
// - the key was TPM-generated and resides within TPM
// - the key can sign/decrypt outside-TPM objects
// - the signature is successfuly verified against the passed public key
// For now, it accepts only RSA verification keys.
func (p *CertificationParameters) Verify(opts VerifyOpts) error {
	pub, err := tpm2.DecodePublic(p.Public)
	if err != nil {
		return fmt.Errorf("DecodePublic() failed: %v", err)
	}
	_, err = tpm2.DecodeCreationData(p.CreateData)
	if err != nil {
		return fmt.Errorf("DecodeCreationData() failed: %v", err)
	}
	att, err := tpm2.DecodeAttestationData(p.CreateAttestation)
	if err != nil {
		return fmt.Errorf("DecodeAttestationData() failed: %v", err)
	}
	if att.Type != tpm2.TagAttestCreation {
		return fmt.Errorf("attestation does not apply to creation data, got tag %x", att.Type)
	}

	switch pub.Type {
	case tpm2.AlgRSA:
		if pub.RSAParameters.KeyBits < minRSABits {
			return fmt.Errorf("attested key too small: must be at least %d bits but was %d bits", minRSABits, pub.RSAParameters.KeyBits)
		}
	case tpm2.AlgECC:
		if !secureCurves[pub.ECCParameters.CurveID] {
			return fmt.Errorf("attested key uses insecure curve")
		}
	default:
		return fmt.Errorf("public key of alg 0x%x not supported", pub.Type)
	}

	// Compute & verify that the creation data matches the digest in the
	// attestation structure.
	nameHash, err := pub.NameAlg.Hash()
	if err != nil {
		return fmt.Errorf("HashConstructor() failed: %v", err)
	}
	h := nameHash.New()
	h.Write(p.CreateData)
	if !bytes.Equal(att.AttestedCreationInfo.OpaqueDigest, h.Sum(nil)) {
		return errors.New("attestation refers to different public key")
	}

	// Make sure the key has sane parameters (e.g., attestation can be faked if an AK
	// can be used for arbitrary signatures).
	// We verify the following:
	// - Key is TPM backed.
	// - Key is TPM generated.
	// - Key is not restricted (means it can do arbitrary signing/decrypt ops).
	// - Key cannot be duplicated.
	// - Key was generated by a call to TPM_Create*.
	if att.Magic != tpm20GeneratedMagic {
		return errors.New("creation attestation was not produced by a TPM")
	}
	if (pub.Attributes & tpm2.FlagFixedTPM) == 0 {
		return errors.New("provided key is exportable")
	}
	if (pub.Attributes & tpm2.FlagRestricted) != 0 {
		return errors.New("provided key is restricted")
	}
	if (pub.Attributes & tpm2.FlagFixedParent) == 0 {
		return errors.New("provided key can be duplicated to a different parent")
	}
	if (pub.Attributes & tpm2.FlagSensitiveDataOrigin) == 0 {
		return errors.New("provided key is not created by TPM")
	}

	// Verify the attested creation name matches what is computed from
	// the public key.
	match, err := att.AttestedCreationInfo.Name.MatchesPublic(pub)
	if err != nil {
		return err
	}
	if !match {
		return errors.New("creation attestation refers to a different key")
	}

	// Check the signature over the attestation data verifies correctly.
	// TODO: Support ECC certifying keys
	pk, ok := opts.Public.(*rsa.PublicKey)
	if !ok {
		return fmt.Errorf("Only RSA verification keys are supported")
	}
	hsh := opts.Hash.New()
	hsh.Write(p.CreateAttestation)

	if len(p.CreateSignature) < 8 {
		return fmt.Errorf("signature invalid: length of %d is shorter than 8", len(p.CreateSignature))
	}

	sig, err := tpm2.DecodeSignature(bytes.NewBuffer(p.CreateSignature))
	if err != nil {
		return fmt.Errorf("DecodeSignature() failed: %v", err)
	}

	if err := rsa.VerifyPKCS1v15(pk, opts.Hash, hsh.Sum(nil), sig.RSA.Signature); err != nil {
		return fmt.Errorf("could not verify attestation: %v", err)
	}

	return nil
}