package main

import (
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/google/go-tpm/tpm"
	"github.com/urfave/cli/v2"
)

func buildGenerateKeyCommand() *cli.Command {
	return &cli.Command{
		Name:  "generate_key",
		Usage: "Generate key",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:        "parameter",
				Usage:       "The path to the TPM device to use",
				DefaultText: "/dev/tpm0",
			},
			&cli.StringFlag{
				Name:        "keyblob",
				Usage:       "Output path of the generated keyblob",
				DefaultText: "keyblob",
			},
			&cli.StringFlag{
				Name:        "public-key",
				Usage:       "Output path of the generated keyblob's public key",
				DefaultText: "publickey",
			},
			&cli.StringFlag{
				Name:  "pcrs",
				Usage: "A comma-separated list of PCR numbers against which the generated key will be bound. If blank, it will not be bound to any PCR values.",
			},
		},
		Action: func(ctx *cli.Context) error {
			tpmname := ctx.String("tpmname")
			keyblob := ctx.String("keyblob")
			publicKey := ctx.String("public-key")
			pcrs := ctx.String("pcrs")

			generateAction(tpmname, keyblob, publicKey, pcrs)
			return nil
		},
	}
}

func generateAction(tpmname, keyblobPath, pubKeyPath, pcrsStr string) {

	var pcrs []int
	if pcrsStr != "" {
		for _, pcr := range strings.Split(pcrsStr, ",") {
			pcrNum, err := strconv.Atoi(pcr)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Bad value in pcrs argument: %s\n", pcr)
				return
			}
			pcrs = append(pcrs, pcrNum)
		}
	}

	rwc, err := tpm.OpenTPM(tpmname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't open the TPM file %s: %s\n", tpmname, err)
		return
	}
	defer rwc.Close()

	// Compute the auth values as needed.
	var srkAuth [20]byte
	srkInput := os.Getenv(srkAuthEnvVar)
	if srkInput != "" {
		sa := sha1.Sum([]byte(srkInput))
		copy(srkAuth[:], sa[:])
	}

	var usageAuth [20]byte
	usageInput := os.Getenv(usageAuthEnvVar)
	if usageInput != "" {
		ua := sha1.Sum([]byte(usageInput))
		copy(usageAuth[:], ua[:])
	}

	var migrationAuth [20]byte
	migrationInput := os.Getenv(migrationAuthEnvVar)
	if migrationInput != "" {
		ma := sha1.Sum([]byte(migrationInput))
		copy(migrationAuth[:], ma[:])
	}

	keyblob, err := tpm.CreateWrapKey(rwc, srkAuth[:], usageAuth, migrationAuth, pcrs)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't make a new signing key: %s\n", err)
		return
	}
	fmt.Printf("Writing keyblob to %s\n", keyblobPath)
	if err = os.WriteFile(keyblobPath, keyblob, 0644); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing keyblob file: %s\n", err)
		return
	}

	pubKey, err := tpm.UnmarshalRSAPublicKey(keyblob)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not get public key: %s\n", err)
		return
	}

	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not marshal public key: %s\n", err)
		return
	}
	fmt.Printf("Writing public key to %s\n", pubKeyPath)
	if err = os.WriteFile(pubKeyPath, pubKeyBytes, 0644); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing public key file: %s\n", err)
		return
	}

	tpmKey := &TpmKey{
		Tpm:     tpmname,
		Keyblob: base64.StdEncoding.EncodeToString(keyblob),
	}

	tpmKeyJson, err := json.Marshal(&tpmKey)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error marshal tmp key: %s\n", err)
		return
	}
	println(base64.StdEncoding.EncodeToString(tpmKeyJson))
}