// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package firmware

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path/filepath"
	"reflect"
	"strings"
	"time"

	"boringssl.googlesource.com/boringssl/util/fipstools/acvp/acvptool/subprocess"

	"chromiumos/tast/common/testexec"
	"chromiumos/tast/errors"
	"chromiumos/tast/shutil"
	"chromiumos/tast/testing"
)

type acv struct {
	Version string `json:"acvVersion"`
}

type vectors struct {
	Retry    uint64 `json:"retry,omitempty"`
	ID       uint64 `json:"vsId"`
	Algo     string `json:"algorithm,omitempty"`
	Revision string `json:"revision,omitempty"`
	Mode     string `json:"mode,omitempty"`
	IsSample bool   `json:"isSample,omitempty"`
}

func init() {
	testing.AddTest(&testing.Test{
		Func: ACVP,
		Contacts: []string{
			"gurleengrewal@chromium.org", // Test author
			"sukhomlinov@chromium.org",   // CR50 certification lead
		},
		Desc:         "Takes a JSON generated by the ACVP server and runs the test cases in it",
		SoftwareDeps: []string{"chrome", "tpm"},
		Params: []testing.Param{{
			Name: "aes_ecb_full",
			Val: data{
				inputFile:    "aes-ecb-full.json",
				expectedFile: "aes-ecb-full-expected.json",
			},
			ExtraData: []string{
				"aes-ecb-full.json",
				"aes-ecb-full-expected.json",
			},
		}, {
			Name: "aes_ecb_short",
			Val: data{
				inputFile:    "aes-ecb-short.json",
				expectedFile: "aes-ecb-short-expected.json",
			},
			ExtraData: []string{
				"aes-ecb-short.json",
				"aes-ecb-short-expected.json",
			},
		}, {
			Name: "aes_cbc_full",
			Val: data{
				inputFile:    "aes-cbc-full.json",
				expectedFile: "aes-cbc-full-expected.json",
			},
			ExtraData: []string{
				"aes-cbc-full.json",
				"aes-cbc-full-expected.json",
			},
		}, {
			Name: "aes_cbc_short",
			Val: data{
				inputFile:    "aes-cbc-short.json",
				expectedFile: "aes-cbc-short-expected.json",
			},
			ExtraData: []string{
				"aes-cbc-short.json",
				"aes-cbc-short-expected.json",
			},
		}, {
			Name: "sha2_256_full",
			Val: data{
				inputFile:    "sha2-256-full.json",
				expectedFile: "sha2-256-full-expected.json",
			},
			ExtraData: []string{
				"sha2-256-full.json",
				"sha2-256-full-expected.json",
			},
		}, {
			Name: "sha2_256_short",
			Val: data{
				inputFile:    "sha2-256-short.json",
				expectedFile: "sha2-256-short-expected.json",
			},
			ExtraData: []string{
				"sha2-256-short.json",
				"sha2-256-short-expected.json",
			},
		}, {
			Name: "sha2_384_full",
			Val: data{
				inputFile:    "sha2-384-full.json",
				expectedFile: "sha2-384-full-expected.json",
			},
			ExtraData: []string{
				"sha2-384-full.json",
				"sha2-384-full-expected.json",
			},
		}, {
			Name: "sha2_384_short",
			Val: data{
				inputFile:    "sha2-384-short.json",
				expectedFile: "sha2-384-short-expected.json",
			},
			ExtraData: []string{
				"sha2-384-short.json",
				"sha2-384-short-expected.json",
			},
		}, {
			Name: "sha2_512_full",
			Val: data{
				inputFile:    "sha2-512-full.json",
				expectedFile: "sha2-512-full-expected.json",
			},
			ExtraData: []string{
				"sha2-512-full.json",
				"sha2-512-full-expected.json",
			},
		}, {
			Name: "sha2_512_short",
			Val: data{
				inputFile:    "sha2-512-short.json",
				expectedFile: "sha2-512-short-expected.json",
			},
			ExtraData: []string{
				"sha2-512-short.json",
				"sha2-512-short-expected.json",
			},
		}, {
			Name: "hmac_drbg",
			Val: data{
				inputFile:    "drbg-test.json",
				expectedFile: "drbg-expected.json",
			},
			ExtraData: []string{
				"drbg-test.json",
				"drbg-expected.json",
			},
		}, {
			Name: "hmac_sha2_256",
			Val: data{
				inputFile:    "hmac-sha2-256-test.json",
				expectedFile: "hmac-sha2-256-expected.json",
			},
			ExtraData: []string{
				"hmac-sha2-256-test.json",
				"hmac-sha2-256-expected.json",
			},
		}, {
			Name: "ecdsa_keygen",
			Val: data{
				inputFile: "ecdsa-keygen-test.json",
			},
			ExtraData: []string{
				"ecdsa-keygen-test.json",
			},
		}, {
			Name: "ecdsa_keyver",
			Val: data{
				inputFile:    "ecdsa-keyver-test.json",
				expectedFile: "ecdsa-keyver-expected.json",
			},
			ExtraData: []string{
				"ecdsa-keyver-test.json",
				"ecdsa-keyver-expected.json",
			},
		}, {
			Name: "ecdsa_siggen",
			Val: data{
				inputFile: "ecdsa-siggen-test.json",
			},
			ExtraData: []string{
				"ecdsa-siggen-test.json",
			},
		}, {
			Name: "ecdsa_sigver",
			Val: data{
				inputFile:    "ecdsa-sigver-test.json",
				expectedFile: "ecdsa-sigver-expected.json",
			},
			ExtraData: []string{
				"ecdsa-sigver-test.json",
				"ecdsa-sigver-expected.json",
			},
		}, {
			Name: "u2f_hmac",
			Val: data{
				inputFile:    "U2F_request_HMAC-SHA2-256_555822.json",
				expectedFile: "U2F_expected_HMAC-SHA2-256_555822.json",
			},
			ExtraData: []string{
				"U2F_request_HMAC-SHA2-256_555822.json",
				"U2F_expected_HMAC-SHA2-256_555822.json",
			},
		}},
		Timeout: time.Hour * 10,
	})
}

const (
	cr50HeaderSize     = 12
	cr50RespHeaderSize = 12
	ecb                = "AES"
	cbc                = "AES-CBC"
	// Trunks hash cmd modes are as follows:
	// 0 - start, 1 - cont., 2 - finish, 3 - single
	// 4 - SW HMAC single shot (TPM code)
	// 5 - HW HMAC SHA256 single shot (dcrypto code)
	tpmHashCmdMode = "03"
	tpmHMACCmdMode = "05"
	tpmSHA256      = "0B"
)

// Holds test data information for each test type.
type data struct {
	inputFile    string
	expectedFile string
}

type opType int

const (
	_ opType = iota
	tpmSigGen
	tpmSigVer
	tpmKeyGen
	tpmKeyVer
)

func newOp(s string) (opType, error) {
	switch s {
	case "sigGen":
		return tpmSigGen, nil
	case "sigVer":
		return tpmSigVer, nil
	case "keyGen":
		return tpmKeyGen, nil
	case "keyVer":
		return tpmKeyVer, nil
	default:
		return 0, errors.Errorf("unknown ECDSA op type %q", s)
	}
}

func (o opType) hex() string {
	switch o {
	case tpmSigGen:
		return "06"
	case tpmSigVer:
		return "05"
	case tpmKeyGen:
		return "02"
	case tpmKeyVer:
		return "04"
	default:
		panic(fmt.Sprintf("Unknown open type %v", o))
	}
}

type option func(*ecdsa) error

// Holds trunks parameters for an ECDSA operation.
type ecdsa struct {
	curve  string
	r      string
	s      string
	qx     string
	qy     string
	digest string
	hash   string
	d      string
	op     opType
}

// ecdsaCurve returns an option that can be passed to newECDSA
// to set the elliptic curve for an ECDSA operation.
func ecdsaCurve(curvehex string) option {
	return func(e *ecdsa) error {
		curve, err := hex.DecodeString(curvehex)
		if err != nil {
			return errors.Errorf("unable to decode curve hex value: %q", curvehex)
		}
		if string(curve) == "P-256" {
			e.curve = "03" // 03 is the trunks constant to indicate the value of P-256
			return nil
		}
		return errors.Errorf("unsupported curve: %q", curve)
	}
}

// hashAlg returns an option that can be passed to newECDSA to
// set the hashing algorithm for the ECDSA operation.
func hashAlg(algBytes string) option {
	return func(e *ecdsa) error {
		switch e.op {
		case tpmSigVer, tpmSigGen:
			hash, err := hex.DecodeString(algBytes)
			if err != nil {
				return errors.Errorf("unable to decode hash hex value: %q", algBytes)
			}
			if string(hash) == "SHA2-256" {
				e.hash = tpmSHA256
			} else {
				return errors.Errorf("unsupported hash algorithm for ECDSA: %q", hash)
			}
		}
		return nil
	}
}

// ecdsaOp returns an option that can be passed to newECDSA to
// set a supported ECDSA operation to its associated trunks command values.
// Trunks operations values are:
// TEST_SIGN = 0, TEST_VERIFY = 1, TEST_KEYGEN = 2, TEST_KEYDERIVE = 3, TEST_POINT = 4, TEST_VERIFY_ANY = 5
func ecdsaOp(op string) option {
	return func(e *ecdsa) error {
		switch op {
		case "keyGen":
			e.op = tpmKeyGen
		case "keyVer":
			e.op = tpmKeyVer
		case "sigGen":
			e.op = tpmSigGen
		case "sigVer":
			e.op = tpmSigVer
		default:
			return errors.Errorf("unsupported ECDSA operation: %q", op)
		}
		return nil
	}
}

// digest returns an option that can be passed to newECDSA to
// compute and set the digest of an ECDSA message.
func digest(msg string) option {
	return func(e *ecdsa) error {
		if e.op != tpmSigGen && e.op != tpmSigVer {
			return nil //No digest needed
		}
		msgBytes, err := hex.DecodeString(msg)
		if err != nil {
			return err
		}
		msgDigest := sha256.Sum256(msgBytes)
		e.digest = hex.EncodeToString(msgDigest[:])
		return nil
	}
}

// ecdsaSigParams returns an option that can be passed to newECDSA to
// set the r, s, d, and Q values in an ECDSA operation.
func ecdsaSigParams(r, s, d, qx, qy string) option {
	return func(e *ecdsa) error {
		e.r = r
		e.s = s
		e.d = d
		e.qx = qx
		e.qy = qy
		return nil
	}
}

// newECDSA returns a new ECDSA struct for a signing or verification operation.
// The op (type of operation) option is mandatory and must be set before other
// options are set.
func newECDSA(op opType, opts ...option) (*ecdsa, error) {
	e := &ecdsa{op: op}
	for _, opt := range opts {
		if err := opt(e); err != nil {
			return nil, err
		}
	}
	return e, nil
}

// Holds trunks parameters for a hash operation.
type hashPrimitive struct {
	msg string
	alg string
	key string
	cmd string
}

func (hp *hashPrimitive) setAlg(alg string) error {
	// Set the hash_mode for use in call to trunks command
	// Command structure is as follows:
	// field     |    size  |                  note
	// ===================================================================
	// mode      |    1     | 0 - start, 1 - cont., 2 - finish, 3 - single
	// hash_mode |    1     | 0 - sha1, 1 - sha256, 2 - sha384, 3 - sha512
	// handle    |    1     | seassion handle, ignored in 'single' mode
	// text_len  |    2     | size of the text to process, big endian
	// text      | text_len | text to hash
	switch alg {
	case "SHA-1", "HMAC-SHA-1":
		hp.alg = "00"
	case "SHA2-256", "HMAC-SHA2-256":
		hp.alg = "01"
	case "SHA2-384", "HMAC-SHA2-384":
		hp.alg = "02"
	case "SHA2-512", "HMAC-SHA2-512":
		hp.alg = "03"
	default:
		return errors.Errorf("unsupported algorithm: %s", alg)
	}
	return nil
}

// setCmd sets the hash cmd for use in call to trunks command
func (hp *hashPrimitive) setCmd(alg, key string) error {
	switch alg {
	case "SHA-1", "SHA2-256", "SHA2-384", "SHA2-512":
		hp.cmd = tpmHashCmdMode
		if key != "" {
			return errors.Errorf("unexpected key value in hash algorithm: %q", key)
		}
	case "HMAC-SHA-1", "HMAC-SHA2-256", "HMAC-SHA2-384", "HMAC-SHA2-512":
		hp.cmd = tpmHMACCmdMode
		if key == "" {
			return errors.New("missing key value for HMAC")
		}
	default:
		return errors.Errorf("unsupported algorithm: %s", alg)
	}
	return nil
}

// newSHA returns a new hashPrimitive struct for SHA
func newSHA(msg, alg, key string) (*hashPrimitive, error) {
	hp := &hashPrimitive{msg: msg, key: key}
	if err := hp.setAlg(alg); err != nil {
		return nil, err
	}
	if err := hp.setCmd(alg, key); err != nil {
		return nil, err
	}
	return hp, nil
}

// Holds trunks parameters for a block cipher operation.
type blockCipher struct {
	encrypt string
	mode    string
	key     string
	in      string
	iv      string
}

func (bc *blockCipher) setMode(mode string) error {
	switch mode {
	case cbc:
		bc.mode = "02"
	case ecb:
		bc.mode = "00"
	default:
		return errors.Errorf("unsupported mode: %s", mode)
	}
	return nil
}

func (bc *blockCipher) setEnc(encrypt string) error {
	switch encrypt {
	case "encrypt":
		bc.encrypt = "01"
	case "decrypt":
		bc.encrypt = "00"
	default:
		return errors.New("encryption string should contain 'encrypt' or 'decrypt'")
	}
	return nil
}

// newAES returns a new blockCipher struct for AES.
func newAES(dir, key, in, iv, mode string) (*blockCipher, error) {
	bc := &blockCipher{key: key, in: in}
	if mode == ecb && iv != "" {
		return nil, errors.Errorf("have IV in ECB mode: %q", iv)
	}
	bc.iv = iv
	if err := bc.setMode(mode); err != nil {
		return nil, err
	}
	if err := bc.setEnc(dir); err != nil {
		return nil, err
	}
	return bc, nil
}

// Holds trunks parameters for an HMAC DRBG based on SHA256 operation.
type drbgSHA256 struct {
	entropy string
	perso   string
	input   string
	input2  string
	nonce   string
	outLen  uint32
}

// newDRBGSHA256 returns a new HMAC-SHA256-DRBG struct.
func newDRBGSHA256(primitive, outLenStr, entropy, perso, input, input2, nonce string) (*drbgSHA256, error) {
	d := &drbgSHA256{entropy: entropy, perso: perso, input: input, input2: input2, nonce: nonce}
	if primitive != "SHA2-256" {
		return nil, errors.Errorf("HMAC DRBG requesting unsupported hash primitive: %q", primitive)
	}
	outLenBytes, err := hex.DecodeString(outLenStr)
	if err != nil {
		return nil, errors.Errorf("unable to decode required output length from argument: %q", outLenStr)
	}
	d.outLen = binary.LittleEndian.Uint32(outLenBytes)
	if d.outLen > 128 {
		return nil, errors.Errorf("DRBG requested too many bytes: %d, maximum is 128", d.outLen)
	}
	return d, nil
}

// Writes commands to DUT and reads the output.
// This interface is required by acvptool.
type cr50IO struct {
	ctx    context.Context
	outBuf bytes.Buffer // ACVPtool output goes here
}

// Runs a trunks command on the DUT.
func (w *cr50IO) Write(b []byte) (int, error) {
	cmdArgs, numRes, err := getTrunksCmds(b)
	if err != nil {
		return 0, errors.Wrap(err, "getTrunksCmd failed")
	}
	for i, cmdArg := range cmdArgs {
		cmd := testexec.CommandContext(w.ctx, "trunks_send", "--raw", cmdArg)
		out, err := cmd.Output(testexec.DumpLogOnError)
		if err != nil {
			return 0, errors.Wrap(err, shutil.EscapeSlice(cmd.Args))
		}
		// For some reason there is an extra byte at the end of command output
		if len(out) > 0 && len(out)%2 != 0 {
			out = out[:len(out)-1]
		}
		respBytes, err := extractTrunksResp(out)
		if err != nil {
			return 0, errors.Wrapf(err, "extractTrunksResp failed after sending trunks command %q", cmdArg)
		}
		// Populate OutBuf with response bytes if we are at the last command.
		if i == len(cmdArgs)-1 {
			if err := w.populateOutBuf(respBytes, numRes); err != nil {
				return 0, errors.Wrapf(err, "populateOutBuf failed after sending trunks command %q and parsing response %q", cmdArg, respBytes)
			}
		}
	}
	return len(b), nil
}

// Close does nothing, but is required to implement the ReadWriteCloser interface.
func (w *cr50IO) Close() error {
	return nil
}

// Read the output of the trunks command.
func (w *cr50IO) Read(b []byte) (int, error) {
	return w.outBuf.Read(b)
}

// populateOutBuf takes a cr50 command response and
// converts to output consumable by ACVPtool.
// b is the cr50 command response.
// n is the expected number of parameters in the response.
func (w *cr50IO) populateOutBuf(b []byte, n uint8) error {
	respSizes := make([]uint32, n)
	responses := make([][]byte, n)
	if n == 1 {
		respSizes[0] = uint32(len(b))
		responses[0] = b
	} else {
		// Check that first byte is hard coded to 1 and skip it.
		if b[0] != 1 {
			return errors.Errorf("unexpected first byte of response: got %q; want 0x01", b[0])
		}
		b = b[1:]
		for i := uint8(0); i < n; i++ {
			respSizes[i] = uint32(binary.BigEndian.Uint16(b))
			responses[i] = b[2 : respSizes[i]+2]
			b = b[respSizes[i]+2:]
		}
	}

	w.outBuf.Write([]byte{byte(n), 00, 00, 00}) // num responses
	respSizeBytes := make([]byte, 4)
	for i := uint8(0); i < n; i++ {
		binary.LittleEndian.PutUint32(respSizeBytes, respSizes[i])
		w.outBuf.Write(respSizeBytes)
	}
	for i := uint8(0); i < n; i++ {
		w.outBuf.Write(responses[i])
	}
	return nil
}

// getTrunksCmds converts contents of b into one or more trunks commands.
// Returns command, number of operations, and possible error value.
func getTrunksCmds(b []byte) ([]string, uint8, error) {
	args, err := parseInBuf(b)
	if err != nil {
		return nil, 0, err
	}

	// The first index in array contains algorithm arguments separated by '/'
	algArgs := strings.Split(args[0], "/")
	algType := algArgs[0]
	algArgs = algArgs[1:]
	algArgs = append(algArgs, args[1:]...)
	argLen := len(algArgs)
	switch strings.Split(algType, "-")[0] {
	case "AES":
		var iv string
		switch argLen {
		case 3:
			iv = ""
		case 4:
			iv = algArgs[3]
		default:
			return nil, 0, errors.Errorf("incorrect number of args for AES: got %d, want 3 or 4", argLen)
		}
		bc, err := newAES(algArgs[0], algArgs[1], algArgs[2], iv, algType)
		if err != nil {
			return nil, 0, err
		}
		return []string{getAESCommand(bc)}, 1, nil
	case "SHA", "SHA2", "HMAC":
		var key string
		switch argLen {
		case 1:
			key = ""
		case 2:
			key = algArgs[1]
		default:
			return nil, 0, errors.Errorf("incorrect number of args for hash/HMAC operation: got %d, want 1 or 2", argLen)
		}
		hp, err := newSHA(algArgs[0], algType, key)
		if err != nil {
			return nil, 0, err
		}
		return []string{getHashCommand(hp)}, 1, nil
	case "hmacDRBG":
		if argLen != 7 {
			return nil, 0, errors.Errorf("incorrect number of args for DRBG: got %d, want 7", argLen)
		}
		d, err := newDRBGSHA256(algArgs[0], algArgs[1], algArgs[2], algArgs[3], algArgs[4], algArgs[5], algArgs[6])
		if err != nil {
			return nil, 0, err
		}
		return getDRBGCommands(d), 1, nil
	case "ECDSA":
		var curve, r, s, d, qx, qy, msg, hash, op string
		if argLen > 1 {
			op = algArgs[0]
			curve = algArgs[1]
		}
		var numRes uint8 = 1
		switch argLen {
		case 2: //keygen
			numRes = 3
		case 4: //keyver
			qx = algArgs[2]
			qy = algArgs[3]
		case 5: //siggen
			d = algArgs[2]
			hash = algArgs[3]
			msg = algArgs[4]
			numRes = 2
		case 8: //sigver
			hash = algArgs[2]
			msg = algArgs[3]
			qx = algArgs[4]
			qy = algArgs[5]
			r = algArgs[6]
			s = algArgs[7]
		default:
			return nil, 0, errors.Errorf("incorrect number of args for ECDSA operation: got %d, want 2, 4, 5, or 8", argLen)
		}
		o, err := newOp(op)
		if err != nil {
			return nil, 0, err
		}
		e, err := newECDSA(o, ecdsaCurve(curve), ecdsaSigParams(r, s, d, qx, qy), digest(msg), hashAlg(hash))
		if err != nil {
			return nil, 0, err
		}
		return []string{getECDSACommand(e)}, numRes, nil
	default:
		return nil, 0, errors.Errorf("unrecognized algorithm: %s", algType)
	}
}

// parseInBuf is a generic parser for ACVPTool input format.
// It returns an array of args passed.
func parseInBuf(b []byte) ([]string, error) {
	if len(b) < 4 {
		return nil, errors.Errorf("input buffer too short: %d bytes", len(b))
	}
	numArgs := binary.LittleEndian.Uint32(b[0:4])
	var res []string
	startInd := numArgs*4 + 4
	if uint32(len(b)) < startInd {
		return nil, errors.Errorf("input buffer too short: %d bytes, cannot extract arguments", len(b))
	}
	// first arg is already a string
	argLen := binary.LittleEndian.Uint32(b[4:8])
	endInd := startInd + argLen
	if uint32(len(b)) < endInd {
		return nil, errors.Errorf("input buffer too short: %d bytes, expected at least %d bytes",
			len(b), endInd)
	}
	res = append(res, string(b[startInd:endInd]))
	startInd = endInd
	for i := uint32(8); i < numArgs*4+4; i += 4 {
		endInd = startInd + binary.LittleEndian.Uint32(b[i:i+4])
		if uint32(len(b)) < endInd {
			return nil, errors.Errorf("input buffer too short: %d bytes, expected at least % bytes",
				len(b), endInd)
		}
		res = append(res, hex.EncodeToString(b[startInd:endInd]))
		startInd = endInd
	}
	return res, nil
}

// getHashCommand constructs a trunks Hash command
// Trunks command gets executed via hash_command_handler in cr50
func getHashCommand(hp *hashPrimitive) string {
	// 8001      TPM_ST_NO_SESSIONS
	// 00000000  Command/response size
	// 20000000  Cr50 Vendor Command (Constant, TPM Command Code)
	// 0000      Vendor Command Code (VENDOR_CC_ enum) 0001 for SHA1, SHA-256
	// Command structure, shared out of band with the test driver running
	// on the host:
	//
	// field     |    size  |                  note
	// ===================================================================
	// hash_cmd  |    1     | 0 - start, 1 - cont., 2 - finish, 3 - single
	//           |          | 4 - SW HMAC single shot (TPM code)
	//           |          | 5 - HW HMAC SHA256 single shot (dcrypto code)
	// hash_alg  |    1     | 0 - sha1, 1 - sha256, 2 - sha384, 3 - sha512
	// handle    |    1     | session handle, ignored in 'single' mode
	// text_len  |    2     | size of the text to process, big endian
	// text      | text_len | text to hash
	// for HMAC single shot only:
	// key_len   |    2     | size of the key for HMAC, big endian
	// key       | key_len  | key for HMAC single shot

	var cmdBody, cmdHeader bytes.Buffer
	cmdHeader.WriteString("8001")
	cmdBody.WriteString(hp.cmd)
	cmdBody.WriteString(hp.alg)
	cmdBody.WriteString("00")
	cmdBody.WriteString(fmt.Sprintf("%04x", len(hp.msg)/2))
	cmdBody.WriteString(hp.msg)
	if hp.cmd == tpmHMACCmdMode {
		cmdBody.WriteString(fmt.Sprintf("%04x", len(hp.key)/2))
		cmdBody.WriteString(hp.key)
	}
	cmdHeader.WriteString(fmt.Sprintf("%08x", cmdBody.Len()/2+cr50HeaderSize))
	cmdHeader.WriteString("200000000001")
	return cmdHeader.String() + cmdBody.String()
}

// getAESCommand constructs a trunks AES command
// Trunks command gets executed via aes_command_handler in cr50
func getAESCommand(bc *blockCipher) string {
	// Cipher modes being tested
	// CIPHER_MODES = {'ECB': '00', 'CTR': '01', 'CBC': '02',
	// 				'GCM': '03', 'OFB': '04', 'CFB': '05'}
	// 8001      TPM_ST_NO_SESSIONS
	// 00000000  Command/response size
	// 20000000  Cr50 Vendor Command (Constant, TPM Command Code)
	// 0000      Vendor Command Code (VENDOR_CC_ enum) 0000 for AES
	// Command body: test_mode|cipher_mode|
	// Command structure, shared out of band with the test driver running
	// on the host:

	//  field       |    size  |              note
	//  ================================================================
	//  mode        |    1     | 0 - decrypt, 1 - encrypt
	//  cipher_mode |    1     | as per aes_test_cipher_mode
	//  key_len     |    1     | key size in bytes (16, 24 or 32)
	//  key         | key len  | key to use
	//  iv_len      |    1     | either 0 or 16
	//  iv          | 0 or 16  | as defined by iv_len
	//  aad_len     |  <= 127  | additional authentication data length
	//  aad         |  aad_len | additional authentication data
	//  text_len    |    2     | size of the text to process, big endian
	//  text        | text_len | text to encrypt/decrypt

	var cmdBody, cmdHeader bytes.Buffer
	cmdHeader.WriteString("8001")
	cmdBody.WriteString(bc.encrypt)
	cmdBody.WriteString(bc.mode)
	cmdBody.WriteString(fmt.Sprintf("%02x", len(bc.key)/2))
	cmdBody.WriteString(bc.key)
	cmdBody.WriteString(fmt.Sprintf("%02x", len(bc.iv)/2))
	cmdBody.WriteString(bc.iv)
	// AAD is always empty for currently supported modes
	cmdBody.WriteString("00")
	cmdBody.WriteString(fmt.Sprintf("%04x", len(bc.in)/2))
	cmdBody.WriteString(bc.in)
	cmdHeader.WriteString(fmt.Sprintf("%08x", cmdBody.Len()/2+cr50HeaderSize))
	cmdHeader.WriteString("200000000000")

	return cmdHeader.String() + cmdBody.String()
}

// getDRBGGenCommand constructs a trunks command to
// generate output from an already initialized DRBG.
func getDRBGGenCommand(input string, outLen uint32) string {
	var cmdBody, cmdHeader bytes.Buffer
	cmdHeader.WriteString("8001")
	cmdBody.WriteString("02")
	cmdBody.WriteString(fmt.Sprintf("%04x", len(input)/2))
	cmdBody.WriteString(input)
	cmdBody.WriteString(fmt.Sprintf("%04x", outLen))
	cmdHeader.WriteString(fmt.Sprintf("%08x", cmdBody.Len()/2+cr50HeaderSize))
	cmdHeader.WriteString("200000000032")
	return cmdHeader.String() + cmdBody.String()
}

// getDRBGCommands constructs the sequence of trunks commands
// that need to be executed for a given NIST test case.
func getDRBGCommands(d *drbgSHA256) []string {
	// 8001      TPM_ST_NO_SESSIONS
	// 00000000  Command/response size
	// 20000000  Cr50 Vendor Command (Constant, TPM Command Code)
	// 0032      Vendor Command Code (VENDOR_CC_ enum) 0x32 for DRBG
	//
	// DRBG_TEST command structure:
	//
	// field       |    size  |              note
	// ==========================================================================
	// mode        |    1     | 0 - DRBG_INIT, 1 - DRBG_RESEED, 2 - DRBG_GENERATE
	// p0_len      |    2     | size of first input in bytes
	// p0          |  p0_len  | entropy for INIT & SEED, input for GENERATE
	// p1_len      |    2     | size of second input in bytes (for INIT & RESEED)
	//             |          | or size of expected output for GENERATE
	// p1          |  p1_len  | nonce for INIT & SEED
	// p2_len      |    2     | size of third input in bytes for DRBG_INIT
	// p2          |  p2_len  | personalization for INIT & SEED
	//
	// DRBG_INIT (entropy, nonce, perso)
	// DRBG_RESEED (entropy, additional input 1, additional input 2)
	// DRBG_INIT and DRBG_RESEED returns empty response
	// DRBG_GENERATE (p0_len, p0 - additional input 1, p1_len - size of output)
	// DRBG_GENERATE returns p1_len bytes of generated data
	// (up to a maximum of 128 bytes)

	result := make([]string, 3)
	var cmdBody, cmdHeader bytes.Buffer
	cmdHeader.WriteString("8001")
	cmdBody.WriteString("00")
	cmdBody.WriteString(fmt.Sprintf("%04x", len(d.entropy)/2))
	cmdBody.WriteString(d.entropy)
	cmdBody.WriteString(fmt.Sprintf("%04x", len(d.nonce)/2))
	cmdBody.WriteString(d.nonce)
	cmdBody.WriteString(fmt.Sprintf("%04x", len(d.perso)/2))
	cmdBody.WriteString(d.perso)
	cmdHeader.WriteString(fmt.Sprintf("%08x", cmdBody.Len()/2+cr50HeaderSize))
	cmdHeader.WriteString("200000000032")
	result[0] = cmdHeader.String() + cmdBody.String()
	result[1] = getDRBGGenCommand(d.input, d.outLen)
	result[2] = getDRBGGenCommand(d.input2, d.outLen)
	return result
}

// getECDSACommand constructs a trunks ECDSA command.
// Trunks command gets executed via ecc_command_handler in cr50.
func getECDSACommand(e *ecdsa) string {
	// Command format:

	// TEST_SIGN:
	// OP | CURVE_ID | SIGN_MODE | HASHING | DIGEST_LEN | DIGEST
	//    @returns R | S

	// TEST_SIGN_ANY:
	// OP | CURVE_ID | SIGN_MODE | HASHING | DIGEST_LEN | DIGEST |
	//   D_LEN | D
	//    @returns R | S

	// TEST_VERIFY:
	// OP | CURVE_ID | SIGN_MODE | HASHING | R_LEN | R | S_LEN | S
	//   DIGEST_LEN | DIGEST
	//    @returns 1 if successful

	// TEST_VERIFY_ANY:
	// OP | CURVE_ID | SIGN_MODE | HASHING | R_LEN | R | S_LEN | S |
	//   DIGEST_LEN | DIGEST | QX_LEN | QX | QY_LEN | QY
	//    @returns 1 if successful

	// TEST_KEYDERIVE:
	// OP | CURVE_ID | SEED_LEN | SEED
	//    @returns 1 if successful

	// TEST_POINT:
	// OP | CURVE_ID | QX_LEN | QX | QY_LEN | QY
	//    @returns 1 if point is on curve

	// TEST_KEYGEN:
	// OP | CURVE_ID
	//    @returns 1 if successful

	// FIELD          LENGTH
	// OP             1
	// CURVE_ID       1
	// SIGN_MODE      1
	// HASHING        1
	// MSG_LEN        2 (big endian)
	// MSG            MSG_LEN
	// SEED_LEN       2 (big endian)
	// SEED           SEED_LEN
	// R_LEN          2 (big endian)
	// R              R_LEN
	// S_LEN          2 (big endian)
	// S              S_LEN
	// DIGEST_LEN     2 (big endian)
	// DIGEST         DIGEST_LEN
	// D_LEN          2 (big endian)
	// D              D_LEN
	// QX_LEN         2 (big endian)
	// QX             QX_LEN
	// QY_LEN         2 (big endian)
	// QY             QX_LEN

	var cmdBody bytes.Buffer
	cmdBody.WriteString(e.op.hex())
	cmdBody.WriteString(e.curve)
	switch e.op {
	case tpmKeyVer:
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.qx)/2))
		cmdBody.WriteString(e.qx)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.qy)/2))
		cmdBody.WriteString(e.qy)
	case tpmSigGen:
		cmdBody.WriteString("18") // 0x18 indicates the ECDSA algo to the trunks command
		cmdBody.WriteString(e.hash)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.digest)/2))
		cmdBody.WriteString(e.digest)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.d)/2))
		cmdBody.WriteString(e.d)
	case tpmSigVer:
		cmdBody.WriteString("18") // 0x18 indicates the ECDSA algo to the trunks command
		cmdBody.WriteString(e.hash)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.r)/2))
		cmdBody.WriteString(e.r)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.s)/2))
		cmdBody.WriteString(e.s)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.digest)/2))
		cmdBody.WriteString(e.digest)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.qx)/2))
		cmdBody.WriteString(e.qx)
		cmdBody.WriteString(fmt.Sprintf("%04x", len(e.qy)/2))
		cmdBody.WriteString(e.qy)
	}

	var cmdHeader bytes.Buffer
	cmdHeader.WriteString("8001")
	cmdHeader.WriteString(fmt.Sprintf("%08x", cmdBody.Len()/2+cr50HeaderSize))
	cmdHeader.WriteString("200000000003") // 03 is the ECC command code

	return cmdHeader.String() + cmdBody.String()
}

// verifyResult verifies that the result groups returned by subprocess equal expected result groups.
func verifyResult(actual, expected []byte) (bool, error) {
	actual = bytes.ToLower(actual)
	expected = bytes.ToLower(expected)
	var actualGroups, expectedGroups interface{}
	if err := json.Unmarshal(actual, &actualGroups); err != nil {
		return false, err
	}
	if err := json.Unmarshal(expected, &expectedGroups); err != nil {
		return false, err
	}
	return reflect.DeepEqual(actualGroups, expectedGroups), nil
}

// extractTrunksResp validates the trunks response is a success and
// extracts the response bytes from raw trunks output
func extractTrunksResp(b []byte) ([]byte, error) {
	// Responses to TPM vendor commands have the following header structure:
	// 8001      TPM_ST_NO_SESSIONS
	// 00000000  Response size
	// 00000000  Response code
	// 0000      Vendor Command Code
	b, err := hex.DecodeString(string(b))
	if err != nil {
		return nil, errors.Wrap(err, "failed to decode string from byte slice")
	}

	if len(b) < cr50RespHeaderSize {
		return nil, errors.Errorf("trunks response too small: %d bytes", len(b))
	}

	respCode := binary.LittleEndian.Uint32(b[6:10])
	if respCode != 0 {
		return nil, errors.Errorf("unexpected response code from Cr50: %x", respCode)
	}

	return b[12:], nil
}

func extractACV(dec *json.Decoder, a *acv) error {
	// Parse the starting '['.
	arrayStart, err := dec.Token()
	if err != nil {
		return err
	}
	if delim, ok := arrayStart.(json.Delim); !ok || delim != '[' {
		return errors.Errorf("found %#v when expecting initial array in input file", arrayStart)
	}
	// Extract the ACV version.
	if err := dec.Decode(a); err != nil {
		return err
	}
	return nil
}

// ACVP takes a JSON generated by the ACVP server and runs the test cases in it.
func ACVP(ctx context.Context, s *testing.State) {
	d := s.Param().(data)
	vectorsBytes, err := ioutil.ReadFile(s.DataPath(d.inputFile))
	if err != nil {
		s.Fatal("Failed reading internal data file: ", err)
	} else {
		s.Log("Read data file: ", d.inputFile)
	}

	newFormat := false
	var v vectors
	err = json.Unmarshal(vectorsBytes, &v)
	var a acv
	// Check if parsing worked, otherwise try the lab format.
	// Input file may be obtained from lab, in which case it uses a slightly
	// different format.
	if err != nil || v.ID == 0 {
		// Assume new format.
		newFormat = true
		dec := json.NewDecoder(bytes.NewReader(vectorsBytes))
		if err := extractACV(dec, &a); err != nil {
			s.Fatal("Error parsing ACV version from input file: ", err)
		}
		// Extract the test vector.
		var tv interface{}
		if err := dec.Decode(&tv); err != nil {
			s.Fatal("Parse error while decoding vector: " + err.Error())
		}
		vectorsBytes, err = json.Marshal(tv)
		if err != nil {
			s.Fatal("Can't marshal interface to bytes: ", err)
		}
		if err := json.Unmarshal(vectorsBytes, &v); err != nil {
			s.Fatal("Failed to parse vector set: ", err)
		}
	}

	inout := cr50IO{
		ctx: ctx,
	}

	cmd := testexec.CommandContext(ctx, "trunks_send", "--raw")

	// TODO(b/141372763): cmd.Cmd should actually be something empty.
	middle := subprocess.NewWithIO(cmd.Cmd, &inout, &inout)
	defer middle.Close()

	replyGroups, err := middle.Process(v.Algo, vectorsBytes)
	if err != nil {
		s.Errorf("Failed to process middle: %s", err)
	}

	outFile := strings.Split(d.inputFile, ".")[0] + "-output.json"
	if newFormat {
		outResult := map[string]interface{}{
			"vsId":      v.ID,
			"algorithm": v.Algo,
			"isSample":  v.IsSample,
		}

		if v.Revision != "" {
			outResult["revision"] = v.Revision
		}
		if v.Mode != "" {
			outResult["mode"] = v.Mode
		}
		var r interface{}
		err := json.Unmarshal(replyGroups, &r)
		if err != nil {
			s.Fatal("Can't unmarshal replyGroups from acvpTool: ", err)
		}
		outResult["testGroups"] = r
		outResultBytes, err := json.MarshalIndent([]interface{}{a, outResult}, "", "    ")
		if err != nil {
			s.Fatal("Can't marshal out results interface to bytes")
		}
		if err := ioutil.WriteFile(filepath.Join(s.OutDir(), outFile),
			outResultBytes, 0644); err != nil {
			s.Error("Unable to write output file in new format: ", err)
		}
	} else {
		if err := ioutil.WriteFile(filepath.Join(s.OutDir(), outFile),
			replyGroups, 0644); err != nil {
			s.Error("Unable to write output file: ", err)
		}
	}

	// Verify results are correct if an expected NIST results file is provided.
	// Expected results file may not always be present, for example when the
	// result relies on random numbers generated by Cr50.
	if d.expectedFile != "" {
		expectedBytes, err := ioutil.ReadFile(s.DataPath(d.expectedFile))
		if err != nil {
			s.Fatal("Failed reading results data file: ", err)
		}
		if newFormat {
			expectedDec := json.NewDecoder(bytes.NewReader(expectedBytes))
			var b acv
			if err := extractACV(expectedDec, &b); err != nil {
				s.Fatal("Error parsing ACV version from expected output file: ", err)
			}
			// Extract the test vector
			var r map[string]interface{}
			if err := expectedDec.Decode(&r); err != nil {
				s.Fatal("parse error while decoding expected results: " + err.Error())
			}
			rGroups, ok := r["testGroups"]
			if !ok {
				log.Fatal("can't extract reply groups frome expected results.")
			}
			expectedBytes, err = json.Marshal(rGroups)
			if err != nil {
				log.Fatal("can't marshal expected reply groups interface to bytes")
			}
		}

		match, err := verifyResult(replyGroups, expectedBytes)
		if err != nil {
			s.Fatal("Failed to verify expected result matches processed result: ", err)
		}
		if !match {
			s.Fatal("Result returned by Cr50 does not match expected result")
		}
	}
}