From a627cb77eb4b1888e6eb3a8349d378b18a9261cf Mon Sep 17 00:00:00 2001 From: Roman Date: Tue, 15 Oct 2024 10:16:00 -0700 Subject: [PATCH 1/3] add PasswordError and bound to python --- src/errors.rs | 34 +++++++++++++++++++++++++++++++++- src/keyfile.rs | 12 ++++++------ src/lib.rs | 1 + 3 files changed, 40 insertions(+), 7 deletions(-) diff --git a/src/errors.rs b/src/errors.rs index 0ee77e8..3e95fa3 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -2,13 +2,14 @@ use pyo3::exceptions::PyException; use pyo3::prelude::*; use std::{error, fmt}; +// KeyFileError #[pyclass(extends=PyException)] #[derive(Debug)] pub struct KeyFileError { pub message: String, } -/// Error thrown when the keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid. +/// Error thrown when the keyfile is corrupt, non-writable, non-readable. #[pymethods] impl KeyFileError { #[new] @@ -31,6 +32,7 @@ impl fmt::Display for KeyFileError { impl error::Error for KeyFileError {} +// ConfigurationError #[pyclass(extends=PyException)] #[derive(Debug)] pub struct ConfigurationError { @@ -59,3 +61,33 @@ impl fmt::Display for ConfigurationError { } impl error::Error for ConfigurationError {} + +// PasswordError +#[pyclass(extends=PyException)] +#[derive(Debug)] +pub struct PasswordError { + pub message: String, +} + +/// PasswordError occurs if the password used for decryption is invalid. +#[pymethods] +impl PasswordError { + #[new] + #[pyo3(signature = (message=None))] + pub fn new(message: Option) -> Self { + let msg = message.unwrap_or_default(); + PasswordError { message: msg } + } + + pub fn __str__(&self) -> String { + self.message.clone() + } +} + +impl fmt::Display for PasswordError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PasswordError: {}", self.message) + } +} + +impl error::Error for PasswordError {} diff --git a/src/keyfile.rs b/src/keyfile.rs index 90a7a27..38cdc15 100644 --- a/src/keyfile.rs +++ b/src/keyfile.rs @@ -19,7 +19,7 @@ use passwords::analyzer; use passwords::scorer; use serde_json::json; -use crate::errors::KeyFileError; +use crate::errors::{KeyFileError, PasswordError}; use crate::keypair::Keypair; use crate::utils; @@ -414,7 +414,7 @@ pub fn decrypt_keyfile_data( .ok_or(PyErr::new::("Invalid nonce."))?; let ciphertext = &data[secretbox::NONCEBYTES..]; secretbox::open(ciphertext, &nonce, key) - .map_err(|_| PyErr::new::("Wrong password for nacl decryption.")) + .map_err(|_| PyErr::new::("Wrong password for nacl decryption.")) } // decrypt of keyfile_data with legacy way @@ -428,7 +428,7 @@ pub fn decrypt_keyfile_data( let keyfile_data_str = from_utf8(keyfile_data)?; fernet .decrypt(keyfile_data_str) - .map_err(|_| PyErr::new::("Wrong password for nacl decryption.")) + .map_err(|_| PyErr::new::("Wrong password for legacy decryption.")) } let mut password = password; @@ -453,21 +453,21 @@ pub fn decrypt_keyfile_data( if keyfile_data_is_encrypted_nacl(py, keyfile_data)? { let key = derive_key(password.as_bytes()); let decrypted_data = nacl_decrypt(keyfile_data, &key) - .map_err(|_| PyErr::new::("Wrong password for decryption."))?; + .map_err(|_| PyErr::new::("Wrong password for decryption."))?; return Ok(PyBytes::new_bound(py, &decrypted_data).into_py(py)); } // Ansible Vault decryption if keyfile_data_is_encrypted_ansible(py, keyfile_data)? { let decrypted_data = decrypt_vault(keyfile_data, password.as_str()) - .map_err(|_| PyErr::new::("Wrong password for decryption."))?; + .map_err(|_| PyErr::new::("Wrong password for decryption."))?; return Ok(PyBytes::new_bound(py, &decrypted_data).into_py(py)); } // Legacy decryption if keyfile_data_is_encrypted_legacy(py, keyfile_data)? { let decrypted_data = legacy_decrypt(&password, keyfile_data) - .map_err(|_| PyErr::new::("Wrong password for decryption."))?; + .map_err(|_| PyErr::new::("Wrong password for decryption."))?; return Ok(PyBytes::new_bound(py, &decrypted_data).into_py(py)); } diff --git a/src/lib.rs b/src/lib.rs index 98db472..c54af28 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,6 +36,7 @@ fn register_errors_module(main_module: &Bound<'_, PyModule>) -> PyResult<()> { let errors_module = PyModule::new_bound(main_module.py(), "errors")?; errors_module.add_class::()?; errors_module.add_class::()?; + errors_module.add_class::()?; main_module.add_submodule(&errors_module) } From 21b5a89d98057a3d7716e8d985b36d5498f07486 Mon Sep 17 00:00:00 2001 From: Roman Date: Tue, 15 Oct 2024 10:16:07 -0700 Subject: [PATCH 2/3] formatter --- src/wallet.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/wallet.rs b/src/wallet.rs index 69169bb..36d5470 100644 --- a/src/wallet.rs +++ b/src/wallet.rs @@ -34,7 +34,11 @@ pub fn display_mnemonic_msg(mnemonic: String, key_type: &str) { } // Function to safely retrieve attribute as Option from passed python object -fn get_attribute_string(py: Python, obj: &Bound, attr_name: &str) -> PyResult> { +fn get_attribute_string( + py: Python, + obj: &Bound, + attr_name: &str, +) -> PyResult> { match obj.getattr(attr_name) { Ok(attr) => { if attr.is_none() { @@ -82,9 +86,8 @@ impl Wallet { hotkey: Option, path: Option, config: Option, - py: Python + py: Python, ) -> PyResult { - // default config's values if config and config.wallet exist let mut conf_name: Option = None; let mut conf_hotkey: Option = None; @@ -128,7 +131,10 @@ impl Wallet { let final_path = if let Some(path) = path { path } else if let Some(conf_path) = conf_path { - conf_path.strip_prefix("~/").unwrap_or(&conf_path).to_string() + conf_path + .strip_prefix("~/") + .unwrap_or(&conf_path) + .to_string() } else { BT_WALLET_PATH.to_string() }; @@ -183,8 +189,8 @@ impl Wallet { env::var("BT_WALLET_NAME").unwrap_or_else(|_| BT_WALLET_NAME.to_string()); let default_hotkey = env::var("BT_WALLET_HOTKEY").unwrap_or_else(|_| BT_WALLET_HOTKEY.to_string()); - let default_path = - env::var("BT_WALLET_PATH").unwrap_or_else(|_| format!("~/{}", BT_WALLET_PATH.to_string())); + let default_path = env::var("BT_WALLET_PATH") + .unwrap_or_else(|_| format!("~/{}", BT_WALLET_PATH.to_string())); let prefix_str = if let Some(value) = prefix { format!("\"{}\"", value) From b61872e3a1bb8c131e6d467968fc8be6da3a88c7 Mon Sep 17 00:00:00 2001 From: Roman Date: Tue, 15 Oct 2024 10:20:15 -0700 Subject: [PATCH 3/3] add `from bittensor_wallet.errors import PasswordError` --- bittensor_wallet/errors/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bittensor_wallet/errors/__init__.py b/bittensor_wallet/errors/__init__.py index 9bd501f..2e86768 100644 --- a/bittensor_wallet/errors/__init__.py +++ b/bittensor_wallet/errors/__init__.py @@ -2,3 +2,4 @@ ConfigurationError = _.ConfigurationError KeyFileError = _.KeyFileError +PasswordError = _.PasswordError