[wallet CLI] Supporting system variables in interactive shell (Mysten…

…Labs#621)

* add env variable support

* substitute all occurrence of tokens

sui/src/lib.rs

@@ -2,6 +2,8 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+extern crate core;
+
 pub mod config;
 pub mod keystore;
 pub mod shell;

sui/src/shell.rs

@@ -11,17 +11,16 @@ use rustyline::validate::Validator;
 use rustyline::{Config, Context, Editor};
 use rustyline_derive::Helper;
 use std::fmt::Display;
-use std::io;
 use std::io::Write;
-use structopt::clap::App;
+use std::{env, io};
+use structopt::clap::{App, SubCommand};
 use unescape::unescape;
 
 /// A interactive command line shell with history and completion support
 pub struct Shell<P: Display, S, H> {
     pub prompt: P,
     pub state: S,
     pub handler: H,
-    pub description: String,
     pub command: CommandStructure,
 }
 
@@ -42,17 +41,11 @@ impl<P: Display, S: Send, H: AsyncHandler<S>> Shell<P, S, H> {
             children: vec![],
         };
         command.children.push(help);
-        command.completions.extend([
-            "help".to_string(),
-            "exit".to_string(),
-            "quit".to_string(),
-            "clear".to_string(),
-        ]);
+        command.completions.extend(["help".to_string()]);
 
         rl.set_helper(Some(ShellHelper { command }));
 
         let mut stdout = io::stdout();
-
         'shell: loop {
             print!("{}", self.prompt);
             stdout.flush()?;
@@ -65,28 +58,42 @@ impl<P: Display, S: Send, H: AsyncHandler<S>> Shell<P, S, H> {
                 Err(err) => return Err(err.into()),
             };
 
+            let line = Self::substitution_env_variables(line);
+
             // Runs the line
             match Self::split_and_unescape(line.trim()) {
                 Ok(line) => {
-                    // do nothing if line is empty
-                    if line.is_empty() {
-                        continue 'shell;
-                    };
-                    // safe to unwrap with the above is_empty check.
-                    if *line.first().unwrap() == "quit" || *line.first().unwrap() == "exit" {
-                        println!("Bye!");
-                        break 'shell;
-                    };
-                    if *line.first().unwrap() == "clear" {
-                        // Clear screen and move cursor to top left
-                        print!("\x1B[2J\x1B[1;1H");
+                    if let Some(s) = line.first() {
+                        // These are shell only commands.
+                        match s.as_str() {
+                            "quit" | "exit" => {
+                                println!("Bye!");
+                                break 'shell;
+                            }
+                            "clear" => {
+                                // Clear screen and move cursor to top left
+                                print!("\x1B[2J\x1B[1;1H");
+                                continue 'shell;
+                            }
+                            "echo" => {
+                                let out = line.as_slice()[1..line.len()].join(" ");
+                                println!("{}", out);
+                                continue 'shell;
+                            }
+                            "env" => {
+                                for (key, var) in env::vars() {
+                                    println!("{}={}", key, var);
+                                }
+                                continue 'shell;
+                            }
+                            _ => {}
+                        }
+                    } else {
+                        // do nothing if line is empty
                         continue 'shell;
-                    };
-                    if self
-                        .handler
-                        .handle_async(line, &mut self.state, &self.description)
-                        .await
-                    {
+                    }
+
+                    if self.handler.handle_async(line, &mut self.state).await {
                         break 'shell;
                     };
                 }
@@ -96,6 +103,28 @@ impl<P: Display, S: Send, H: AsyncHandler<S>> Shell<P, S, H> {
         Ok(())
     }
 
+    fn substitution_env_variables(s: String) -> String {
+        if !s.contains('$') {
+            return s;
+        }
+        let mut env = env::vars().collect::<Vec<_>>();
+        // Sort variable name by the length in descending order, to prevent wrong substitution by variable with partial same name.
+        env.sort_by(|(k1, _), (k2, _)| Ord::cmp(&k2.len(), &k1.len()));
+
+        for (key, value) in env {
+            let var = format!("${}", key);
+            if s.contains(&var) {
+                let result = s.replace(var.as_str(), value.as_str());
+                return if result.contains('$') {
+                    Self::substitution_env_variables(result)
+                } else {
+                    result
+                };
+            }
+        }
+        s
+    }
+
     fn split_and_unescape(line: &str) -> Result<Vec<String>, String> {
         let mut commands = Vec::new();
         for word in line.split_whitespace() {
@@ -107,6 +136,17 @@ impl<P: Display, S: Send, H: AsyncHandler<S>> Shell<P, S, H> {
     }
 }
 
+pub fn install_shell_plugins<'a>(clap: App<'a, 'a>) -> App<'a, 'a> {
+    clap.subcommand(
+        SubCommand::with_name("exit")
+            .alias("quit")
+            .about("Exit the interactive shell"),
+    )
+    .subcommand(SubCommand::with_name("clear").about("Clear screen"))
+    .subcommand(SubCommand::with_name("echo").about("Write arguments to the console output"))
+    .subcommand(SubCommand::with_name("env").about("Print environment"))
+}
+
 #[derive(Helper)]
 struct ShellHelper {
     pub command: CommandStructure,
@@ -217,5 +257,5 @@ impl CommandStructure {
 
 #[async_trait]
 pub trait AsyncHandler<T: Send> {
-    async fn handle_async(&self, args: Vec<String>, state: &mut T, description: &str) -> bool;
+    async fn handle_async(&self, args: Vec<String>, state: &mut T) -> bool;
 }

sui/src/wallet.rs

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 use structopt::clap::{App, AppSettings};
 use structopt::StructOpt;
 use sui::config::{Config, WalletConfig};
-use sui::shell::{AsyncHandler, CommandStructure, Shell};
+use sui::shell::{install_shell_plugins, AsyncHandler, CommandStructure, Shell};
 use sui::wallet_commands::*;
 use tracing::error;
 
@@ -67,7 +67,7 @@ async fn main() -> Result<(), anyhow::Error> {
     if !options.no_shell {
         let app: App = WalletCommands::clap();
         println!("{}", SUI.cyan().bold());
-        print!("--- Sui");
+        print!("--- ");
         app.write_long_version(&mut io::stdout())?;
         println!(" ---");
         println!("{}", context.config);
@@ -79,9 +79,9 @@ async fn main() -> Result<(), anyhow::Error> {
             prompt: "sui>-$ ",
             state: context,
             handler: ClientCommandHandler,
-            description: String::new(),
-            command: CommandStructure::from_clap(&app),
+            command: CommandStructure::from_clap(&install_shell_plugins(app)),
         };
+
         shell.run_async().await?;
     } else if let Some(mut cmd) = options.cmd {
         cmd.execute(&mut context).await?.print(!options.json);
@@ -93,17 +93,21 @@ struct ClientCommandHandler;
 
 #[async_trait]
 impl AsyncHandler<WalletContext> for ClientCommandHandler {
-    async fn handle_async(&self, args: Vec<String>, context: &mut WalletContext, _: &str) -> bool {
-        let command: Result<WalletOpts, _> = WalletOpts::from_iter_safe(args);
-        if let Err(e) = handle_command(command, context).await {
+    async fn handle_async(&self, args: Vec<String>, context: &mut WalletContext) -> bool {
+        if let Err(e) = handle_command(get_command(args), context).await {
             error!("{}", e.to_string().red());
         }
         false
     }
 }
 
+fn get_command(args: Vec<String>) -> Result<WalletOpts, anyhow::Error> {
+    let app: App = install_shell_plugins(WalletOpts::clap());
+    Ok(WalletOpts::from_clap(&app.get_matches_from_safe(args)?))
+}
+
 async fn handle_command(
-    wallet_opts: Result<WalletOpts, structopt::clap::Error>,
+    wallet_opts: Result<WalletOpts, anyhow::Error>,
     context: &mut WalletContext,
 ) -> Result<(), anyhow::Error> {
     let mut wallet_opts = wallet_opts?;

sui/src/wallet_commands.rs

@@ -378,18 +378,18 @@ impl Display for WalletCommandResult {
         let mut writer = String::new();
         match self {
             WalletCommandResult::Publish(cert, effects) => {
-                writeln!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
+                write!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
             }
             WalletCommandResult::Object(object_read) => {
                 let object = object_read.object().map_err(fmt::Error::custom)?;
                 writeln!(writer, "{}", object)?;
             }
             WalletCommandResult::Call(cert, effects) => {
-                writeln!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
+                write!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
             }
             WalletCommandResult::Transfer(time_elapsed, cert, effects) => {
                 writeln!(writer, "Transfer confirmed after {} us", time_elapsed)?;
-                writeln!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
+                write!(writer, "{}", write_cert_and_effects(cert, effects)?)?;
             }
             WalletCommandResult::Addresses(addresses) => {
                 writeln!(writer, "Showing {} results.", addresses.len())?;
@@ -441,9 +441,9 @@ fn write_cert_and_effects(
 ) -> Result<String, fmt::Error> {
     let mut writer = String::new();
     writeln!(writer, "{}", "----- Certificate ----".bold())?;
-    writeln!(writer, "{}", cert)?;
+    write!(writer, "{}", cert)?;
     writeln!(writer, "{}", "----- Transaction Effects ----".bold())?;
-    writeln!(writer, "{}", effects)?;
+    write!(writer, "{}", effects)?;
     Ok(writer)
 }