From 7ac9d6c60458576f1f6b5bf5e8ae34290d2f50f8 Mon Sep 17 00:00:00 2001
From: Tyera Eulberg <tyera@solana.com>
Date: Mon, 17 Sep 2018 10:48:04 -0600
Subject: [PATCH] Create keygen helper function for use in Wallet CLI, print
 keypair statement

---
 src/bin/keygen.rs | 21 ++++-----------------
 src/bin/wallet.rs | 11 +++--------
 src/wallet.rs     | 37 +++++++++++++++++++++++++++++++++++--
 3 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/src/bin/keygen.rs b/src/bin/keygen.rs
index 41ff70193..10a49ca48 100644
--- a/src/bin/keygen.rs
+++ b/src/bin/keygen.rs
@@ -3,14 +3,11 @@ extern crate clap;
 extern crate dirs;
 extern crate ring;
 extern crate serde_json;
+extern crate solana;
 
 use clap::{App, Arg};
-use ring::rand::SystemRandom;
-use ring::signature::Ed25519KeyPair;
+use solana::wallet::gen_keypair_file;
 use std::error;
-use std::fs::{self, File};
-use std::io::Write;
-use std::path::Path;
 
 fn main() -> Result<(), Box<error::Error>> {
     let matches = App::new("solana-keygen")
@@ -24,10 +21,6 @@ fn main() -> Result<(), Box<error::Error>> {
                 .help("Path to generated file"),
         ).get_matches();
 
-    let rnd = SystemRandom::new();
-    let pkcs8_bytes = Ed25519KeyPair::generate_pkcs8(&rnd)?;
-    let serialized = serde_json::to_string(&pkcs8_bytes.to_vec())?;
-
     let mut path = dirs::home_dir().expect("home directory");
     let outfile = if matches.is_present("outfile") {
         matches.value_of("outfile").unwrap()
@@ -36,15 +29,9 @@ fn main() -> Result<(), Box<error::Error>> {
         path.to_str().unwrap()
     };
 
+    let serialized_keypair = gen_keypair_file(outfile.to_string())?;
     if outfile == "-" {
-        println!("{}", serialized);
-    } else {
-        if let Some(outdir) = Path::new(outfile).parent() {
-            fs::create_dir_all(outdir)?;
-        }
-        let mut f = File::create(outfile)?;
-        f.write_all(&serialized.into_bytes())?;
+        println!("{}", serialized_keypair);
     }
-
     Ok(())
 }
diff --git a/src/bin/wallet.rs b/src/bin/wallet.rs
index 4f8c65c7a..67e21e25d 100644
--- a/src/bin/wallet.rs
+++ b/src/bin/wallet.rs
@@ -10,10 +10,9 @@ use solana::drone::DRONE_PORT;
 use solana::logger;
 use solana::signature::{read_keypair, KeypairUtil};
 use solana::thin_client::poll_gossip_for_leader;
-use solana::wallet::{parse_command, process_command, WalletConfig, WalletError};
+use solana::wallet::{gen_keypair_file, parse_command, process_command, WalletConfig, WalletError};
 use std::error;
 use std::net::SocketAddr;
-use std::process::Command;
 
 pub fn parse_args(matches: &ArgMatches) -> Result<WalletConfig, Box<error::Error>> {
     let network = if let Some(addr) = matches.value_of("network") {
@@ -37,12 +36,8 @@ pub fn parse_args(matches: &ArgMatches) -> Result<WalletConfig, Box<error::Error
     } else {
         path.extend(&[".config", "solana", "id.json"]);
         if !path.exists() {
-            Command::new("cargo")
-                .arg("run")
-                .arg("--bin")
-                .arg("solana-keygen")
-                .status()
-                .expect("failed to execute process");
+            gen_keypair_file(path.to_str().unwrap().to_string())?;
+            println!("New keypair generated at: {:?}", path.to_str().unwrap());
         }
 
         path.to_str().unwrap()
diff --git a/src/wallet.rs b/src/wallet.rs
index ad70a0851..6b3f2f7be 100644
--- a/src/wallet.rs
+++ b/src/wallet.rs
@@ -4,13 +4,16 @@ use clap::ArgMatches;
 use crdt::NodeInfo;
 use drone::DroneRequest;
 use fullnode::Config;
+use ring::rand::SystemRandom;
+use ring::signature::Ed25519KeyPair;
 use serde_json;
 use signature::{Keypair, KeypairUtil, Pubkey, Signature};
-use std::fs::File;
+use std::fs::{self, File};
 use std::io::prelude::*;
 use std::io::{Error, ErrorKind, Write};
 use std::mem::size_of;
 use std::net::{Ipv4Addr, SocketAddr, TcpStream};
+use std::path::Path;
 use std::thread::sleep;
 use std::time::Duration;
 use std::{error, fmt, mem};
@@ -236,6 +239,21 @@ pub fn request_airdrop(
     Ok(signature)
 }
 
+pub fn gen_keypair_file(outfile: String) -> Result<String, Box<error::Error>> {
+    let rnd = SystemRandom::new();
+    let pkcs8_bytes = Ed25519KeyPair::generate_pkcs8(&rnd)?;
+    let serialized = serde_json::to_string(&pkcs8_bytes.to_vec())?;
+
+    if outfile != "-" {
+        if let Some(outdir) = Path::new(&outfile).parent() {
+            fs::create_dir_all(outdir)?;
+        }
+        let mut f = File::create(outfile)?;
+        f.write_all(&serialized.clone().into_bytes())?;
+    }
+    Ok(serialized)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -247,7 +265,7 @@ mod tests {
     use fullnode::Fullnode;
     use ledger::LedgerWriter;
     use mint::Mint;
-    use signature::{Keypair, KeypairUtil};
+    use signature::{read_keypair, read_pkcs8, Keypair, KeypairUtil};
     use std::net::UdpSocket;
     use std::sync::mpsc::channel;
 
@@ -441,4 +459,19 @@ mod tests {
         assert!(signature.is_ok());
         assert!(client.check_signature(&signature.unwrap()));
     }
+    #[test]
+    fn test_gen_keypair_file() {
+        let outfile = "test_gen_keypair_file.json";
+        let serialized_keypair = gen_keypair_file(outfile.to_string()).unwrap();
+        let keypair_vec: Vec<u8> = serde_json::from_str(&serialized_keypair).unwrap();
+        assert!(Path::new(outfile).exists());
+        assert_eq!(keypair_vec, read_pkcs8(&outfile).unwrap());
+        assert!(read_keypair(&outfile).is_ok());
+        assert_eq!(
+            read_keypair(&outfile).unwrap().pubkey().as_ref().len(),
+            mem::size_of::<Pubkey>()
+        );
+        fs::remove_file(outfile).unwrap();
+        assert!(!Path::new(outfile).exists());
+    }
 }