From 8336d6c634c3c3385a997e93e267b02299ca89f4 Mon Sep 17 00:00:00 2001
From: Philip Vu <pvu2002@outlook.com>
Date: Mon, 21 Oct 2024 16:09:29 -0700
Subject: [PATCH] Allow for AIP-80 compliant strings in imports

---
 crates/aptos/CHANGELOG.md        |  1 +
 crates/aptos/src/common/init.rs  | 12 ++++++++----
 crates/aptos/src/common/types.rs |  4 ++--
 crates/aptos/src/common/utils.rs | 12 ++++++++++++
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/crates/aptos/CHANGELOG.md b/crates/aptos/CHANGELOG.md
index 5f56758e8c928..1dec577ea8e51 100644
--- a/crates/aptos/CHANGELOG.md
+++ b/crates/aptos/CHANGELOG.md
@@ -8,6 +8,7 @@ All notable changes to the Aptos CLI will be captured in this file. This project
 - Update the default version of `movefmt` to be installed from 1.0.4 to 1.0.5
 - Update the local-testnet logs to use `println` for regular output and reserve `eprintln` for errors.
 - Set compiler V2 as default when using `aptos move prove`.
+- Add support for AIP-80 compliant strings when importing using the CLI arguments or manual input.
 
 ## [4.2.3] - 2024/09/20
 - Fix the broken indexer in localnet in 4.2.2, which migrates table info from sycn to async ways.
diff --git a/crates/aptos/src/common/init.rs b/crates/aptos/src/common/init.rs
index db590ff97faea..cb3e526d2713f 100644
--- a/crates/aptos/src/common/init.rs
+++ b/crates/aptos/src/common/init.rs
@@ -9,7 +9,10 @@ use crate::{
             ConfigSearchMode, EncodingOptions, HardwareWalletOptions, PrivateKeyInputOptions,
             ProfileConfig, ProfileOptions, PromptOptions, RngArgs, DEFAULT_PROFILE,
         },
-        utils::{explorer_account_link, fund_account, prompt_yes_with_override, read_line},
+        utils::{
+            explorer_account_link, fund_account, prompt_yes_with_override, read_line,
+            strip_private_key_prefix,
+        },
     },
 };
 use aptos_crypto::{ed25519::Ed25519PrivateKey, PrivateKey, ValidCryptoMaterialStringExt};
@@ -218,9 +221,10 @@ impl CliCommand<()> for InitTool {
                             .generate_ed25519_private_key()
                     }
                 } else {
-                    Ed25519PrivateKey::from_encoded_string(input).map_err(|err| {
-                        CliError::UnableToParse("Ed25519PrivateKey", err.to_string())
-                    })?
+                    Ed25519PrivateKey::from_encoded_string(&strip_private_key_prefix(
+                        &input.to_string(),
+                    ))
+                    .map_err(|err| CliError::UnableToParse("Ed25519PrivateKey", err.to_string()))?
                 }
             };
 
diff --git a/crates/aptos/src/common/types.rs b/crates/aptos/src/common/types.rs
index 0b5fc13047670..6c96e01e3e4b3 100644
--- a/crates/aptos/src/common/types.rs
+++ b/crates/aptos/src/common/types.rs
@@ -1,7 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use super::utils::{explorer_transaction_link, fund_account};
+use super::utils::{explorer_transaction_link, fund_account, strip_private_key_prefix};
 use crate::{
     common::{
         init::Network,
@@ -675,7 +675,7 @@ pub trait ParsePrivateKey {
                 encoding.load_key("--private-key-file", file.as_path())?,
             ))
         } else if let Some(ref key) = private_key {
-            let key = key.as_bytes().to_vec();
+            let key = strip_private_key_prefix(key).as_bytes().to_vec();
             Ok(Some(encoding.decode_key("--private-key", key)?))
         } else {
             Ok(None)
diff --git a/crates/aptos/src/common/utils.rs b/crates/aptos/src/common/utils.rs
index 90e3ddd983787..94c12570b25be 100644
--- a/crates/aptos/src/common/utils.rs
+++ b/crates/aptos/src/common/utils.rs
@@ -597,3 +597,15 @@ pub fn explorer_transaction_link(
         )
     }
 }
+
+/// Strips the private key prefix for a given key string if it is AIP-80 compliant.
+///
+/// [Read about AIP-80](https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md)
+pub fn strip_private_key_prefix(key: &String) -> String {
+    let prefixes = ["ed25519-priv-", "secp256k1-priv-"];
+    prefixes
+        .iter()
+        .find_map(|prefix| key.strip_prefix(prefix))
+        .unwrap_or(key) // If no prefix is found, return the original key
+        .to_string()
+}