package main

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

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

type TpmKey struct {
	Tpm     string `json:"tpm"`
	Keyblob string `json:"keyblob"`
}

type ErrorResponse struct {
	Success bool   `json:"success"`
	Error   string `json:"error"`
}

type ExternalSpecResponse struct {
	Success       bool     `json:"success"`
	Agent         string   `json:"agent"`
	Specification string   `json:"specification"`
	Commands      []string `json:"commands"`
}

type ExternalPublicKeyResponse struct {
	Success         bool   `json:"success"`
	PublicKeyBase64 string `json:"public_key_base64"`
}

type ExternalSignatureResponse struct {
	Success         bool   `json:"success"`
	SignatureBase64 string `json:"signature_base64"`
}

func printResponse(response any) error {
	errorMessageJsonBytes, jsonErr := json.Marshal(response)
	if jsonErr != nil {
		return jsonErr
	} else {
		fmt.Println(string(errorMessageJsonBytes))
		return nil
	}
}

func printResponseSlient(response any) {
	if err := printResponse(response); err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
	}
}

func main() {
	err := innerMain()
	if err != nil {
		errorMessage := &ErrorResponse{
			Success: false,
			Error:   fmt.Sprintf("%v", err),
		}
		if printResponse(errorMessage) != nil {
			_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
		}
		os.Exit(1)
	}
}

func innerMain() error {
	app := cli.App{
		Name:  "external-signer-tpm",
		Usage: "External signer TPM",
		Commands: []*cli.Command{
			buildExternalSpecCommand(),
			buildExternalPublicKeyCommand(),
			buildExternalSignCommand(),
			buildGenerateKeyCommand(),
		},
		Action: func(ctx *cli.Context) error {
			fmt.Println("External signer TPM, specification: https://openwebstandard.org/rfc1")
			return nil
		},
	}

	if err := app.Run(os.Args); err != nil {
		return err
	}
	return nil
}

func buildExternalSpecCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_spec",
		Usage: "External specification",
		Flags: []cli.Flag{},
		Action: func(ctx *cli.Context) error {
			externalSpecResponse := &ExternalSpecResponse{
				Success:       true,
				Agent:         "external-signer-tpm/0.1.0",
				Specification: "External/1.0.0-alpha",
				Commands: []string{
					"external_public_key",
					"external_sign",
				},
			}
			printResponseSlient(externalSpecResponse)
			return nil
		},
	}
}

func buildExternalPublicKeyCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_public_key",
		Usage: "External public key",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "parameter",
				Usage:    "Parameter",
				Required: true,
			},
		},
		Action: func(ctx *cli.Context) error {
			parameter := ctx.String("parameter")

			tpmKey, err := parseParameter(parameter)
			if err != nil {
				return fmt.Errorf("cannot parse parameter: %s", err)
			}

			keyblob, err := base64.StdEncoding.DecodeString(tpmKey.Keyblob)
			if err != nil {
				return fmt.Errorf("cannot decode keyblob: %s", err)
			}
			pubKey, err := tpm.UnmarshalRSAPublicKey(keyblob)
			if err != nil {
				return fmt.Errorf("cannot get public key: %s", err)
			}

			pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
			if err != nil {
				return fmt.Errorf("cannot marshal public key: %s", err)
			}

			externalPublicKeyResponse := &ExternalPublicKeyResponse{
				Success:         true,
				PublicKeyBase64: base64.StdEncoding.EncodeToString(pubKeyBytes),
			}

			printResponseSlient(externalPublicKeyResponse)
			return nil
		},
	}
}

func buildExternalSignCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_sign",
		Usage: "External sign",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "parameter",
				Usage:    "Parameter",
				Required: true,
			},
			&cli.StringFlag{
				Name:     "alg",
				Usage:    "Algorithm",
				Required: true,
			},
			&cli.StringFlag{
				Name:     "message-base64",
				Usage:    "Message base64 encoded",
				Required: true,
			},
		},
		Action: func(ctx *cli.Context) error {
			parameter := ctx.String("parameter")
			alg := ctx.String("alg")
			messageBase64 := ctx.String("message-base64")
			message, err := base64.StdEncoding.DecodeString(messageBase64)
			if err != nil {
				return fmt.Errorf("unmarshal --message-base64 failed: %v", err)
			}

			println("ALG: ", alg)

			tpmKey, err := parseParameter(parameter)
			if err != nil {
				return fmt.Errorf("cannot parse parameter: %s", err)
			}

			rwc, err := tpm.OpenTPM(tpmKey.Tpm)
			if err != nil {
				return fmt.Errorf("cannot open TPM %s, error: %s", tpmKey.Tpm, err)
			}
			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 := base64.StdEncoding.DecodeString(tpmKey.Keyblob)
			if err != nil {
				return fmt.Errorf("cannot decode keyblob: %s", err)
			}
			keyHandle, err := tpm.LoadKey2(rwc, keyblob, srkAuth[:])
			if err != nil {
				return fmt.Errorf("cannot load keyblob: %s", err)
			}
			defer tpm.CloseKey(rwc, keyHandle)

			hash := crypto.SHA256.New()
			if _, err = hash.Write(message); err != nil {
				return fmt.Errorf("error building hash of data: %s", err)
			}
			hashed := hash.Sum(nil)
			signature, err := tpm.Sign(rwc, usageAuth[:], keyHandle, crypto.SHA256, hashed[:])
			if err != nil {
				return fmt.Errorf("cannot perform sign operation: %s", err)
			}

			externalSignatureResponse := ExternalSignatureResponse{
				Success:         true,
				SignatureBase64: base64.StdEncoding.EncodeToString(signature),
			}

			printResponseSlient(externalSignatureResponse)
			return nil
		},
	}
}

func parseParameter(parameter string) (*TpmKey, error) {
	parameterBytes, err := base64.StdEncoding.DecodeString(parameter)
	if err != nil {
		return nil, fmt.Errorf("parse --parameter base64 failed: %v", err)
	}
	var tpmKey TpmKey
	err = json.Unmarshal(parameterBytes, &tpmKey)
	if err != nil {
		return nil, fmt.Errorf("parse --parameter base64 failed: %v", err)
	}
	return &tpmKey, nil
}