Merge pull request #8 from pshenmic/feat/private-key-types

Accept all possible private key encoding types (base64, hex, wif)
pshenmic · Oct 20, 2024 · 6bba5f3 · 6bba5f3
2 parents 4c88e0f + 10f8e30
commit 6bba5f3
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 29 deletions.
diff --git a/README.md b/README.md
@@ -42,9 +42,9 @@ You can use --help to get more info about command flags
 
 Example commands:
 ```bash
-$ platform-cli withdraw --dapi-url https://127.0.0.1:1443 --identity A1rgGVjRGuznRThdAA316VEEpKuVQ7mV8mBK1BFJvXnb --private-key private_key.txt --withdrawal-address yifJkXaxe7oM1NgBDTaXnWa6kXZAazBfjk --amount 40000
-$ platform-cli register-dpns-name --identity 8eTDkBhpQjHeqgbVeriwLeZr1tCa6yBGw76SckvD1cwc --private-key private_key.txt --dapi-url https://52.43.13.92:1443 --name tesstst32423sts
-$ platform-cli masternode-vote-dpns-name --dapi-url https://52.43.13.92:1443 --private-key voting_key.txt --pro-tx-hash 7a1ae04de7582262d9dea3f4d72bc24a474c6f71988066b74a41f17be5552652 --normalized-label testc0ntested --choice 8eTDkBhpQjHeqgbVeriwLeZr1tCa6yBGw76SckvD1cwc
+$ platform-cli withdraw --network testnet --dapi-url https://127.0.0.1:1443 --identity A1rgGVjRGuznRThdAA316VEEpKuVQ7mV8mBK1BFJvXnb --private-key private_key.txt --withdrawal-address yifJkXaxe7oM1NgBDTaXnWa6kXZAazBfjk --amount 40000
+$ platform-cli register-dpns-name --network testnet --dapi-url https://52.43.13.92:1443 --identity 8eTDkBhpQjHeqgbVeriwLeZr1tCa6yBGw76SckvD1cwc --private-key private_key.txt --name tesstst32423sts
+$ platform-cli masternode-vote-dpns-name --network testnet --dapi-url https://52.43.13.92:1443 --private-key voting_key.txt --pro-tx-hash 7a1ae04de7582262d9dea3f4d72bc24a474c6f71988066b74a41f17be5552652 --normalized-label testc0ntested --choice 8eTDkBhpQjHeqgbVeriwLeZr1tCa6yBGw76SckvD1cwc
 ```
 
 ### Credits Withdrawal
@@ -54,12 +54,14 @@ Withdraw credits from the Identity to the L1 Core chain
 Usage: platform-cli withdraw [OPTIONS]
 
 Options:
+      --network <NETWORK>
+          Network, mainnet or testnet [default: ]
       --dapi-url <DAPI_URL>
           DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443 [default: ]
       --identity <IDENTITY>
           Identity address, that initiate withdrawal [default: ]
       --private-key <PRIVATE_KEY>
-          Identity private key in WIF format [default: ]
+          Path to file with private key from Identity in WIF format [default: ]
       --withdrawal-address <WITHDRAWAL_ADDRESS>
           Core withdrawal address (P2PKH / P2SH) [default: ]
       --amount <AMOUNT>
@@ -82,9 +84,10 @@ Register an Identity Name in the Dash Platform DPNS system
 Usage: platform-cli register-dpns-name [OPTIONS]
 
 Options:
+      --network <NETWORK>          Network, mainnet or testnet [default: ]
       --dapi-url <DAPI_URL>        DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443 [default: ]
       --identity <IDENTITY>        Identity address that registers a name [default: ]
-      --private-key <PRIVATE_KEY>  Identity private key in WIF format [default: ]
+      --private-key <PRIVATE_KEY>  Path to file with private key from Identity in WIF format [default: ]
       --name <NAME>                Name to register (excluding .dash) [default: ]
   -h, --help                       Print help
 ```
@@ -100,12 +103,14 @@ Perform a masternode vote towards contested DPNS name
 Usage: platform-cli masternode-vote-dpns-name [OPTIONS]
 
 Options:
+      --network <NETWORK>
+          Network, mainnet or testnet [default: ]
       --dapi-url <DAPI_URL>
           DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443 [default: ]
       --pro-tx-hash <PRO_TX_HASH>
           ProTxHash of the Masternode performing a Vote, in hex [default: ]
       --private-key <PRIVATE_KEY>
-          Voting (or Owner) private key in WIF format [default: ]
+          Path to file with voting (or owner) private key in WIF format [default: ]
       --normalized-label <NORMALIZED_LABEL>
           Normalized label to vote upon (can be grabbed from https//dash.vote) [default: ]
       --choice <CHOICE>

diff --git a/src/commands/masternode_vote_dpns_name.rs b/src/commands/masternode_vote_dpns_name.rs
@@ -1,8 +1,9 @@
 use std::fs;
+use std::str::FromStr;
 use clap::Parser;
 use dpp::dashcore::hashes::Hash;
 use dpp::dashcore::key::Secp256k1;
-use dpp::dashcore::{PrivateKey, ProTxHash};
+use dpp::dashcore::{Network, ProTxHash};
 use dpp::identifier::{Identifier, MasternodeIdentifiers};
 use dpp::identity::accessors::IdentityGettersV0;
 use dpp::identity::hash::IdentityPublicKeyHashMethodsV0;
@@ -18,10 +19,15 @@ use crate::errors::identity_public_key_hash_mismatch_error::IdentityPublicKeyHas
 use crate::factories::Factories;
 use crate::grpc::PlatformGRPCClient;
 use crate::MockBLS;
+use crate::utils::Utils;
 
 /// Perform a masternode vote towards contested DPNS name
 #[derive(Parser)]
 pub struct MasternodeVoteDPNSNameCommand {
+    /// Network, mainnet or testnet
+    #[clap(long, default_value(""))]
+    network: String,
+
     /// DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443
     #[clap(long, default_value(""))]
     dapi_url: String,
@@ -48,6 +54,9 @@ const DPNS_DATA_CONTRACT_IDENTIFIER: &str = "GWRSAVFMjXx8HpQFaNJMqBV7MBgMK4br5UE
 
 impl MasternodeVoteDPNSNameCommand {
     pub async fn run(&self) -> Result<(), Error> {
+        if self.network.is_empty() {
+            return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("network")));
+        }
         if self.pro_tx_hash.is_empty() {
             return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("pro_tx_hash")));
         }
@@ -66,9 +75,9 @@ impl MasternodeVoteDPNSNameCommand {
 
         let secp = Secp256k1::new();
 
-        let private_key_data = fs::read_to_string(&self.private_key).expect("Unable to read file");
-        let (private_key_data_stripped, _) = private_key_data.split_at(52);
-        let private_key = PrivateKey::from_wif(&private_key_data_stripped).expect("Could not load private key from WIF");
+        let network_type = Network::from_str(&self.network).expect("Could not parse network");
+        let private_key_data = fs::read_to_string(&self.private_key).expect("Unable to read private key file");
+        let private_key = Utils::decode_private_key_from_input_string(private_key_data.as_str(), network_type)?;
         let public_key = private_key.public_key(&secp);
         let pro_tx_hash = ProTxHash::from_hex(&self.pro_tx_hash).expect("Could not decode pro tx hash");
         let voting_address = public_key.pubkey_hash().to_byte_array();

diff --git a/src/commands/register_dpns_name.rs b/src/commands/register_dpns_name.rs
@@ -4,7 +4,7 @@ use std::str::FromStr;
 use std::time::Duration;
 use clap::{ Parser};
 use dpp::dashcore::hashes::Hash;
-use dpp::dashcore::PrivateKey;
+use dpp::dashcore::{Network};
 use dpp::dashcore::secp256k1::Secp256k1;
 use dpp::data_contract::accessors::v0::DataContractV0Getters;
 use dpp::data_contract::{DataContract, JsonValue};
@@ -33,13 +33,17 @@ use crate::errors::identity_public_key_hash_mismatch_error::IdentityPublicKeyHas
 use crate::factories::create_documents_batch::IdentityStateTransition;
 use crate::factories::Factories;
 use crate::grpc::PlatformGRPCClient;
-use crate::utils::MyDefaultEntropyGenerator;
+use crate::utils::{MyDefaultEntropyGenerator, Utils};
 use regex::Regex;
 use crate::MockBLS;
 
 /// Register an Identity Name in the Dash Platform DPNS system.
 #[derive(Parser)]
 pub struct RegisterDPNSNameCommand {
+    /// Network, mainnet or testnet
+    #[clap(long, default_value(""))]
+    network: String,
+
     /// DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443
     #[clap(long, default_value(""))]
     dapi_url: String,
@@ -59,6 +63,10 @@ pub struct RegisterDPNSNameCommand {
 
 impl RegisterDPNSNameCommand {
     pub async fn run(&self) -> Result<(), Error> {
+        if self.network.is_empty() {
+            return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("network")));
+        }
+
         if self.private_key.is_empty() {
             return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("private_key")));
         }
@@ -77,9 +85,9 @@ impl RegisterDPNSNameCommand {
 
         let secp = Secp256k1::new();
 
-        let private_key_data = fs::read_to_string(&self.private_key).expect("Unable to read file");
-        let (private_key_data_stripped, _) = private_key_data.split_at(52);
-        let private_key = PrivateKey::from_wif(&private_key_data_stripped).expect("Could not load private key from WIF");
+        let network_type = Network::from_str(&self.network).expect("Could not parse network");
+        let private_key_data = fs::read_to_string(&self.private_key).expect("Unable to read private key file");
+        let private_key = Utils::decode_private_key_from_input_string(private_key_data.as_str(), network_type)?;
         let public_key = private_key.public_key(&secp);
         let identifier = Identifier::from_string(&self.identity, Base58).unwrap();
 
@@ -190,4 +198,5 @@ impl RegisterDPNSNameCommand {
 
         Ok(())
     }
-}
+}
+
diff --git a/src/commands/withdraw.rs b/src/commands/withdraw.rs
@@ -1,6 +1,7 @@
 use std::fs;
+use std::str::FromStr;
 use clap::Parser;
-use dpp::dashcore::{PrivateKey};
+use dpp::dashcore::{Network};
 use dpp::dashcore::hashes::Hash;
 use dpp::dashcore::secp256k1::Secp256k1;
 use dpp::identifier::Identifier;
@@ -20,10 +21,15 @@ use crate::errors::identity_public_key_hash_mismatch_error::IdentityPublicKeyHas
 use crate::errors::Error;
 use crate::grpc::PlatformGRPCClient;
 use crate::MockBLS;
+use crate::utils::Utils;
 
 /// Withdraw credits from the Identity to the L1 Core chain
 #[derive(Parser)]
 pub struct WithdrawCommand {
+    /// Network, mainnet or testnet
+    #[clap(long, default_value(""))]
+    network: String,
+
     /// DAPI GRPC Endpoint URL, ex. https://127.0.0.1:1443
     #[clap(long, default_value(""))]
     dapi_url: String,
@@ -47,6 +53,10 @@ pub struct WithdrawCommand {
 
 impl WithdrawCommand {
     pub async fn run(&self) -> Result<(), Error> {
+        if self.network.is_empty() {
+            return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("network")));
+        }
+
         if self.dapi_url.is_empty() {
             return Err(Error::CommandLineArgumentMissingError(CommandLineArgumentMissingError::from("dapi_url")));
         }
@@ -69,9 +79,9 @@ impl WithdrawCommand {
 
         let secp = Secp256k1::new();
 
-        let private_key_data = fs::read_to_string(&self.private_key.trim()).expect("Unable to read file");
-        let (private_key_data_stripped, _) = private_key_data.split_at(52);
-        let private_key = PrivateKey::from_wif(&private_key_data_stripped).expect("Could not load private key from WIF");
+        let network_type = Network::from_str(&self.network).expect("Could not parse network");
+        let private_key_data = fs::read_to_string(&self.private_key).expect("Unable to read private key file");
+        let private_key = Utils::decode_private_key_from_input_string(private_key_data.as_str(), network_type)?;
         let public_key = private_key.public_key(&secp);
 
         let platform_grpc_client = PlatformGRPCClient::new(&self.dapi_url);

diff --git a/src/errors/cli_argument_invalid_input.rs b/src/errors/cli_argument_invalid_input.rs
@@ -0,0 +1,16 @@
+use std::fmt;
+
+#[derive(Debug)]
+pub struct CommandLineArgumentInvalidInput(String);
+
+impl fmt::Display for CommandLineArgumentInvalidInput {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Input data is not valid: {}", &self.0)
+    }
+}
+
+impl From<&str> for CommandLineArgumentInvalidInput {
+    fn from(value: &str) -> Self {
+        CommandLineArgumentInvalidInput(String::from(value))
+    }
+}
diff --git a/src/errors/mod.rs b/src/errors/mod.rs
@@ -1,4 +1,5 @@
 use std::fmt::{Display, Formatter};
+use crate::errors::cli_argument_invalid_input::CommandLineArgumentInvalidInput;
 use crate::errors::cli_argument_missing_error::CommandLineArgumentMissingError;
 use crate::errors::dapi_response_error::DapiResponseError;
 use crate::errors::identity_not_found_error::{IdentityNotFoundError};
@@ -8,10 +9,12 @@ pub mod cli_argument_missing_error;
 pub mod identity_not_found_error;
 pub mod dapi_response_error;
 pub mod identity_public_key_hash_mismatch_error;
+pub mod cli_argument_invalid_input;
 
 
 pub enum Error {
     CommandLineArgumentMissingError(CommandLineArgumentMissingError),
+    CommandLineArgumentInvalidInput(CommandLineArgumentInvalidInput),
     IdentityNotFoundError(IdentityNotFoundError),
     IdentityPublicKeyHashMismatchError(IdentityPublicKeyHashMismatchError),
     DapiResponseError(DapiResponseError),
@@ -24,15 +27,6 @@ impl Display for Error {
                 write!(f, "{}", err)
             }
             Error::IdentityNotFoundError(err) => {
-                // match err{
-                //     IdentifierOrPublicKey::Identifier(identifier) => {
-                //         write!(f, "Identity with identifier {} not found.", identifier.to_string(Base58))
-                //     }
-                //     IdentifierOrPublicKey::PublicKey(pkh) => {
-                //         write!(f, "Identity by public key hash {} not found.", pkh)
-                //     }
-                // }
-                //
                 write!(f, "{}", err)
             }
             Error::DapiResponseError(err) => {
@@ -41,6 +35,9 @@ impl Display for Error {
             Error::IdentityPublicKeyHashMismatchError(err) => {
                 write!(f, "{}", err)
             }
+            Error::CommandLineArgumentInvalidInput(err) => {
+                write!(f, "{}", err)
+            }
         }
     }
 }

diff --git a/src/utils.rs b/src/utils.rs
@@ -1,6 +1,11 @@
 use anyhow::Context;
+use base64::Engine;
+use base64::engine::general_purpose;
+use dpp::dashcore::{Network, PrivateKey};
 use dpp::util::entropy_generator::EntropyGenerator;
 use getrandom::getrandom;
+use crate::errors::cli_argument_invalid_input::CommandLineArgumentInvalidInput;
+use crate::errors::Error;
 
 
 pub struct MyDefaultEntropyGenerator;
@@ -11,4 +16,33 @@ impl EntropyGenerator for MyDefaultEntropyGenerator {
         getrandom(&mut buffer).context("generating entropy failed").unwrap();
         Ok(buffer)
     }
+}
+
+pub struct Utils;
+
+impl Utils {
+    pub fn decode_private_key_from_input_string(input: &str, network: Network) -> Result<PrivateKey, Error> {
+        let trimmed_input = input.replace("\n", "");
+
+        let base58: Vec<u8> = match PrivateKey::from_wif(&trimmed_input) {
+            Ok(private_key) => private_key.to_bytes(),
+            Err(_) => Vec::from([])
+        };
+        let hex: Vec<u8> = hex::decode(&trimmed_input).unwrap_or(Vec::from([]));
+        let base64: Vec<u8> = general_purpose::STANDARD.decode(hex::decode(&trimmed_input).unwrap_or(Vec::from([]))).unwrap_or(Vec::from([]));
+
+        let private_key: PrivateKey = {
+            if base58.len() > 0 {
+                PrivateKey::from_wif(&trimmed_input).expect("Unexpected error, could not construct private key from hex after validation")
+            } else if hex.len() > 0 {
+                PrivateKey::from_slice(hex.as_slice(), network).expect("Unexpected error, could not construct private key from hex after validation")
+            } else if base64.len() > 0 {
+                PrivateKey::from_slice(base64.as_slice(), network).expect("Unexpected error, could not construct private key from base64 after validation")
+            } else {
+                return Err(Error::CommandLineArgumentInvalidInput(CommandLineArgumentInvalidInput::from("Could not decode private key type from file (should be in WIF or hex)")))
+            }
+        };
+
+        Ok(private_key)
+    }
 }