[aptos-tool] Cleanup common code into proper locations

There was a lot of common code spread across key infromation and utils, so
I've compiled the common types and utils accordingly.

Closes: aptos-labs#433

crates/aptos-faucet/src/main.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
-use aptos::{common::types::EncodingType, op::key::load_key};
+use aptos::common::types::EncodingType;
 use aptos_crypto::ed25519;
 use aptos_logger::info;
 use aptos_sdk::types::{
@@ -67,8 +67,9 @@ async fn main() {
         args.maximum_amount,
     );
 
-    let key: ed25519::Ed25519PrivateKey =
-        load_key(Path::new(&args.mint_key_file_path), EncodingType::BCS).unwrap();
+    let key: ed25519::Ed25519PrivateKey = EncodingType::BCS
+        .load_key(Path::new(&args.mint_key_file_path))
+        .unwrap();
 
     let faucet_address: AccountAddress =
         args.mint_account_address.unwrap_or_else(aptos_root_address);

crates/aptos/src/common/types.rs

@@ -1,10 +1,14 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::op::key::{decode_key, load_key};
-use aptos_crypto::x25519;
-use clap::Parser;
-use std::{fmt::Debug, path::PathBuf, str::FromStr};
+use crate::common::utils::{check_if_file_exists, write_to_file};
+use aptos_crypto::{x25519, ValidCryptoMaterial, ValidCryptoMaterialStringExt};
+use clap::{ArgEnum, Parser};
+use std::{
+    fmt::Debug,
+    path::{Path, PathBuf},
+    str::FromStr,
+};
 use thiserror::Error;
 
 /// A common result to be returned to users
@@ -44,7 +48,7 @@ pub enum Error {
 }
 
 /// Types of Keys used by the blockchain
-#[derive(Clone, Copy, Debug, Parser)]
+#[derive(ArgEnum, Clone, Copy, Debug)]
 pub enum KeyType {
     Ed25519,
     X25519,
@@ -63,7 +67,7 @@ impl FromStr for KeyType {
 }
 
 /// Types of encodings used by the blockchain
-#[derive(Clone, Copy, Debug, Parser)]
+#[derive(ArgEnum, Clone, Copy, Debug)]
 pub enum EncodingType {
     /// Binary Canonical Serialization
     BCS,
@@ -73,6 +77,53 @@ pub enum EncodingType {
     Base64,
 }
 
+impl EncodingType {
+    /// Encodes `Key` into one of the `EncodingType`s
+    pub fn encode_key<Key: ValidCryptoMaterial>(
+        &self,
+        key: &Key,
+        key_name: &str,
+    ) -> Result<Vec<u8>, Error> {
+        Ok(match self {
+            EncodingType::Hex => hex::encode_upper(key.to_bytes()).into_bytes(),
+            EncodingType::BCS => {
+                bcs::to_bytes(key).map_err(|err| Error::BCS(key_name.to_string(), err))?
+            }
+            EncodingType::Base64 => base64::encode(key.to_bytes()).into_bytes(),
+        })
+    }
+
+    /// Loads a key from a file
+    pub fn load_key<Key: ValidCryptoMaterial>(&self, path: &Path) -> Result<Key, Error> {
+        let data = std::fs::read(&path).map_err(|err| {
+            Error::UnableToReadFile(path.to_str().unwrap().to_string(), err.to_string())
+        })?;
+
+        self.decode_key(data)
+    }
+
+    /// Decodes an encoded key given the known encoding
+    pub fn decode_key<Key: ValidCryptoMaterial>(&self, data: Vec<u8>) -> Result<Key, Error> {
+        match self {
+            EncodingType::BCS => {
+                bcs::from_bytes(&data).map_err(|err| Error::BCS("Key".to_string(), err))
+            }
+            EncodingType::Hex => {
+                let hex_string = String::from_utf8(data).unwrap();
+                Key::from_encoded_string(hex_string.trim())
+                    .map_err(|err| Error::UnableToParse("Key", err.to_string()))
+            }
+            EncodingType::Base64 => {
+                let string = String::from_utf8(data).unwrap();
+                let bytes = base64::decode(string.trim())
+                    .map_err(|err| Error::UnableToParse("Key", err.to_string()))?;
+                Key::try_from(bytes.as_slice())
+                    .map_err(|err| Error::UnexpectedError(format!("Failed to parse key {}", err)))
+            }
+        }
+    }
+}
+
 impl FromStr for EncodingType {
     type Err = &'static str;
 
@@ -115,10 +166,10 @@ pub struct PrivateKeyInputOptions {
 impl PrivateKeyInputOptions {
     pub fn extract_private_key(&self, encoding: EncodingType) -> Result<x25519::PrivateKey, Error> {
         if let Some(ref file) = self.private_key_file {
-            load_key(file.as_path(), encoding)
+            encoding.load_key(file.as_path())
         } else if let Some(ref key) = self.private_key {
             let key = key.as_bytes().to_vec();
-            decode_key(key, encoding)
+            encoding.decode_key(key)
         } else {
             Err(Error::CommandArgumentError(
                 "One of ['--private-key', '--private-key-file'] must be used".to_string(),
@@ -140,10 +191,10 @@ pub struct PublicKeyInputOptions {
 impl PublicKeyInputOptions {
     pub fn extract_public_key(&self, encoding: EncodingType) -> Result<x25519::PublicKey, Error> {
         if let Some(ref file) = self.public_key_file {
-            load_key(file.as_path(), encoding)
+            encoding.load_key(file.as_path())
         } else if let Some(ref key) = self.public_key {
             let key = key.as_bytes().to_vec();
-            decode_key(key, encoding)
+            encoding.decode_key(key)
         } else {
             Err(Error::CommandArgumentError(
                 "One of ['--public-key', '--public-key-file'] must be used".to_string(),
@@ -161,6 +212,7 @@ 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);
@@ -175,3 +227,25 @@ impl KeyInputOptions {
         }
     }
 }
+
+#[derive(Debug, Parser)]
+pub struct SaveFile {
+    /// Output file name
+    #[clap(long, parse(from_os_str))]
+    pub output_file: PathBuf,
+
+    #[clap(flatten)]
+    pub prompt_options: PromptOptions,
+}
+
+impl SaveFile {
+    /// Check if the key file exists already
+    pub fn check_file(&self) -> Result<(), Error> {
+        check_if_file_exists(self.output_file.as_path(), self.prompt_options.assume_yes)
+    }
+
+    /// Save to the `output_file`
+    pub fn save_to_file(&self, name: &str, bytes: &[u8]) -> Result<(), Error> {
+        write_to_file(self.output_file.as_path(), name, bytes)
+    }
+}

crates/aptos/src/common/utils.rs

@@ -3,7 +3,11 @@
 
 use crate::{CliResult, Error};
 use serde::Serialize;
-use std::path::{Path, PathBuf};
+use std::{
+    fs::File,
+    io::Write,
+    path::{Path, PathBuf},
+};
 
 /// Prompts for confirmation until a yes or no is given explicitly
 /// TODO: Capture interrupts
@@ -72,6 +76,31 @@ impl<T> From<Result<T, Error>> for ResultWrapper<T> {
     }
 }
 
+/// Checks if a file exists, being overridden by `--assume-yes`
+pub fn check_if_file_exists(file: &Path, assume_yes: bool) -> Result<(), Error> {
+    if file.exists()
+        && !assume_yes
+        && !prompt_yes(
+            format!(
+                "{:?} already exists, are you sure you want to overwrite it?",
+                file.as_os_str()
+            )
+            .as_str(),
+        )
+    {
+        Err(Error::AbortedError)
+    } else {
+        Ok(())
+    }
+}
+
+/// 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)
+        .map_err(|e| Error::IO(name.to_string(), e))
+}
+
 /// Appends a file extension to a `Path` without overwriting the original extension.
 pub fn append_file_extension(
     file: &Path,