Merge pull request #99 from opentensor/fix/roman/nuul-string-problem-with-password-in-env

basfroman · web-flow · commit bc0ac5f4fec6 · 2025-02-14T10:50:11.000-08:00
Fix `pyo3_runtime.PanicException` when encrypted password has `NUL byte`
diff --git a/src/errors.rs b/src/errors.rs
@@ -34,6 +34,10 @@ pub enum KeyFileError {
     EnvVarError(String),
     #[error("Password error: {0}")]
     PasswordError(String),
+    #[error("Base64 decoding error: {0}")]
+    Base64DecodeError(String),
+    #[error("Base64 encoding error: {0}")]
+    Base64EncodeError(String),
     #[error("Generic error: {0}")]
     Generic(String),
 }
diff --git a/src/keyfile.rs b/src/keyfile.rs
@@ -9,6 +9,7 @@ use std::str::from_utf8;
 use ansible_vault::{decrypt_vault, encrypt_vault};
 use fernet::Fernet;
 
+use base64::{engine::general_purpose, Engine as _};
 use passwords::analyzer;
 use passwords::scorer;
 use pyo3::pyfunction;
@@ -227,8 +228,11 @@ pub fn legacy_encrypt_keyfile_data(
 /// Retrieves the cold key password from the environment variables.
 pub fn get_password_from_environment(env_var_name: String) -> Result<Option<String>, KeyFileError> {
     match env::var(&env_var_name) {
-        Ok(encrypted_password) => {
-            let decrypted_password = decrypt_password(encrypted_password, env_var_name);
+        Ok(encrypted_password_base64) => {
+            let encrypted_password = general_purpose::STANDARD
+                .decode(&encrypted_password_base64)
+                .map_err(|_| KeyFileError::Base64DecodeError("Invalid Base64".to_string()))?;
+            let decrypted_password = decrypt_password(&encrypted_password, &env_var_name);
             Ok(Some(decrypted_password))
         }
         Err(_) => Ok(None),
@@ -373,23 +377,25 @@ fn expand_tilde(path: &str) -> String {
 }
 
 // Encryption password
-fn encrypt_password(key: String, value: String) -> String {
-    let mut encrypted = String::new();
-    for (i, c) in value.chars().enumerate() {
-        let encrypted_char = (c as u8) ^ (key.chars().nth(i % key.len()).unwrap() as u8);
-        encrypted.push(encrypted_char as char);
-    }
-    encrypted
+fn encrypt_password(key: &str, value: &str) -> Vec<u8> {
+    let key_bytes = key.as_bytes();
+    value
+        .as_bytes()
+        .iter()
+        .enumerate()
+        .map(|(i, &c)| c ^ key_bytes[i % key_bytes.len()])
+        .collect()
 }
 
 // Decrypting password
-fn decrypt_password(data: String, key: String) -> String {
-    let mut decrypted = String::new();
-    for (i, c) in data.chars().enumerate() {
-        let decrypted_char = (c as u8) ^ (key.chars().nth(i % key.len()).unwrap() as u8);
-        decrypted.push(decrypted_char as char);
-    }
-    decrypted
+fn decrypt_password(data: &[u8], key: &str) -> String {
+    let key_bytes = key.as_bytes();
+    let decrypted_bytes: Vec<u8> = data
+        .iter()
+        .enumerate()
+        .map(|(i, &c)| c ^ key_bytes[i % key_bytes.len()])
+        .collect();
+    String::from_utf8(decrypted_bytes).unwrap_or_else(|_| String::new())
 }
 
 #[derive(Clone)]
@@ -901,28 +907,13 @@ impl Keyfile {
             },
         };
         // saving password
-        match self.env_var_name() {
-            Ok(env_var_name) => {
-                // encrypt password
-                let encrypted_password = encrypt_password(self.env_var_name()?, password);
-                // store encrypted password
-                env::set_var(&env_var_name, &encrypted_password);
-
-                let message = format!(
-                    "The password has been saved to environment variable '{}'.\n",
-                    env_var_name
-                );
-                utils::print(message);
-                Ok(encrypted_password)
-            }
-            Err(e) => {
-                utils::print(format!(
-                    "Error saving environment variable name: {:?}.\n",
-                    e
-                ));
-                Ok("".to_string())
-            }
-        }
+        let env_var_name = self.env_var_name()?;
+        // encrypt password
+        let encrypted_password = encrypt_password(&env_var_name, &password);
+        let encrypted_password_base64 = general_purpose::STANDARD.encode(&encrypted_password);
+        // store encrypted password
+        env::set_var(&env_var_name, &encrypted_password_base64);
+        Ok(encrypted_password_base64)
     }
 
     /// Removes the password associated with the Keyfile from the local environment.
diff --git a/src/python_bindings.rs b/src/python_bindings.rs
@@ -118,6 +118,12 @@ impl PyKeyfile {
             .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
     }
 
+    fn env_var_name(&self) -> PyResult<String> {
+        self.inner
+            .env_var_name()
+            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
+    }
+
     #[pyo3(signature = (password=None))]
     fn save_password_to_env(&self, password: Option<String>) -> PyResult<String> {
         self.inner
@@ -846,6 +852,54 @@ except argparse.ArgumentError:
         Ok(Wallet { inner: result })
     }
 
+    /// Checks for existing coldkeypub and hotkeys, and creates them if non-existent.
+    ///     Arguments:
+    ///         coldkey_use_password (bool): Whether to use a password for coldkey. Defaults to ``True``.
+    ///         hotkey_use_password (bool): Whether to use a password for hotkey. Defaults to ``False``.
+    ///         save_coldkey_to_env (bool): Whether to save a coldkey password to local env. Defaults to ``False``.
+    ///         save_hotkey_to_env (bool): Whether to save a hotkey password to local env. Defaults to ``False``.
+    ///         coldkey_password (Optional[str]): Coldkey password for encryption. Defaults to ``None``. If `coldkey_password` is passed, then `coldkey_use_password` is automatically ``True``.
+    ///         hotkey_password (Optional[str]): Hotkey password for encryption. Defaults to ``None``. If `hotkey_password` is passed, then `hotkey_use_password` is automatically ``True``.
+    ///         overwrite (bool): Whether to overwrite an existing keys. Defaults to ``False``.
+    ///         suppress (bool): If ``True``, suppresses the display of the keys mnemonic message. Defaults to ``False``.
+    ///
+    ///     Returns:
+    ///         Wallet instance with created keys.
+
+    #[pyo3(signature = (coldkey_use_password=true, hotkey_use_password=false, save_coldkey_to_env=false, save_hotkey_to_env=false, coldkey_password=None, hotkey_password=None, overwrite=false, suppress=false))]
+    pub fn create(
+        &mut self,
+        coldkey_use_password: Option<bool>,
+        hotkey_use_password: Option<bool>,
+        save_coldkey_to_env: Option<bool>,
+        save_hotkey_to_env: Option<bool>,
+        coldkey_password: Option<String>,
+        hotkey_password: Option<String>,
+        overwrite: Option<bool>,
+        suppress: Option<bool>,
+    ) -> PyResult<Self> {
+        let result = self
+            .inner
+            .create(
+                coldkey_use_password.unwrap_or(true),
+                hotkey_use_password.unwrap_or(false),
+                save_coldkey_to_env.unwrap_or(false),
+                save_hotkey_to_env.unwrap_or(false),
+                coldkey_password,
+                hotkey_password,
+                overwrite.unwrap_or(false),
+                suppress.unwrap_or(false),
+            )
+            .map_err(|e| match e {
+                WalletError::InvalidInput(_) | WalletError::KeyGeneration(_) => {
+                    PyErr::new::<PyValueError, _>(e.to_string())
+                }
+                _ => PyErr::new::<PyKeyFileError, _>(format!("Failed to create wallet: {:?}", e)),
+            })?;
+
+        Ok(Wallet { inner: result })
+    }
+
     #[pyo3(
         signature = (coldkey_use_password=Some(true), hotkey_use_password=Some(false), save_coldkey_to_env=Some(false), save_hotkey_to_env=Some(false), coldkey_password=None, hotkey_password=None, overwrite=Some(false), suppress=Some(false))
     )]
diff --git a/tests/test_keyfile.py b/tests/test_keyfile.py
@@ -28,6 +28,7 @@
 from bittensor_wallet.keyfile import Keyfile
 from bittensor_wallet.keyfile import get_coldkey_password_from_environment
 from bittensor_wallet.keypair import Keypair
+from bittensor_wallet import Wallet
 
 
 def test_generate_mnemonic():
@@ -62,8 +63,8 @@ def test_only_provide_ss58_address():
     keypair = Keypair(ss58_address="16ADqpMa4yzfmWs3nuTSMhfZ2ckeGtvqhPWCNqECEGDcGgU2")
 
     assert (
-        f"0x{keypair.public_key.hex()}"
-        == "0xe4359ad3e2716c539a1d663ebd0a51bdc5c98a12e663bb4c4402db47828c9446"
+            f"0x{keypair.public_key.hex()}"
+            == "0xe4359ad3e2716c539a1d663ebd0a51bdc5c98a12e663bb4c4402db47828c9446"
     )
 
 
@@ -197,8 +198,8 @@ def test_create_keypair_from_private_key():
         private_key="0x1f1995bdf3a17b60626a26cfe6f564b337d46056b7a1281b64c649d592ccda0a9cffd34d9fb01cae1fba61aeed184c817442a2186d5172416729a4b54dd4b84e",
     )
     assert (
-        f"0x{keypair.public_key.hex()}"
-        == "0xe4359ad3e2716c539a1d663ebd0a51bdc5c98a12e663bb4c4402db47828c9446"
+            f"0x{keypair.public_key.hex()}"
+            == "0xe4359ad3e2716c539a1d663ebd0a51bdc5c98a12e663bb4c4402db47828c9446"
     )
 
 
@@ -326,24 +327,24 @@ def test_create(keyfile_setup_teardown):
     str(keyfile)
 
     assert (
-        keyfile.get_keypair(password="thisisafakepassword").ss58_address
-        == alice.ss58_address
+            keyfile.get_keypair(password="thisisafakepassword").ss58_address
+            == alice.ss58_address
     )
     assert (
-        keyfile.get_keypair(password="thisisafakepassword").public_key
-        == alice.public_key
+            keyfile.get_keypair(password="thisisafakepassword").public_key
+            == alice.public_key
     )
 
     bob = Keypair.create_from_uri("/Bob")
     keyfile.set_keypair(
         bob, encrypt=True, overwrite=True, password="thisisafakepassword"
     )
     assert (
-        keyfile.get_keypair(password="thisisafakepassword").ss58_address
-        == bob.ss58_address
+            keyfile.get_keypair(password="thisisafakepassword").ss58_address
+            == bob.ss58_address
     )
     assert (
-        keyfile.get_keypair(password="thisisafakepassword").public_key == bob.public_key
+            keyfile.get_keypair(password="thisisafakepassword").public_key == bob.public_key
     )
 
     repr(keyfile)
@@ -457,18 +458,36 @@ def test_deserialize_keypair_from_keyfile_data(keyfile_setup_teardown):
 
 
 @pytest.mark.parametrize(
-    "env_name,encrypted,decrypted",
+    "encrypted,decrypted",
     [
-        ("BT_PW_COLD_WALLET", "61,$>18", "testin{"),
-        ("BT_PW_COLD_WALLET", " =+$21,:!t``", "bittenoum0?7"),
+        ("c2ZsZGJpaG1q", "123456789"),
+        ("ID0rJDIxLDoh", "bittensor"),
+        ("NjEsJD4xOA==", "testing"),
     ],
 )
 def test_get_coldkey_password_from_environment(
-    monkeypatch, env_name, encrypted, decrypted
+        tmp_path, encrypted, decrypted
 ):
     # Preps
-    monkeypatch.setenv(env_name, encrypted)
+    assert tmp_path.exists()
+    assert tmp_path.is_dir()
+
+    wallet_name = "test_wallet"
+
+    wallet = Wallet(name=wallet_name, path=str(tmp_path))
+    wallet.create(
+        coldkey_use_password=True,
+        hotkey_use_password=False,
+        save_coldkey_to_env=True,
+        save_hotkey_to_env=False,
+        coldkey_password=decrypted,
+        overwrite=True,
+        suppress=True
+    )
+
+    # Call
+    wallet.coldkey_file.save_password_to_env(decrypted)
 
     # Calls + Assertions
-    assert get_coldkey_password_from_environment(env_name) == decrypted
+    assert get_coldkey_password_from_environment(wallet.coldkey_file.env_var_name()) == decrypted
     assert get_coldkey_password_from_environment("non_existent_env_variable") is None
