package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"os"
	"reflect"
	"strings"

	"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{
			buildParameterDecodeCommand(),
			buildParameterEncodeCommand(),
			buildExternalSpecCommand(),
			buildExternalPublicKeyCommand(),
			buildExternalSignCommand(),
		},
		Action: func(ctx *cli.Context) error {
			fmt.Println("External signer PKCS#11, specification: https://openwebstandard.org/rfc1")
			return nil
		},
	}

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

func buildParameterEncodeCommand() *cli.Command {
	return &cli.Command{
		Name:  "parameter_encode",
		Usage: "Encode parameter, show PKCS#11 URIs: pkcs11-tool --module /usr/local/lib/libykcs11.dylib --login --list-objects --type privkey",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "uri",
				Usage:    "PKCS#11 URI",
				Required: true,
			},
			&cli.StringFlag{
				Name:  "library",
				Usage: "Library",
			},
			&cli.StringFlag{
				Name:  "pin",
				Usage: "PIN",
			},
		},
		Action: func(ctx *cli.Context) error {
			uri := ctx.String("uri")
			library := ctx.String("library")
			pin := ctx.String("pin")
			if !strings.HasPrefix(uri, "pkcs11:") {
				return fmt.Errorf("invalid uri: %s", uri)
			}
			if library == "" {
				library = "/usr/local/lib/libykcs11.dylib"
			}
			uriMap := parseUri(uri)
			pkcs11Key := Pkcs11Key{
				Library:    library,
				Pin:        pin,
				TokenLabel: uriMap["token"],
				KeyLabel:   uriMap["object"],
			}
			parameter, err := json.Marshal(&pkcs11Key)
			if err != nil {
				return nil
			}
			fmt.Println(base64.StdEncoding.EncodeToString(parameter))
			return nil
		},
	}
}

func parseUri(uri string) map[string]string {
	uriMap := map[string]string{}
	pkcs11UriLeftPart := strings.TrimPrefix(uri, "pkcs11:")
	keyValus := strings.Split(pkcs11UriLeftPart, ";")
	for _, kv := range keyValus {
		keyAndValue := strings.Split(kv, "=")
		if len(keyAndValue) == 2 {
			decodedValue, err := url.QueryUnescape(keyAndValue[1])
			if err == nil {
				uriMap[keyAndValue[0]] = decodedValue
			} else {
				uriMap[keyAndValue[0]] = keyAndValue[1]
			}
		}
	}
	return uriMap
}

func buildParameterDecodeCommand() *cli.Command {
	return &cli.Command{
		Name:  "parameter_decode",
		Usage: "Decode parameter",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "parameter",
				Usage:    "Parameter",
				Required: true,
			},
		},
		Action: func(ctx *cli.Context) error {
			parameter := ctx.String("parameter")
			pkcs11Key, err := parseParameter(parameter)
			if err != nil {
				return err
			}
			return printResponse(pkcs11Key)
		},
	}
}

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-pkcs11/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")
			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)
			}
			return tryWithPkcs11(parameter, true, func(privateKey crypto11.Signer) error {
				publicKey := privateKey.Public()

				if ecdsaPublicKey, ok := publicKey.(*ecdsa.PublicKey); ok {
					ecName := ecdsaPublicKey.Curve.Params().Name
					es256Matches := alg == "ES256" && ecName == "P-256"
					es384Matches := alg == "ES384" && ecName == "P-384"
					es512Matches := alg == "ES512" && ecName == "P-521"

					esMatches := es256Matches || es384Matches || es512Matches

					if !esMatches {
						return fmt.Errorf("algorithm mismatch %s vs %s", alg, ecName)
					}
				} else if _, ok := publicKey.(*rsa.PublicKey); ok {
					rsValid := alg == "RS256" || alg == "RS384" || alg == "RS512"
					if !rsValid {
						return fmt.Errorf("rsa algorithm invalid %s", alg)
					}
				} else {
					return fmt.Errorf("invalid public key type: %+v", reflect.TypeOf(publicKey))
				}

				var digest []byte
				var hash crypto.Hash
				if alg == "ES256" || alg == "RS256" {
					digest32 := sha256.Sum256(message)
					digest = digest32[:]
					hash = crypto.SHA256
				} else if alg == "ES384" || alg == "RS384" {
					sha384 := crypto.SHA384.New()
					digest = sha384.Sum(message)
					hash = crypto.SHA384
				} else if alg == "ES512" || alg == "RS512" {
					digest64 := sha512.Sum512(message)
					digest = digest64[:]
					hash = crypto.SHA512
				} else {
					return fmt.Errorf("invalid algorithm: %s", alg)
				}

				signature, err := privateKey.Sign(rand.Reader, digest, hash)
				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)
}