More wallet fun

Cargo.toml

@@ -74,3 +74,4 @@ tokio = "0.1"
 tokio-codec = "0.1"
 tokio-io = "0.1"
 itertools = "0.7.8"
+bs58 = "0.2.0"

src/bin/wallet.rs

@@ -1,16 +1,17 @@
 extern crate atty;
 extern crate bincode;
+extern crate bs58;
 extern crate env_logger;
 extern crate getopts;
 extern crate serde_json;
 extern crate solana;
 
 use bincode::serialize;
-use getopts::Options;
+use getopts::{Matches, Options};
 use solana::crdt::{get_ip_addr, ReplicatedData};
 use solana::drone::DroneRequest;
 use solana::mint::Mint;
-use solana::signature::Signature;
+use solana::signature::{PublicKey, Signature};
 use solana::thin_client::ThinClient;
 use std::env;
 use std::error;
@@ -24,10 +25,11 @@ use std::thread::sleep;
 use std::time::Duration;
 
 enum WalletCommand {
+    Address,
     Balance,
     AirDrop,
-    Pay,
-    Confirm,
+    Pay(i64, PublicKey),
+    Confirm(Signature),
 }
 
 #[derive(Debug, Clone)]
@@ -58,7 +60,6 @@ struct WalletConfig {
     client_addr: SocketAddr,
     drone_addr: SocketAddr,
     command: WalletCommand,
-    sig: Option<Signature>,
 }
 
 impl Default for WalletConfig {
@@ -70,26 +71,49 @@ impl Default for WalletConfig {
             client_addr: default_addr.clone(),
             drone_addr: default_addr.clone(),
             command: WalletCommand::Balance,
-            sig: None,
         }
     }
 }
 
 fn print_usage(program: &str, opts: Options) {
     let mut brief = format!("Usage: {} [options]\n\n", program);
-    brief += "  solana-wallet allows you to perform basic actions, including\n";
+    brief += "  solana-wallet allows you to perform basic actions, including";
     brief += "  requesting an airdrop, checking your balance, and spending tokens.";
     brief += "  Takes json formatted mint file to stdin.";
 
     print!("{}", opts.usage(&brief));
+    display_actions();
 }
 
-fn parse_command(input: &str) -> Result<WalletCommand, WalletError> {
-    match input {
+fn parse_command(matches: &Matches) -> Result<WalletCommand, WalletError> {
+    let input = &matches.free[0];
+    match input.as_ref() {
+        "address" => Ok(WalletCommand::Address),
         "balance" => Ok(WalletCommand::Balance),
         "airdrop" => Ok(WalletCommand::AirDrop),
-        "pay" => Ok(WalletCommand::Pay),
-        "confirm" => Ok(WalletCommand::Confirm),
+        "pay" => {
+            if matches.free.len() < 3 {
+                eprintln!("No tokens and public key provided");
+                exit(1);
+            }
+            let tokens = matches.free[1].parse().expect("parse integer");
+            let pubkey_vec = bs58::decode(&matches.free[2])
+                .into_vec()
+                .expect("base58-encoded public key");
+            let to = PublicKey::clone_from_slice(&pubkey_vec);
+            Ok(WalletCommand::Pay(tokens, to))
+        }
+        "confirm" => {
+            if matches.free.len() < 2 {
+                eprintln!("No signature provided");
+                exit(1);
+            }
+            let sig_vec = bs58::decode(&matches.free[1])
+                .into_vec()
+                .expect("base58-encoded signature");
+            let sig = Signature::clone_from_slice(&sig_vec);
+            Ok(WalletCommand::Confirm(sig))
+        }
         _ => Err(WalletError::CommandNotRecognized(input.to_string())),
     }
 }
@@ -110,10 +134,8 @@ fn parse_args(args: Vec<String>) -> Result<WalletConfig, Box<error::Error>> {
         }
     };
 
-    if matches.opt_present("h") {
-        let program = args[0].clone();
-        print_usage(&program, opts);
-        display_actions();
+    if matches.opt_present("h") || matches.free.len() < 1 {
+        print_usage(&args[0], opts);
         return Ok(WalletConfig::default());
     }
 
@@ -144,15 +166,13 @@ fn parse_args(args: Vec<String>) -> Result<WalletConfig, Box<error::Error>> {
     let mut drone_addr = leader.transactions_addr.clone();
     drone_addr.set_port(9900);
 
-    let command = parse_command(&matches.free[0])?;
-
+    let command = parse_command(&matches)?;
     Ok(WalletConfig {
         leader,
         id,
         client_addr,
         drone_addr, // TODO: Add an option for this.
         command,
-        sig: None, // TODO: Add an option for this.
     })
 }
 
@@ -162,6 +182,9 @@ fn process_command(
 ) -> Result<(), Box<error::Error>> {
     match config.command {
         // Check client balance
+        WalletCommand::Address => {
+            println!("{}", bs58::encode(config.id.pubkey()).into_string());
+        }
         WalletCommand::Balance => {
             println!("Balance requested...");
             let balance = client.poll_get_balance(&config.id.pubkey());
@@ -190,52 +213,31 @@ fn process_command(
             );
         }
         // If client has positive balance, spend tokens in {balance} number of transactions
-        WalletCommand::Pay => {
+        WalletCommand::Pay(tokens, to) => {
             let last_id = client.get_last_id();
-            let balance = client.poll_get_balance(&config.id.pubkey());
-            match balance {
-                Ok(0) => {
-                    println!("You don't have any tokens!");
-                }
-                Ok(balance) => {
-                    println!("Sending {:?} tokens to self...", balance);
-                    let sig = client
-                        .transfer(balance, &config.id.keypair(), config.id.pubkey(), &last_id)
-                        .expect("transfer return signature");
-                    println!("Transaction sent!");
-                    println!("Signature: {:?}", sig);
-                }
-                Err(ref e) if e.kind() == std::io::ErrorKind::Other => {
-                    println!("No account found! Request an airdrop to get started.");
-                }
-                Err(error) => {
-                    println!("An error occurred: {:?}", error);
-                }
-            }
+            let sig = client.transfer(tokens, &config.id.keypair(), to, &last_id)?;
+            println!("{}", bs58::encode(sig).into_string());
         }
         // Confirm the last client transaction by signature
-        WalletCommand::Confirm => match config.sig {
-            Some(sig) => {
-                if client.check_signature(&sig) {
-                    println!("Signature found at bank id {:?}", config.id);
-                } else {
-                    println!("Uh oh... Signature not found!");
-                }
+        WalletCommand::Confirm(sig) => {
+            if client.check_signature(&sig) {
+                println!("Confirmed");
+            } else {
+                println!("Not found");
             }
-            None => {
-                println!("No recent signature. Make a payment to get started.");
-            }
-        },
+        }
     }
     Ok(())
 }
 
 fn display_actions() {
     println!("");
-    println!("  `balance` - Get your account balance");
-    println!("  `airdrop` - Request a batch of tokens");
-    println!("  `pay` - Spend your tokens as fast as possible");
-    println!("  `confirm` - Confirm your last payment by signature");
+    println!("Commands:");
+    println!("  address   Get your public key");
+    println!("  balance   Get your account balance");
+    println!("  airdrop   Request a batch of tokens");
+    println!("  pay       Spend your tokens as fast as possible");
+    println!("  confirm   Confirm your last payment by signature");
     println!("");
 }