players := []int{2, 3, 5}
threshold := 1
sharing, err := tsmutils.RSASecretShare(threshold, players, privateKey)

wrappingKey, err := client.WrappingKey().WrappingKey(ctx)
publicKey, err := x509.ParsePKIXPublicKey(wrappingKey)
publicRSAKey, ok := publicKey.(*rsa.PublicKey)
wrappedKeyShare, err = tsmutils.EnvelopeWrap(publicRSAKey, sharing[playerID])

keyID, err := client.RSA().ImportKeyShares(ctx, sessionConfig, wrappedKeyShare, desiredKeyID)

publicKey, err := client.RSA().PublicKey(ctx, keyID)

wrappingKey, err := targetClient.WrappingKey().WrappingKey(ctx)
publicKey, err := x509.ParsePKIXPublicKey(wrappingKey)
publicWrappingKey, ok := publicKey.(*rsa.PublicKey)

privateWrappingKey, err := rsa.GenerateKey(rand.Reader, 4096)
publicWrappingKey, err := x509.MarshalPKIXPublicKey(&privateWrappingKey.PublicKey)

res, err := client.RSA().ExportKeyShares(ctx, exportSessionConfig, keyID, publicWrappingKey)

// If you do this for all key shares
keyShare, err := tsmutils.EnvelopeUnwrap(privateWrappingKey, res.WrappedKeyShare)

// Then you can recover the RSA key as follows
exportedPrivateRSAKey, err := tsmutils.RSARecombine([][]byte{ keyShare1, keyShare2, keyShare3})

message := []byte("some message to be signed")
sha256Hashed := sha256.Sum256(message)
hashedMessage := sha256Hashed[:]
hashFunction := tsm.HashFunctionSHA256
partialSignature, err := client.RSA().SignPSS(ctx, signSessionConfig, keyID, hashFunction, hashedMessage)

signature, err := tsm.RSAFinalizeSignaturePSS(hashFunction, hashedMessage, partialSignatures)

// Call this concurrently on each SDK
partialSignature, err := client.RSA().SignPKCS1v15(ctx, keyID, hashFunction, hashedMessage)

// Then call this somewhere, after collecting all partial signatures returned from the SDKs 
signature, err := tsm.RSAFinalizeSignaturePKCS1v15(hashFunction, hashedMessage, partialSignatures)

partialDecryption, err := client.RSA().Decrypt(ctx, keyID, ciphertext)

plaintext, err := tsm.RSAFinalizeDecryptionRaw(partialDecryptions)
plaintext, err := tsm.RSAFinalizeDecryptionPKCS1v15(partialDecryptions)
plaintext, err := tsm.RSAFinalizeDecryptionOAEP(hashFunction, label, partialDecryptions)

package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"fmt"
	"gitlab.com/Blockdaemon/go-tsm-sdkv2/v65/tsm"
	"gitlab.com/Blockdaemon/go-tsm-sdkv2/v65/tsm/tsmutils"
	"golang.org/x/sync/errgroup"
	"sync"
)

func main() {

	// Generate an RSA key to be imported into Builder Vault

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}

	// Split the private RSA key into a secret sharing

	players := []int{0, 1, 2} // The key should be secret shared among the MPC three nodes with these player IDs
	threshold := 1            // The security threshold that we want for the imported key

	sharing, err := tsmutils.RSASecretShare(threshold, players, privateKey)
	if err != nil {
		panic(err)
	}

	// Create a client for each of the nodes

	configs := []*tsm.Configuration{
		tsm.Configuration{URL: "http://localhost:8500"}.WithAPIKeyAuthentication("apikey0"),
		tsm.Configuration{URL: "http://localhost:8501"}.WithAPIKeyAuthentication("apikey1"),
		tsm.Configuration{URL: "http://localhost:8502"}.WithAPIKeyAuthentication("apikey2"),
	}

	clients := make([]*tsm.Client, len(configs))
	for i, config := range configs {
		var err error
		if clients[i], err = tsm.NewClient(config); err != nil {
			panic(err)
		}
	}

	// Wrap each key share with the wrapping key of the MPC node that should receive the share

	ctx := context.Background()
	wrappedKeyShares := make([][]byte, len(players))
	for i, client := range clients {

		// Get the wrapping key from the MPC node

		wrappingKey, err := client.WrappingKey().WrappingKey(ctx)
		if err != nil {
			panic(err)
		}
		publicKey, err := x509.ParsePKIXPublicKey(wrappingKey)
		if err != nil {
			panic(err)
		}
		publicRSAKey, ok := publicKey.(*rsa.PublicKey)
		if !ok {
			panic("error parsing wrapping key: not an rsa key")
		}

		// Wrap the key share using the wrapping key

		keyShare := sharing[players[i]]
		wrappedKeyShares[i], err = tsmutils.EnvelopeWrap(publicRSAKey, keyShare)
		if err != nil {
			panic(err)
		}

	}

	// Import the wrapped RSA key shares into the Builder Vault

	importSessionConfig := tsm.NewSessionConfig(tsm.GenerateSessionID(), players, nil)
	keyIDs := map[int]string{}
	var eg errgroup.Group
	var lock sync.Mutex
	for playerID, client := range clients {
		playerID, client := playerID, client
		eg.Go(func() error {
			res, err := client.RSA().ImportKeyShares(ctx, importSessionConfig, wrappedKeyShares[playerID], "")
			if err != nil {
				return err
			}
			lock.Lock()
			keyIDs[playerID] = res
			lock.Unlock()
			return nil
		})
	}

	// Validate key ID; make sure none of the nodes aborted during import, before you use the key.

	if err := eg.Wait(); err != nil {
		panic(err)
	}
	for i := 1; i < len(keyIDs); i++ {
		if keyIDs[0] != keyIDs[i] {
			panic("key IDs do not match")
		}
	}
	keyID := keyIDs[0]
	fmt.Println("Key ID of imported key:", keyID)

	// CAVEAT: The public RSA key can be obtained from any one MPC node as follows. It is important to only use the
	// public key if you trust that MPC node. Generally, you need to obtain the public key from at least threshold+1
	// MPC nodes and make sure the returned public keys are equal, before you start using the key. Obtaining the
	// public key from threshold + 1 nodes ensures that at least one honest MPC node is involved.

	publicKey, err := clients[0].RSA().PublicKey(ctx, keyID)
	if err != nil {
		panic(err)
	}

	fmt.Println("Public key (ASN.1 DER encoding of SubjectPublicKeyInfo):", hex.EncodeToString(publicKey))

	// Now we can use the imported key, for example to generate an RSA-PSS signature on a message

	message := []byte("some message")
	sha256Hashed := sha256.Sum256(message)
	hashedMessage := sha256Hashed[:]
	hashFunction := tsm.HashFunctionSHA256

	fmt.Println("Hashed message          :", hex.EncodeToString(hashedMessage))

	var partialSignatures [][]byte = make([][]byte, 0)
	signSessionConfig := tsm.NewSessionConfig(tsm.GenerateSessionID(), players, nil)
	for _, client := range clients {
		client := client
		eg.Go(func() error {
			partialSignature, err := client.RSA().SignPSS(ctx, signSessionConfig, keyID, hashFunction, hashedMessage)
			if err != nil {
				return err
			}
			lock.Lock()
			partialSignatures = append(partialSignatures, partialSignature)
			lock.Unlock()
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		panic(err)
	}

	// Combine partial signatures into the final RSASSA-PSS signature.

	signature, err := tsm.RSAFinalizeSignaturePSS(hashFunction, hashedMessage, partialSignatures)

	fmt.Println("RSA-PSS signature       :", hex.EncodeToString(signature))

	pubKey := privateKey.Public().(*rsa.PublicKey)
	err = rsa.VerifyPSS(pubKey, crypto.SHA256, hashedMessage, signature, nil)
	if err != nil {
		panic("invalid signature")
	}

	// OAEP e a message using the public RSA key

	label := []byte("hello label")
	ciphertext, err := rsa.EncryptOAEP(crypto.SHA256.New(), rand.Reader, pubKey, hashedMessage, label)
	if err != nil {
		panic(err)
	}

	// Use Builder Vault to decrypt

	var partialDecryptions [][]byte
	for _, client := range clients {
		client := client
		eg.Go(func() error {
			partialDecryption, err := client.RSA().Decrypt(ctx, keyID, ciphertext)
			if err != nil {
				return err
			}
			lock.Lock()
			partialDecryptions = append(partialDecryptions, partialDecryption)
			lock.Unlock()
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		panic(err)
	}

	// Recombine the partial decrypt results

	decryptedHashedMessage, err := tsm.RSAFinalizeDecryptionOAEP(hashFunction, label, partialDecryptions)
	if err != nil {
		panic(err)
	}

	fmt.Println("Decrypted hashed message:", hex.EncodeToString(decryptedHashedMessage))

	if !bytes.Equal(hashedMessage, decryptedHashedMessage) {
		panic("decrypted message does not match original message")
	}

	// We can export the RSA key from the Builder Vault as follows

	// Key shares are wrapped before being returned to the SDKs.
	// In a real application, you would probably use different wrapping keys for each MPC node. For example, you would
	// use the wrapping key obtained from another instance of the MPC node with the same player ID.
	// For the sake of simplicity, we use the same wrapping key for all MPC nodes in this example.
	// In any case, the wrapping key must be whitelisted and key export must be enabled in the MPC node configuration,
	// before the export can take place.

	privateWrappingKey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		panic(err)
	}
	publicWrappingKey, err := x509.MarshalPKIXPublicKey(&privateWrappingKey.PublicKey)
	if err != nil {
		panic(err)
	}

	exportSessionConfig := tsm.NewSessionConfig(tsm.GenerateSessionID(), players, nil)
	exportedWrappedKeyShares := make(map[int]*tsm.RSAWrappedKeyShare, len(clients))
	for playerID, client := range clients {
		playerID, client := playerID, client
		eg.Go(func() error {
			res, err := client.RSA().ExportKeyShares(ctx, exportSessionConfig, keyID, publicWrappingKey)
			if err != nil {
				return err
			}
			lock.Lock()
			exportedWrappedKeyShares[playerID] = res
			lock.Unlock()
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		panic(err)
	}

	for i := 1; i < len(keyIDs); i++ {
		if !bytes.Equal(exportedWrappedKeyShares[0].PKIXPublicKey, exportedWrappedKeyShares[i].PKIXPublicKey) {
			panic("checksums do not match")
		}
	}
	pkixPubKey := exportedWrappedKeyShares[0].PKIXPublicKey
	fmt.Println("Exported key checksum: ", hex.EncodeToString(pkixPubKey))
	if !bytes.Equal(publicKey, pkixPubKey) {
		panic("public keys do not match")
	}

	// Given all the key shares and the private wrapping key, you can recover the AES key from the shares as follows.
	// But for maximum security, you can also just import the wrapped key shares and checksum into another
	// Builder Vault instance. This will migrate the key between the Builder Vault instances, without any single point
	// of trust.

	// Unwrap the wrapped key shares

	exportedKeyShares := make(map[int][]byte, 0)
	for playerID, wrappedKeyShare := range exportedWrappedKeyShares {
		keyShare, err := tsmutils.EnvelopeUnwrap(privateWrappingKey, wrappedKeyShare.WrappedKeyShare)
		if err != nil {
			panic(err)
		}
		exportedKeyShares[playerID] = keyShare
	}

	// Then recover the HMAC key from the key shares

	exportedPrivateKey, err := tsmutils.RSARecombine(exportedKeyShares)
	if err != nil {
		panic(err)
	}

	// Sanity check: Exported private key should equal original private key

	if exportedPrivateKey.N.Cmp(privateKey.N) != 0 {
		panic("rsa key mismatch")
	}
	if exportedPrivateKey.D.Cmp(privateKey.D) != 0 {
		panic("rsa key mismatch")
	}

}