[aptos-tool] Create init command and configuration file

crates/aptos/Cargo.toml

@@ -24,6 +24,7 @@ tokio-util = { version = "0.6.4", features = ["compat"] }
 
 aptos-config = { path = "../../config" }
 aptos-crypto = { path = "../aptos-crypto" }
+aptos-logger = { path = "../aptos-logger" }
 aptos-secure-storage = { path = "../../secure/storage" }
 aptos-telemetry = { path = "../aptos-telemetry" }
 aptos-temppath = { path = "../aptos-temppath" }

crates/aptos/src/common/init.rs

@@ -0,0 +1,63 @@
+// Copyright (c) Aptos
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{
+    common::{
+        types::CliConfig,
+        utils::{prompt_yes, to_common_success_result},
+    },
+    op::key::GenerateKey,
+    CliResult, Error,
+};
+use aptos_crypto::{x25519::PrivateKey, ValidCryptoMaterialStringExt};
+use clap::Parser;
+
+/// Tool to initialize current directory for the aptos tool
+#[derive(Debug, Parser)]
+pub struct InitTool {}
+
+impl InitTool {
+    pub async fn execute(&self) -> CliResult {
+        to_common_success_result(self.execute_inner().await)
+    }
+
+    async fn execute_inner(&self) -> Result<(), Error> {
+        let mut config = if CliConfig::config_exists()? {
+            if !prompt_yes(
+                "Aptos already initialized, do you want to overwrite the existing config?",
+            ) {
+                eprintln!("Exiting...");
+                return Ok(());
+            }
+            CliConfig::load()?
+        } else {
+            CliConfig::default()
+        };
+
+        eprintln!("Enter your private key as a hex literal (0x...) [No input: Generate new key]");
+        let input = read_line()?;
+        let input = input.trim();
+        let private_key = if input.is_empty() {
+            eprintln!("No key given, generating key...");
+            GenerateKey::generate_x25519_in_memory()?
+        } else {
+            PrivateKey::from_encoded_string(input)
+                .map_err(|err| Error::UnableToParse("PrivateKey", err.to_string()))?
+        };
+        config.private_key = Some(private_key);
+        config.save()?;
+        eprintln!("Aptos is now set up!  Run `aptos help` for more information about commands");
+
+        Ok(())
+    }
+}
+
+/// Reads a line from input
+fn read_line() -> Result<String, Error> {
+    let mut input_buf = String::new();
+    let _ = std::io::stdin()
+        .read_line(&mut input_buf)
+        .map_err(|err| Error::IO("Private key".to_string(), err))?;
+
+    Ok(input_buf)
+}

crates/aptos/src/common/mod.rs

@@ -1,5 +1,6 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod init;
 pub mod types;
 pub mod utils;

crates/aptos/src/common/types.rs

@@ -3,7 +3,9 @@
 
 use crate::common::utils::{check_if_file_exists, write_to_file};
 use aptos_crypto::{x25519, ValidCryptoMaterial, ValidCryptoMaterialStringExt};
+use aptos_logger::{debug, info};
 use clap::{ArgEnum, Parser};
+use serde::{Deserialize, Serialize};
 use std::{
     fmt::Debug,
     path::{Path, PathBuf},
@@ -27,6 +29,8 @@ pub enum Error {
     CommandArgumentError(String),
     #[error("Unable to load config: {0}")]
     ConfigError(String),
+    #[error("Unable to find config {0}, have you run `aptos init`?")]
+    ConfigNotFoundError(String),
     #[error("Error accessing '{0}': {1}")]
     IO(String, #[source] std::io::Error),
     #[error("Error (de)serializing '{0}': {1}")]
@@ -47,6 +51,67 @@ pub enum Error {
     AbortedError,
 }
 
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub struct CliConfig {
+    /// Private key for commands.  TODO: Add vault functionality
+    pub private_key: Option<x25519::PrivateKey>,
+}
+
+impl CliConfig {
+    /// Checks if the config exists in the current working directory
+    pub fn config_exists() -> Result<bool, Error> {
+        Self::aptos_folder().map(|folder| folder.exists())
+    }
+
+    /// Loads the config from the current working directory
+    pub fn load() -> Result<Self, Error> {
+        let config_file = Self::aptos_folder()?.join("config.yml");
+        if !config_file.exists() {
+            return Err(Error::ConfigError(format!("{:?}", config_file)));
+        }
+
+        let bytes = std::fs::read(&config_file)
+            .map_err(|err| Error::IO(format!("Failed to read {:?}", config_file), err))?;
+        serde_yaml::from_slice(&bytes)
+            .map_err(|err| Error::UnableToParseFile(format!("{:?}", config_file), err.to_string()))
+    }
+
+    /// Saves the config to ./.aptos/config.yml
+    pub fn save(&self) -> Result<(), Error> {
+        let aptos_folder = Self::aptos_folder()?;
+
+        // Create if it doesn't exist
+        if !aptos_folder.exists() {
+            std::fs::create_dir(&aptos_folder).map_err(|err| {
+                Error::CommandArgumentError(format!(
+                    "Unable to create {:?} directory {}",
+                    aptos_folder, err
+                ))
+            })?;
+            info!("Created .aptos/ folder");
+        } else {
+            debug!(".aptos/ folder already initialized");
+        }
+
+        // Save over previous config file
+        // TODO: Ask for saving over?
+        let config_file = aptos_folder.join("config.yml");
+        let config_bytes = serde_yaml::to_string(&self)
+            .map_err(|err| Error::UnexpectedError(format!("Failed to serialize config {}", err)))?;
+        write_to_file(&config_file, "config.yml", config_bytes.as_bytes())?;
+        Ok(())
+    }
+
+    /// Finds the current directory's .aptos folder
+    fn aptos_folder() -> Result<PathBuf, Error> {
+        std::env::current_dir()
+            .map_err(|err| {
+                Error::UnexpectedError(format!("Unable to get current directory {}", err))
+            })
+            .map(|dir| dir.join(".aptos"))
+    }
+}
+
 /// Types of Keys used by the blockchain
 #[derive(ArgEnum, Clone, Copy, Debug)]
 pub enum KeyType {
@@ -164,16 +229,17 @@ pub struct PrivateKeyInputOptions {
 }
 
 impl PrivateKeyInputOptions {
-    pub fn extract_private_key(&self, encoding: EncodingType) -> Result<x25519::PrivateKey, Error> {
+    pub fn extract_private_key(
+        &self,
+        encoding: EncodingType,
+    ) -> Result<Option<x25519::PrivateKey>, Error> {
         if let Some(ref file) = self.private_key_file {
-            encoding.load_key(file.as_path())
+            encoding.load_key(file.as_path()).map(Some)
         } else if let Some(ref key) = self.private_key {
             let key = key.as_bytes().to_vec();
-            encoding.decode_key(key)
+            encoding.decode_key(key).map(Some)
         } else {
-            Err(Error::CommandArgumentError(
-                "One of ['--private-key', '--private-key-file'] must be used".to_string(),
-            ))
+            Ok(None)
         }
     }
 }
@@ -189,16 +255,17 @@ pub struct PublicKeyInputOptions {
 }
 
 impl PublicKeyInputOptions {
-    pub fn extract_public_key(&self, encoding: EncodingType) -> Result<x25519::PublicKey, Error> {
+    pub fn extract_public_key(
+        &self,
+        encoding: EncodingType,
+    ) -> Result<Option<x25519::PublicKey>, Error> {
         if let Some(ref file) = self.public_key_file {
-            encoding.load_key(file.as_path())
+            encoding.load_key(file.as_path()).map(Some)
         } else if let Some(ref key) = self.public_key {
             let key = key.as_bytes().to_vec();
-            encoding.decode_key(key)
+            encoding.decode_key(key).map(Some)
         } else {
-            Err(Error::CommandArgumentError(
-                "One of ['--public-key', '--public-key-file'] must be used".to_string(),
-            ))
+            Ok(None)
         }
     }
 }
@@ -214,16 +281,21 @@ pub struct KeyInputOptions {
 impl KeyInputOptions {
     /// Extracts public key from either private or public key options
     pub fn extract_public_key(&self, encoding: EncodingType) -> Result<x25519::PublicKey, Error> {
-        let private_key_result = self.private_key_options.extract_private_key(encoding);
-        let public_key_result = self.public_key_options.extract_public_key(encoding);
+        let private_key_result = self.private_key_options.extract_private_key(encoding)?;
+        let public_key_result = self.public_key_options.extract_public_key(encoding)?;
 
-        if let Ok(private_key) = private_key_result {
+        if let Some(private_key) = private_key_result {
             Ok(private_key.public_key())
-        } else if let Ok(public_key) = public_key_result {
+        } else if let Some(public_key) = public_key_result {
             Ok(public_key)
         } else {
-            // TODO: merge above errors better
-            Err(Error::CommandArgumentError("One of ['--private-key', '--private-key-file', '--public-key', '--public-key-file'] must be used".to_string()))
+            let config = CliConfig::load()?;
+            if let Some(private_key) = config.private_key {
+                println!("Using .aptos/config.yml private key");
+                Ok(private_key.public_key())
+            } else {
+                Err(Error::CommandArgumentError("One of ['--private-key', '--private-key-file', '--public-key', '--public-key-file'] must be used".to_string()))
+            }
         }
     }
 }

crates/aptos/src/common/utils.rs

@@ -95,9 +95,9 @@ pub fn check_if_file_exists(file: &Path, assume_yes: bool) -> Result<(), Error>
 }
 
 /// Write a `&[u8]` to a file
-pub fn write_to_file(key_file: &Path, name: &str, encoded: &[u8]) -> Result<(), Error> {
-    let mut file = File::create(key_file).map_err(|e| Error::IO(name.to_string(), e))?;
-    file.write_all(encoded)
+pub fn write_to_file(path: &Path, name: &str, bytes: &[u8]) -> Result<(), Error> {
+    let mut file = File::create(path).map_err(|e| Error::IO(name.to_string(), e))?;
+    file.write_all(bytes)
         .map_err(|e| Error::IO(name.to_string(), e))
 }
 

crates/aptos/src/lib.rs

@@ -16,6 +16,7 @@ use clap::Parser;
 #[derive(Parser)]
 #[clap(name = "aptos", author, version, propagate_version = true)]
 pub enum Tool {
+    Init(common::init::InitTool),
     List(list::ListResources),
     #[clap(subcommand)]
     Move(move_tool::MoveTool),
@@ -26,9 +27,10 @@ pub enum Tool {
 impl Tool {
     pub async fn execute(self) -> CliResult {
         match self {
-            Tool::List(list_tool) => list_tool.execute().await,
+            Tool::Init(tool) => tool.execute().await,
+            Tool::List(tool) => tool.execute().await,
             Tool::Move(tool) => tool.execute().await,
-            Tool::Op(op_tool) => op_tool.execute().await,
+            Tool::Op(tool) => tool.execute().await,
         }
     }
 }

crates/aptos/src/main.rs

@@ -8,11 +8,20 @@
 #![forbid(unsafe_code)]
 
 use aptos::Tool;
+use aptos_logger::Level;
 use clap::Parser;
 use std::process::exit;
 
 #[tokio::main]
 async fn main() {
+    let mut logger = aptos_logger::Logger::new();
+    logger
+        .channel_size(1000)
+        .is_async(false)
+        .level(Level::Info)
+        .read_env();
+    logger.build();
+
     // Run the corresponding tools
     let result = Tool::parse().execute().await;
 

crates/aptos/src/op/key.rs

@@ -160,6 +160,13 @@ impl GenerateKey {
         let mut rng = rand::rngs::StdRng::from_entropy();
         Ed25519PrivateKey::generate(&mut rng)
     }
+
+    pub fn generate_x25519_in_memory() -> Result<x25519::PrivateKey, Error> {
+        let key = Self::generate_ed25519_in_memory();
+        x25519::PrivateKey::from_ed25519_private_bytes(&key.to_bytes()).map_err(|err| {
+            Error::UnexpectedError(format!("Failed to convert ed25519 to x25519 {:?}", err))
+        })
+    }
 }
 
 #[derive(Debug, Parser)]