package main

import (
	"crypto"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"os"

	"github.com/ThalesGroup/crypto11"
	"github.com/urfave/cli/v2"
)

const (
	EnvPkcs11Pin = "PKCS11_PIN"
)

type Pkcs11Key struct {
	Library    string `json:"library"`
	TokenLabel string `json:"token_label"`
	Pin        string `json:"pin"`
	KeyLabel   string `json:"key_label"`
}

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-pkc11",
		Usage: "External signer PKCS#11",
		Commands: []*cli.Command{
			buildExternalSpecCommand(),
			buildExternalPublicKeyCommand(),
			buildExternalSignCommand(),
		},
		Action: func(ctx *cli.Context) error {
			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-pkc11/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")
			return tryWithPkcs11(parameter, false, func(privateKey crypto11.Signer) error {
				publicKey := privateKey.Public()
				publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
				if err != nil {
					return fmt.Errorf("marshal public key failed: %v", err)
				}
				externalPublicKeyResponse := &ExternalPublicKeyResponse{
					Success:         true,
					PublicKeyBase64: base64.StdEncoding.EncodeToString(publicKeyBytes),
				}

				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")
			messageBase64 := ctx.String("message-base64")
			// TODO check --alg
			message, err := base64.StdEncoding.DecodeString(messageBase64)
			if err != nil {
				return fmt.Errorf("unmarshal --message-base64 failed: %v", err)
			}
			return tryWithPkcs11(parameter, true, func(privateKey crypto11.Signer) error {
				hashed := sha256.Sum256(message)
				signature, err := privateKey.Sign(rand.Reader, hashed[:], crypto.SHA256)
				if err != nil {
					return fmt.Errorf("sign with private key failed: %v", err)
				}
				externalSignatureResponse := ExternalSignatureResponse{
					Success:         true,
					SignatureBase64: base64.StdEncoding.EncodeToString(signature),
				}

				printResponseSlient(externalSignatureResponse)
				return nil
			})
		},
	}
}

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

func getPin(pin string) (string, error) {
	if pin != "" {
		return pin, nil
	}
	envPin := os.Getenv(EnvPkcs11Pin)
	if envPin != "" {
		return envPin, nil
	}
	return "", errors.New("PIN is not set, set PIN: " + EnvPkcs11Pin)
}

func tryWithPkcs11(parameter string, requirePin bool, callback func(crypto11.Signer) error) error {
	pkcs11Key, err := parseParameter(parameter)
	if err != nil {
		return err
	}
	pin, err := getPin(pkcs11Key.Pin)
	if requirePin && err != nil {
		return err
	}
	config := &crypto11.Config{
		Path:       pkcs11Key.Library,
		TokenLabel: pkcs11Key.TokenLabel,
		Pin:        pin,
	}
	keyLabel := pkcs11Key.KeyLabel

	crypto11Ctx, err := crypto11.Configure(config)
	if err != nil {
		return fmt.Errorf("config crypto11 failed: %v", err)
	}
	defer crypto11Ctx.Close()

	privateKey, err := crypto11Ctx.FindKeyPair(nil, []byte(keyLabel))
	if err != nil {
		return fmt.Errorf("find private key failed: %v", err)
	}
	return callback(privateKey)
}