github.com/prysmaticlabs/prysm@v1.4.4/tools/keystores/main_test.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"flag"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"path/filepath"
    10  	"strings"
    11  	"testing"
    12  
    13  	"github.com/google/uuid"
    14  	"github.com/prysmaticlabs/prysm/shared/bls"
    15  	"github.com/prysmaticlabs/prysm/shared/params"
    16  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    17  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    18  	"github.com/prysmaticlabs/prysm/validator/keymanager"
    19  	"github.com/urfave/cli/v2"
    20  	keystorev4 "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4"
    21  )
    22  
    23  const password = "secretPassw0rd$1999"
    24  
    25  type cliConfig struct {
    26  	keystoresPath string
    27  	password      string
    28  	privateKey    string
    29  	outputPath    string
    30  }
    31  
    32  func setupCliContext(
    33  	tb testing.TB,
    34  	conf *cliConfig,
    35  ) *cli.Context {
    36  	app := cli.App{}
    37  	set := flag.NewFlagSet("test", 0)
    38  	set.String(keystoresFlag.Name, conf.keystoresPath, "")
    39  	set.String(passwordFlag.Name, conf.password, "")
    40  	set.String(privateKeyFlag.Name, conf.privateKey, "")
    41  	set.String(outputPathFlag.Name, conf.outputPath, "")
    42  	assert.NoError(tb, set.Set(keystoresFlag.Name, conf.keystoresPath))
    43  	assert.NoError(tb, set.Set(passwordFlag.Name, conf.password))
    44  	assert.NoError(tb, set.Set(privateKeyFlag.Name, conf.privateKey))
    45  	assert.NoError(tb, set.Set(outputPathFlag.Name, conf.outputPath))
    46  	return cli.NewContext(&app, set, nil)
    47  }
    48  
    49  func createRandomKeystore(t testing.TB, password string) (*keymanager.Keystore, bls.SecretKey) {
    50  	encryptor := keystorev4.New()
    51  	id, err := uuid.NewRandom()
    52  	require.NoError(t, err)
    53  	validatingKey, err := bls.RandKey()
    54  	require.NoError(t, err)
    55  	pubKey := validatingKey.PublicKey().Marshal()
    56  	cryptoFields, err := encryptor.Encrypt(validatingKey.Marshal(), password)
    57  	require.NoError(t, err)
    58  	return &keymanager.Keystore{
    59  		Crypto:  cryptoFields,
    60  		Pubkey:  fmt.Sprintf("%x", pubKey),
    61  		ID:      id.String(),
    62  		Version: encryptor.Version(),
    63  		Name:    encryptor.Name(),
    64  	}, validatingKey
    65  }
    66  
    67  func setupRandomDir(t testing.TB) string {
    68  	randDir := t.TempDir()
    69  	require.NoError(t, os.MkdirAll(randDir, os.ModePerm))
    70  	return randDir
    71  }
    72  
    73  func TestDecrypt(t *testing.T) {
    74  	keystoresDir := setupRandomDir(t)
    75  	keystore, privKey := createRandomKeystore(t, password)
    76  	// We write a random keystore to a keystores directory.
    77  	encodedKeystore, err := json.MarshalIndent(keystore, "", "\t")
    78  	require.NoError(t, err)
    79  	keystoreFilePath := filepath.Join(keystoresDir, "keystore.json")
    80  	require.NoError(t, ioutil.WriteFile(
    81  		keystoreFilePath, encodedKeystore, params.BeaconIoConfig().ReadWritePermissions),
    82  	)
    83  
    84  	cliCtx := setupCliContext(t, &cliConfig{
    85  		keystoresPath: keystoreFilePath,
    86  		password:      password,
    87  	})
    88  
    89  	rescueStdout := os.Stdout
    90  	r, w, err := os.Pipe()
    91  	require.NoError(t, err)
    92  	os.Stdout = w
    93  
    94  	// We attempt to decrypt the keystore file we just wrote to disk.
    95  	require.NoError(t, decrypt(cliCtx))
    96  
    97  	require.NoError(t, w.Close())
    98  	out, err := ioutil.ReadAll(r)
    99  	require.NoError(t, err)
   100  
   101  	// We capture output from stdout.
   102  	os.Stdout = rescueStdout
   103  	stringOutput := string(out)
   104  
   105  	// We capture the results of stdout to check the public key and private keys
   106  	// were both printed to stdout.
   107  	assert.Equal(t, true, strings.Contains(stringOutput, keystore.Pubkey))
   108  	assert.Equal(t, true, strings.Contains(stringOutput, fmt.Sprintf("%#x", privKey.Marshal())))
   109  }
   110  
   111  func TestEncrypt(t *testing.T) {
   112  	keystoresDir := setupRandomDir(t)
   113  	keystoreFilePath := filepath.Join(keystoresDir, "keystore.json")
   114  	privKey, err := bls.RandKey()
   115  	require.NoError(t, err)
   116  
   117  	cliCtx := setupCliContext(t, &cliConfig{
   118  		outputPath: keystoreFilePath,
   119  		password:   password,
   120  		privateKey: fmt.Sprintf("%#x", privKey.Marshal()),
   121  	})
   122  
   123  	rescueStdout := os.Stdout
   124  	r, w, err := os.Pipe()
   125  	require.NoError(t, err)
   126  	os.Stdout = w
   127  
   128  	// We attempt to encrypt the secret key and save it to the output path.
   129  	require.NoError(t, encrypt(cliCtx))
   130  
   131  	require.NoError(t, w.Close())
   132  	out, err := ioutil.ReadAll(r)
   133  	require.NoError(t, err)
   134  
   135  	// We capture output from stdout.
   136  	os.Stdout = rescueStdout
   137  	stringOutput := string(out)
   138  
   139  	// We capture the results of stdout to check the public key was printed to stdout.
   140  	res := strings.Contains(stringOutput, fmt.Sprintf("%x", privKey.PublicKey().Marshal()))
   141  	assert.Equal(t, true, res)
   142  }