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 }