package main

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

	"github.com/ThalesGroup/crypto11"
)

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"`
}

func main() {
	parameter := "ewogICJsaWJyYXJ5IjogIi91c3IvbG9jYWwvbGliL2xpYnlrY3MxMS5keWxpYiIsCiAgInRva2VuX2xhYmVsIjogIll1YmlLZXkgUElWICM1MDEwMjIwIiwKICAicGluIjogIiIsCiAgImtleV9sYWJlbCI6ICJQcml2YXRlIGtleSBmb3IgUElWIEF1dGhlbnRpY2F0aW9uIgp9Cg=="
	message := "message"

	parameterBytes, err := base64.StdEncoding.DecodeString(parameter)
	if err != nil {
		println("1", err.Error())
		return
	}
	var pkcs11Key Pkcs11Key
	err = json.Unmarshal(parameterBytes, &pkcs11Key)
	if err != nil {
		println("2", err.Error())
		return
	}

	pin, err := getPin(pkcs11Key.Pin)
	if err != nil {
		println("3", err.Error())
		return
	}
	config := &crypto11.Config{
		Path:       pkcs11Key.Library,
		TokenLabel: pkcs11Key.TokenLabel,
		Pin:        pin,
	}
	keyLabel := pkcs11Key.KeyLabel

	ctx, err := crypto11.Configure(config)
	if err != nil {
		println("4", err.Error())
		return
	}
	defer ctx.Close()

	privateKey, err := ctx.FindKeyPair(nil, []byte(keyLabel))
	if err != nil {
		println("5", err.Error())
		return
	}
	publicKey := privateKey.Public()
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		println("6", err.Error())
		return
	}
	println(base64.StdEncoding.EncodeToString(publicKeyBytes))

	hashed := sha256.Sum256([]byte(message))
	signature, err := privateKey.Sign(rand.Reader, hashed[:], crypto.SHA256)
	if err != nil {
		println("7", err.Error())
		return
	}
	fmt.Printf("%x\n", signature)
}

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)
}