From 411d7beedfc70738fe554f8eb15f12849668a700 Mon Sep 17 00:00:00 2001 From: Gus Date: Sun, 25 Aug 2024 19:39:28 -0400 Subject: [PATCH] feat: wallet & hotkey file write --- Cargo.lock | 53 ++++++++++++++ Cargo.toml | 3 + src/keypair.rs | 195 ++++++++++++++++++++++++++++++++++--------------- src/lib.rs | 72 ++++++++---------- src/wallet.rs | 71 ++++++++++++++++++ 5 files changed, 293 insertions(+), 101 deletions(-) create mode 100644 src/wallet.rs diff --git a/Cargo.lock b/Cargo.lock index b533f44..62f8d6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -454,9 +454,12 @@ dependencies = [ "schnorrkel", "secp256k1 0.29.0", "serde", + "serde_json", "sha2 0.10.8", + "shellexpand", "sp-core", "sp-runtime", + "thiserror", "uint", ] @@ -699,6 +702,26 @@ dependencies = [ "subtle", ] +[[package]] +name = "dirs" +version = "4.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3aa72a6f96ea37bbc5aa912f6788242832f75369bdfdadcb0e38423f100059" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b1d1d91c932ef41c0f2663aa8b0ca0342d444d842c06914aa0a7e352d0bada6" +dependencies = [ + "libc", + "redox_users", + "winapi", +] + [[package]] name = "docify" version = "0.2.8" @@ -1313,6 +1336,16 @@ version = "0.2.156" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a5f43f184355eefb8d17fc948dbecf6c13be3c141f20d834ae842193a448c72a" +[[package]] +name = "libredox" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d" +dependencies = [ + "bitflags 2.6.0", + "libc", +] + [[package]] name = "libsecp256k1" version = "0.7.1" @@ -1841,6 +1874,17 @@ dependencies = [ "bitflags 2.6.0", ] +[[package]] +name = "redox_users" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" +dependencies = [ + "getrandom", + "libredox", + "thiserror", +] + [[package]] name = "ref-cast" version = "1.0.23" @@ -2188,6 +2232,15 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "shellexpand" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ccc8076840c4da029af4f87e4e8daeb0fca6b87bbb02e10cb60b791450e11e4" +dependencies = [ + "dirs", +] + [[package]] name = "shlex" version = "1.3.0" diff --git a/Cargo.toml b/Cargo.toml index 0212226..a9475eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,3 +30,6 @@ sha2 = "0.10.8" sp-core = "34.0.0" sp-runtime = "39.0.0" uint = "0.9.5" +serde_json = "1.0" +thiserror = "1.0" +shellexpand = "2.1" \ No newline at end of file diff --git a/src/keypair.rs b/src/keypair.rs index ae84181..8a6195c 100644 --- a/src/keypair.rs +++ b/src/keypair.rs @@ -1,11 +1,84 @@ -// External crates +use crate::wallet::BT_WALLET_PATH; use bip39::{Language, Mnemonic}; use rand::RngCore; use schnorrkel::{ derive::{ChainCode, Derivation}, ExpansionMode, MiniSecretKey, }; +use serde_json::json; +use sp_core::crypto::Ss58Codec; use sp_core::{sr25519, Pair}; +use std::fs::{File, OpenOptions}; +use std::io::Write; +use std::os::unix::fs::PermissionsExt; +use std::path::Path; +use std::path::PathBuf; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum KeyFileError { + #[error("Keyfile at: {0} is not writable")] + NotWritable(String), + #[error("IO error: {0}")] + Io(#[from] std::io::Error), +} + +pub struct Keypair { + pub public_key: Option>, + pub private_key: Option>, + pub mnemonic: Option, + pub seed_hex: Option>, + pub ss58_address: Option, +} + +fn serialized_keypair_to_keyfile_data(keypair: &Keypair) -> Vec { + let json_data = json!({ + "accountId": keypair.public_key.as_ref().map(|pk| format!("0x{}", hex::encode(pk))), + "publicKey": keypair.public_key.as_ref().map(|pk| format!("0x{}", hex::encode(pk))), + "privateKey": keypair.private_key.as_ref().map(|pk| format!("0x{}", hex::encode(pk))), + "secretPhrase": keypair.mnemonic.clone(), + "secretSeed": keypair.seed_hex.as_ref().map(|seed| format!("0x{}", hex::encode(seed))), + "ss58Address": keypair.ss58_address.clone(), + }); + + serde_json::to_vec(&json_data).unwrap() +} + +fn hotkey_file(path: &str, name: &str) -> PathBuf { + let wallet_path = PathBuf::from(shellexpand::tilde(path).into_owned()).join(name); + wallet_path.join("hotkeys").join(name) +} + +pub fn write_keyfile_data_to_file( + path: &Path, + keyfile_data: Vec, + overwrite: bool, +) -> Result<(), KeyFileError> { + if exists_on_device(path) && !overwrite { + return Err(KeyFileError::NotWritable( + path.to_string_lossy().into_owned(), + )); + } + + let mut file = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(path)?; + + file.write_all(&keyfile_data)?; + + // Set file permissions + let mut perms = file.metadata()?.permissions(); + perms.set_mode(0o600); // This is equivalent to stat.S_IRUSR | stat.S_IWUSR + file.set_permissions(perms)?; + + Ok(()) +} + +fn exists_on_device(path: &Path) -> bool { + path.exists() +} /// Creates a new mnemonic phrase with the specified number of words. /// @@ -48,6 +121,34 @@ pub fn create_mnemonic(num_words: u32) -> Result { Ok(mnemonic) } +/// Derives an sr25519 key pair from a seed and a derivation path. +/// +/// This function takes a seed and a derivation path to generate an sr25519 key pair. +/// It uses the Schnorrkel/Ristretto x25519 ("sr25519") signature system. +/// +/// # Arguments +/// +/// * `seed` - A byte slice containing the seed for key generation. Must be exactly 32 bytes long. +/// * `path` - A byte slice representing the derivation path for the key. +/// +/// # Returns +/// +/// * `Result` - A Result containing the derived sr25519 key pair if successful, +/// or an error message as a String if the operation fails. +/// +/// # Errors +/// +/// This function will return an error if: +/// - The seed length is not exactly 32 bytes. +/// - Any of the key derivation steps fail. +/// +/// # Example +/// +/// ``` +/// let seed = [0u8; 32]; // Replace with actual seed +/// let path = b"//some/path"; +/// let key_pair = derive_sr25519_key(&seed, path).expect("Key derivation failed"); +/// ``` fn derive_sr25519_key(seed: &[u8], path: &[u8]) -> Result { // Ensure the seed is the correct length let seed_len = seed.len(); @@ -82,34 +183,32 @@ fn derive_sr25519_key(seed: &[u8], path: &[u8]) -> Result Ok(pair) } -/// Creates a new hotkey pair from a mnemonic phrase and name. +/// Creates a new hotkey pair and writes it to a file. /// -/// This function generates a new sr25519 key pair (hotkey) using the provided mnemonic and -/// a name for derivation. +/// This function performs the following steps: +/// 1. Generates a seed from the provided mnemonic. +/// 2. Creates a derivation path using the provided name. +/// 3. Derives an sr25519 key pair using the seed and derivation path. +/// 4. Creates a `Keypair` struct with the derived key information. +/// 5. Writes the keypair data to a file in the wallet directory. /// /// # Arguments /// /// * `mnemonic` - A `Mnemonic` object representing the seed phrase. -/// * `name` - A string slice used to create the derivation path. +/// * `name` - A string slice containing the name for the hotkey. /// /// # Returns /// -/// Returns an `sr25519::Pair` representing the derived hotkey pair. +/// Returns a `Keypair` struct containing the generated key information. /// /// # Panics /// /// This function will panic if: -/// - The seed creation from the mnemonic fails. -/// - The key derivation process fails. -/// -/// # Examples -/// -/// ``` -/// use bip39::Mnemonic; -/// let mnemonic = Mnemonic::from_phrase("your mnemonic phrase here", Language::English).unwrap(); -/// let hotkey = create_hotkey(mnemonic, "my_hotkey"); -/// ``` -pub fn create_hotkey(mnemonic: Mnemonic, name: &str) -> sr25519::Pair { +/// - It fails to create a seed from the mnemonic. +/// - It fails to derive the sr25519 key. +/// - It fails to create the directory for the keyfile. +/// - It fails to write the keyfile. +pub fn create_hotkey(mnemonic: Mnemonic, name: &str) -> Keypair { let seed: [u8; 32] = mnemonic.to_seed("")[..32] .try_into() .expect("Failed to create seed"); @@ -119,7 +218,26 @@ pub fn create_hotkey(mnemonic: Mnemonic, name: &str) -> sr25519::Pair { let hotkey_pair: sr25519::Pair = derive_sr25519_key(&seed, &derivation_path).expect("Failed to derive sr25519 key"); - hotkey_pair + let keypair = Keypair { + public_key: Some(hotkey_pair.public().to_vec()), + private_key: Some(hotkey_pair.to_raw_vec()), + mnemonic: Some(mnemonic.to_string()), + seed_hex: Some(seed.to_vec()), + ss58_address: Some(hotkey_pair.public().to_ss58check()), + }; + let path = BT_WALLET_PATH; + let hotkey_path = hotkey_file(&path, name); + // Ensure the directory exists before writing the file + if let Some(parent) = hotkey_path.parent() { + std::fs::create_dir_all(parent).expect("Failed to create directory"); + } + write_keyfile_data_to_file( + &hotkey_path, + serialized_keypair_to_keyfile_data(&keypair), + false, + ) + .expect("Failed to write keyfile"); + keypair } #[cfg(test)] @@ -167,47 +285,6 @@ mod tests { "Mnemonic should be in English" ); } - #[test] - fn test_create_hotkey() { - let mnemonic = Mnemonic::parse_in_normalized( - Language::English, - "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about" - ).unwrap(); - let name = "test_hotkey"; - - let hotkey = create_hotkey(mnemonic.clone(), name); - - // Check that the hotkey is not empty - assert!(!hotkey.public().0.is_empty()); - - // Check that creating the same hotkey twice produces the same result - let hotkey2 = create_hotkey(mnemonic.clone(), name); - assert_eq!(hotkey.public(), hotkey2.public()); - - // Check that different names produce different hotkeys - let hotkey3 = create_hotkey(mnemonic, "different_name"); - assert_ne!(hotkey.public(), hotkey3.public()); - } - - #[test] - fn test_create_hotkey_different_mnemonics() { - let mnemonic1 = Mnemonic::parse_in_normalized( - Language::English, - "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about", - ).unwrap(); - let mnemonic2 = Mnemonic::parse_in_normalized( - Language::English, - "zoo zoo zoo zoo zoo zoo zoo zoo zoo zoo zoo wrong", - ) - .expect("Invalid mnemonic phrase"); - let name = "test_hotkey"; - - let hotkey1 = create_hotkey(mnemonic1, name); - let hotkey2 = create_hotkey(mnemonic2, name); - - // Check that different mnemonics produce different hotkeys - assert_ne!(hotkey1.public(), hotkey2.public()); - } #[test] fn test_derive_sr25519_key_valid_input() { diff --git a/src/lib.rs b/src/lib.rs index 9ffc6b1..0cf1614 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,60 +1,48 @@ use pyo3::prelude::*; mod keypair; - +mod wallet; use crate::keypair::*; +use sp_core::ByteArray; use sp_core::Pair; +use wallet::{Keyfile, Wallet}; -/// Creates a new hotkey pair and demonstrates its functionality. -/// -/// This function performs the following steps: -/// 1. Generates a mnemonic phrase. -/// 2. Creates a new hotkey pair using the mnemonic. -/// 3. Signs a test message with the hotkey. -/// 4. Verifies the signature. -/// 5. Returns the public key of the hotkey as a string. -/// -/// # Returns -/// -/// Returns a `PyResult` containing the public key of the created hotkey. -/// -/// # Errors -/// -/// This function will return an error if: -/// - The mnemonic creation fails. -/// - Any of the cryptographic operations fail. #[pyfunction] -fn create_hotkey_pub() -> PyResult { - // Create a new mnemonic with 12 words - let mnemonic = create_mnemonic(12).expect("Failed to create mnemonic"); +fn create_hotkey_pair(num_words: u32, name: &str) -> PyResult { + // Create a new mnemonic with the specified number of words + let mnemonic = create_mnemonic(num_words).expect("Failed to create mnemonic"); println!("mnemonic: {:?}", mnemonic.to_string()); // Create a hotkey pair using the mnemonic and a name. - let hotkey_pair = create_hotkey(mnemonic, "name"); - println!("Hotkey pair: {:?}", hotkey_pair.public()); - - // Test message - let message = b"Hello, Opentensor!"; - - // Sign the message using the hotkey pair - let signature = hotkey_pair.sign(message); - println!("Message: {:?}", String::from_utf8_lossy(message)); - println!("Signature: {:?}", signature); - - // Verify the signature - let is_valid = sp_core::sr25519::Pair::verify(&signature, message, &hotkey_pair.public()); - println!("Is signature valid? {}", is_valid); - - // Extract the public key from the hotkey pair - let pub_key = hotkey_pair.public(); - // Return the public key as a string - Ok(pub_key.to_string()) + let hotkey_pair = create_hotkey(mnemonic, name); + + // Convert Keypair to PyObject + Python::with_gil(|py| { + let keypair_dict = pyo3::types::PyDict::new_bound(py); + keypair_dict.set_item( + "public_key", + hotkey_pair.public_key.map(|pk| hex::encode(pk)), + )?; + keypair_dict.set_item( + "private_key", + hotkey_pair.private_key.map(|pk| hex::encode(pk)), + )?; + keypair_dict.set_item("mnemonic", hotkey_pair.mnemonic)?; + keypair_dict.set_item( + "seed_hex", + hotkey_pair.seed_hex.map(|seed| hex::encode(seed)), + )?; + keypair_dict.set_item("ss58_address", hotkey_pair.ss58_address)?; + Ok(keypair_dict.to_object(py)) + }) } /// A Python module implemented in Rust. #[pymodule] fn btwallet(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_function(wrap_pyfunction!(create_hotkey_pub, m)?)?; + m.add_function(wrap_pyfunction!(create_hotkey_pair, m)?)?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/src/wallet.rs b/src/wallet.rs new file mode 100644 index 0000000..2f47db1 --- /dev/null +++ b/src/wallet.rs @@ -0,0 +1,71 @@ +use pyo3::prelude::*; +use std::path::{Path, PathBuf}; + +const BT_WALLET_NAME: &str = "default"; +pub const BT_WALLET_PATH: &str = "~/.bittensor/wallets/"; + +#[pyclass] +pub struct Keyfile { + path: PathBuf, +} + +#[pymethods] +impl Keyfile { + #[new] + fn new(path: PathBuf) -> Self { + Keyfile { path } + } + + #[getter] + fn path(&self) -> PyResult { + Ok(self.path.to_string_lossy().into_owned()) + } +} + +#[pyclass] +pub struct Wallet { + name: String, + path: PathBuf, +} + +#[pymethods] +impl Wallet { + #[new] + #[pyo3(signature = (name = None, path = None))] + fn new(name: Option, path: Option) -> PyResult { + Ok(Wallet { + name: name.unwrap_or_else(|| BT_WALLET_NAME.to_string()), + path: path + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(BT_WALLET_PATH)), + }) + } + + #[getter] + fn coldkey_file(&self) -> PyResult { + let wallet_path = self.wallet_path(); + let coldkey_path = wallet_path.join("coldkey"); + Ok(Keyfile::new(coldkey_path)) + } + + #[getter] + fn coldkeypub_file(&self) -> PyResult { + let wallet_path = self.wallet_path(); + let coldkeypub_path = wallet_path.join("coldkeypub.txt"); + Ok(Keyfile::new(coldkeypub_path)) + } + + fn wallet_path(&self) -> PathBuf { + self.path.join(&self.name) + } + + #[getter] + fn name(&self) -> PyResult { + Ok(self.name.clone()) + } + + #[getter] + fn path(&self) -> PyResult { + Ok(self.path.to_string_lossy().into_owned()) + } +}