Repository: opentensor/btwallet (language: python)

Issue:
Updates readme



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.MD b/CHANGELOG.MD
index 793d915..f87d439 100644
--- a/CHANGELOG.MD
+++ b/CHANGELOG.MD
@@ -1,5 +1,14 @@
 # Changelog
 
+## 2.0.2 /2024-10-10
+
+## What's Changed
+
+* Add the ability to use environment variables to store hot/coldkey passwords by @roman-opentensor in https://github.com/opentensor/btwallet/pull/46
+* fix/roman/fix-config-parsing by @roman-opentensor in https://github.com/opentensor/btwallet/pull/47
+
+**Full Changelog**: https://github.com/opentensor/btwallet/compare/v2.0.1...v2.0.2
+
 ## 2.0.1 /2024-10-04
 
 ## What's Changed
@@ -12,18 +21,225 @@
 
 ## What's Changed
 
-* Rust Wallet by @opendansor in https://github.com/opentensor/btwallet/pull/39
-* use pythons stdout by @camfairchild in https://github.com/opentensor/btwallet/pull/30
-* Cam/no excepts by @camfairchild in https://github.com/opentensor/btwallet/pull/32
-* Updated readme docs by @rajkaramchedu in https://github.com/opentensor/btwallet/pull/31
-* Feat/ci/add release by @camfairchild in https://github.com/opentensor/btwallet/pull/38
+# Revamped Bittensor Wallet written in Rust from the ground up
+
+* Updates Changelog by @ibraheem-opentensor in 5966932
+* Merge conflicts. by @opendansor in a5d0ebd
+* Bumps to 2.0.0 by @ibraheem-opentensor in fbc835e
+* Bumps to a5 by @ibraheem-opentensor in ff7ab41
+* Feat/ci/add release (#38) by @Cameron Fairchild in ffc995b
+* remove utils by @Cameron Fairchild in a4521a4
+* Merge pull request #31 from opentensor/rajk/update-README by @Roman in 7904b81
+* update tests by @Roman in 1ba42fc
+* Add re-trying to type the password after wallet creation by @Roman in 60ec77a
+* move `keypair::Keypair.private_key` outside the python scope by @Roman in 241fcd0
+* use pypi trusted release by @Cameron Fairchild in e5d4c49
+* chore: ruff by @Cameron Fairchild in 0c051db
+* Cam/no excepts (#32) by @Cameron Fairchild in 6685852
+* wallet.hotkey.ss58_address doesn't ask typing pass after waller.unlock_hotkey() by @Roman in 591889d
+* wallet.coldkey.ss58_address doesn't ask typing pass after waller.unlock_coldkey() by @Roman in fde4794
+* Fix typos in status messages by @opendansor in 6953184
+* delete double import by @Roman in a02ec3a
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in 1979480
+* fix panic error to KeyFileError by @Roman in 8388ddb
+* remove import by @Cameron Fairchild in f186944
+* wrap with err by @Cameron Fairchild in 246988b
+* Refactor error handling in keyfile.rs by @opendansor in b294357
+* update config::Config -> add getters by @Roman in f03db92
+* Add wallet::Wallet.add_args by @Roman in 192c447
+* try eval bound with import first by @Cameron Fairchild in 3f122b1
+* Adds extra line after print by @ibraheem-opentensor in 2cbf0cb
+* chore: clippy by @Cameron Fairchild in 0aa5877
+* add prompt and prompt password by @Cameron Fairchild in 19dfabf
+* Updated readme docs by @Watchmaker in 7bf89dd
+* Add getters for name and path in Keyfile by @opendansor in 8c31b65
+* use pythons stdout (#30) by @Cameron Fairchild in 9566fc2
+* Add getter methods for wallet name, path, and hotkey by @opendansor in e8eb8ff
+* add tests to readme by @Roman in 984a43f
+* add ScaleCodec to keypair::Keypair.sign and keypair::Keypair.verify by @Roman in de21752
+* fix create_from_seed (add union str, bytes) processing by @Roman in 29fc083
+* Merge remote-tracking branch 'origin/rust-fix-test' into rust-fix-test by @Roman in b8b23e0
+* unused imports by @Roman in 1b5c8c3
+* add version by @Cameron Fairchild in 4863a39
+* add back extension-module feature by @Cameron Fairchild in 25ae0cb
+* move assert into match arm by @Cameron Fairchild in b775256
+* Refactor keypair decoding and wallet initialization by @opendansor in 228886b
+* use new instead of init by @Cameron Fairchild in 9cee953
+* add *args, **kwargs by @Roman in 2659496
+* update for wallets properties by @Roman in f822741
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in f3fd9c4
+* use direct import for error types in utils.rs by @Roman in e8d1124
+* use direct import for error types in wallet.rs by @Roman in 19baa56
+* use direct import for error types in keypair.rs by @Roman in ded234c
+* update .gitignore by @Roman in 562bf4a
+* use direct import for errors types in keyfile.rs by @Roman in 665a2b2
+* Update version to 2.0.0 by @opendansor in 8ff86e5
+* add `rpassword::prompt_password` instead of input for asking password. by @Roman in dd3afcc
+* can't mock functions that are not implemented yet (rust doesn't accept this way). This limits testing from Python. temporarily commenting on mocking tests. by @Roman in 2726151
+* update types of errors in keypair by @Roman in 881790a
+* update types of errors in keypair by @Roman in e06cf62
+* update errors by @Roman in 80bb4f3
+* update ConfigurationError usage for Keypair::new by @Roman in 4f0280a
+* add bittensor_wallet.keypair submodule by @Roman in ed4d619
+* add error into the module by @Roman in dfe99b0
+* add pyo3 signature to  keypair::Keypair methods (functions) by @Roman in ff4b65e
+* improve keyfile::Keyfile.new logic by @Roman in c2d7692
+* add errors::ConfigurationError by @Roman in 490f980
+* fix classic import logic by @Roman in 199c748
+* add export for wallet::display_mnemonic_msg by @Roman in 62b108b
+* improve keyfile::Keyfile.__str__() by @Roman in c26397a
+* add #[pyclass(subclass)] ability for main classes, add mock subpackage by @Roman in 5590e00
+* update README (import logic) by @Roman in 31eb915
+* Fix keypair::Keypair.create_from_encrypted_json by @Roman in 1562c5d
+* lib.rs -> add classes on the package top level by @Roman in 1df5a6f
+* README -> add test `create_from_encrypted_json` by @Roman in 19bcd43
+* README, refactoring by @Roman in 25e9e14
+* README, refactoring by @Roman in 9cb9071
+* update python package structure by @Roman in 3e54d2b
+* add errors::KeyFileError, add its usage in the same way as in python wallet by @Roman in 58728b4
+* create correct python package structure by @Roman in af44b18
+* temporarily remove incorrect private key usage in `create_from_encrypted_json` by @Roman in a4c433e
+* keypair::Keypair.create_from_encrypted_json by @Roman in 1a36312
+* import cleanup by @Roman in f2137d5
+* deps cleanup by @Roman in 82bf3ca
+* cls annotations fix by @Roman in b3c6e24
+* add keypair::Keypair.create_from_encrypted_json by @Roman in e590f1f
+* fixes by @Roman in f964492
+* cleanup by @Roman in a54324d
+* add wallet::Wallet.regenerate_coldkeypub, coldkey_file, coldkeypub_file, set_coldkey, set_coldkeypub, get_coldkeypub, get_coldkey, get_hotkey, regenerate_coldkey by @Roman in 90d343c
+* fix keypair::Keypair.mnemonic() by @Roman in bdd30ba
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in b3f00e6
+* add wallet::Wallet.regenerate_hotkey, wallet::Wallet.set_hotkey, wallet::Wallet.hotkey_file, update constants wallet path by @Roman in 431b7f7
+* Add comprehensive tests for keyfile.rs by @opendansor in 99109e8
+* Refactor wallet initialization to use config defaults by @Roman in 5041c3c
+* fix cc compiler problem by @Roman in c085bce
+* add some empty functions into wallet::Wallet by @Roman in 4a3213a
+* add empty structure to wallet create, create some of the functions by @Roman in 67a1791
+* make keypair::Keypair.create_from_uri public by @Roman in 9e01b53
+* add list of functions into wallet.rs by @Roman in cbef1c1
+* fix keyfile.rs by @Roman in 64fa38f
+* update config by @Roman in a5c699f
+* add wallet::Wallet.create_if_non_existent by @Roman in 247573a
+* activate GPG key by @Roman in 235a0de
+* Refactor keyfile encryption error handling by @Roman in 51c2c11
+* Refactor `decrypt` to handle errors with `PyResult` by @opendansor in db1561b
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @opendansor in faf282f
+* Improve file encryption and decryption error handling by @opendansor in 1090dc1
+* add keypair::Keypair.check_and_update_encryption by @Roman in b8ea5a3
+* add keypair::Keypair.get_keypair by @Roman in 7b69e6b
+* commented keypair getter by @Roman in b9d7b43
+* add keyfile::Keyfile.is_encrypted by @Roman in 72c384d
+* move nacl_decrypt and legacy_decrypt into decrypt_keyfile_data by @Roman in 72f79c8
+* add keypair::Keypair.set_keypair, fix: keypair::encrypt_keyfile_data, add printing messages by @Roman in f96d2fd
+* Refactor `set_keypair` to handle encryption and file writing by @opendansor in 026702c
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in 0509f4f
+* add keypair::decrypt_keyfile_data by @Roman in 9217363
+* Update error handling and getter methods in keyfile.rs by @opendansor in 141b941
+* update docstring + add pyo3 signature by @Roman in 6a163df
+* update docstring by @Roman in 80a4dcf
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in 827bee1
+* add keyfile::encrypt_keyfile_data by @Roman in 86a43f7
+* Add legacy keyfile encryption function by @opendansor in ebd82b2
+* Pins py-bip39-bindings (#25) by @ibraheem-opentensor in 8491a77
+* add keyfile::keyfile_data_encryption_method by @Roman in a6f84aa
+* fix input logic and messages by @Roman in c7023bf
+* Enable encryption detection functions by @opendansor in 8745397
+* Adds changelog (#24) by @ibraheem-opentensor in 7ce313c
+* update README.md -> tests by @Roman in 2e59bf4
+* update README.md -> tests by @Roman in 120a456
+* remove unused prefixes and pub sign by @Roman in 9fc24ac
+* formatter by @Roman in e3c6eae
+* add pyo3 - signature: green build by @Roman in bdd1d02
+* Refactor keyfile functions to use PyBytes instead of Vec<u8> by @opendansor in fc97001
+* Deployment script + updates (#23) by @ibraheem-opentensor in aebb643
+* update deps in Cargo.toml by @Roman in b9d25e1
+* add TODOs by @Roman in 384345c
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in 0550d80
+* add Keyfile::is_encrypted by @Roman in 766fad9
+* Enable keyfile serialization and password validation by @opendansor in 8e8696a
+* add Keyfile::_read_keyfile_data_from_file, Keyfile::is_writable by @Roman in 0544c31
+* add Keyfile::is_readable by @Roman in 8a6fbd2
+* add Keyfile::exists_on_device update by @Roman in 98f45cd
+* add Keyfile::_write_keyfile_data_to_file by @Roman in 5b738f3
+* add Keyfile::_may_overwrite by @Roman in 2689b78
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in 57bc8be
+* add Keyfile::make_dirs by @Roman in 86964e4
+* Add keyfile-related functions to library by @opendansor in 127f3c6
+* add Keyfile::exists_on_device by @Roman in 53ae243
+* add Keyfile.rs structure by @Roman in cf8af98
+* add Keyfile.rs structure by @Roman in 456d2c2
+* cargo files update by @Roman in 7eff490
+* frm by @Roman in bb92e45
+* Add comments with keyfile.py usage across the wallet, btcli, sdk by @Roman in b6ecbae
+* cargo clippy GREEN by @Roman in 777039a
+* cargo fmt by @Roman in 11d5803
+* GREEEEN BUILD by @Roman in bb89aa4
+* update Cargo.toml by @Roman in 5d8e837
+* add funcs to the pymodule in the lib.rs by @Roman in fc8d30f
+* update README.md by @Roman in 9e4f300
+* utils done by @Roman in dbe7f66
+* update README.md by @Roman in 2a0d92d
+* keep support crypto_type 1 only (sr25519) by @Roman in 0fa29bc
+* add utils::get_ss58_format + test by @Roman in 520b1d9
+* fix utils::is_valid_ss58_address function + test by @Roman in d4c91ff
+* fix utils::is_valid_ss58_address function by @Roman in fe5bd41
+* Merge remote-tracking branch 'origin/main-rust' into main-rust by @Roman in cb458e6
+* add verify function + manual test case by @Roman in 67b03f1
+* BROKEN: Add SS58 address validation utility by @opendansor in 7075052
+* add sign function + manual test case by @Roman in a113cb6
+* add the checker `If public_key is missing (ss58_address wasn't created), return an error` by @Roman in 366da8e
+* add creation with ss58_address by @Roman in b35e4c5
+* remove comments from public_key getter by @Roman in 644bb04
+* update manual test script by @Roman in 46dde7d
+* add creating from public_key by @Roman in 78e1dd8
+* fix ci by @Roman in ab8b8c3
+* update README with a test for Keypair by @Roman in 7101925
+* add print for test by @Roman in 26821a2
+* Minimum functionality done by @Roman in 2f18751
+* add __repr__ and __str__ by @Roman in bed0f13
+* all fields in Keypair work by @Roman in a70ce44
+* add getters ss58_address and seed_hex by @Roman in 4829eae
+* add pair by @Roman in 0bfac83
+* impl Default for Keypair by @Roman in 5e0c146
+* empty struct Keypair and getters by @Roman in 6ee130a
+* convert private_key to bytes by @Roman in 0b93c3c
+* convert public_key to bytes by @Roman in 7ae4ad4
+* update structure by @Roman in e5174b1
+* delete pyproject.toml by @Roman in 628291f
+* fix Keypair `#pyclass` macros by @Roman in 6f9005c
+* update structure by @Roman in 7d8db72
+* working by @Roman in 14c8275
+* update python package folder name by @Roman in 07ce2d6
+* update structure by @Roman in 171fdd1
+* wip restructure by @Gus in 258e702
+* Test signed commit by @Gus in 91f2c2c
+* Merge pull request #21 from opentensor/fix/roman/cryptography-version-to-43.0.1 by @Roman in be61bd9
+* bump up cryptography version until 43.0.1 by @Roman in 5e3726d
+* Apply changes from #2237 (#18) by @Benjamin Himes in 83983f9
+* Merge pull request #20 from opentensor/feat/roman/bump-up-version by @Roman in 89ba61f
+* bump up the version by @Roman in 527a7aa
+* Merge pull request #19 from opentensor/feat/thewhaleking/raise-exception-on-unlock-failure by @Benjamin Himes in 8e724c1
+* Raises a KeyFileError when one is encountered when unlocking. by @Benjamin Himes in fdef4ba
+* encrypt using secret box access by @Gus in 2faf3a4
+* wallet integration with python wip by @Gus in 218ccad
+* Merge pull request #17 from opentensor/fix/update-mnemonic-command by @ibraheem-opentensor in 5fa01d9
+* Updates mnemonic command output by @ibraheem-opentensor in 7789a9a
+* clean up by @Gus in a3b58d5
+* adds coldkey gen by @Gus in 88eb46b
+* password working - needs cleanup by @Gus in 6ce4476
+* docs and signing working by @Gus in 742fbb4
+* testing sign by @Gus in 4b668e4
+* clean up and todos by @Gus in 4987d39
+* fmt by @Gus in 722c9e9
+* feat: wallet & hotkey file write by @Gus in 411d7be
+* feat: mnemonic & hotkey pair by @Gus in 43a4704
+* adds tests by @Gus in e8a2ac6
+* adds documentation by @Gus in f903f03
 
 ## New Contributors
-* @roman-opentensor made their first contribution in https://github.com/opentensor/btwallet/pull/1
-* @thewhaleking made their first contribution in https://github.com/opentensor/btwallet/pull/3
-* @ibraheem-opentensor made their first contribution in https://github.com/opentensor/btwallet/pull/6
 * @camfairchild made their first contribution in https://github.com/opentensor/btwallet/pull/30
 * @rajkaramchedu made their first contribution in https://github.com/opentensor/btwallet/pull/31
+* @opendansor made their first contribution in 707505209df4089581d4d2ca17b26a6e162c7f05
 
 **Full Changelog**: https://github.com/opentensor/btwallet/commits/v2.0.0
 
@@ -33,6 +249,7 @@
 
 Bittensor wallet separated as a standalone package:
 
+* Release 1.0.0 - Changelog by @ibraheem-opentensor in https://github.com/opentensor/btwallet/pull/24
 * Deployment script + updates by @ibraheem-opentensor in https://github.com/opentensor/btwallet/pull/23
 * bump up cryptography version until 43.0.1 by @roman-opentensor in https://github.com/opentensor/btwallet/pull/21
 * bump up the version by @roman-opentensor in https://github.com/opentensor/btwallet/pull/20
@@ -48,3 +265,8 @@ Bittensor wallet separated as a standalone package:
 * Fix flaky tests by @roman-opentensor in https://github.com/opentensor/btwallet/pull/5
 * Bump substrate-interface to 1.7.9 by @thewhaleking in https://github.com/opentensor/btwallet/pull/3
 * Massive wallet change after copy-paste by @roman-opentensor in https://github.com/opentensor/btwallet/pull/1
+
+## New Contributors
+* @roman-opentensor made their first contribution in https://github.com/opentensor/btwallet/pull/1
+* @thewhaleking made their first contribution in https://github.com/opentensor/btwallet/pull/3
+* @ibraheem-opentensor made their first contribution in https://github.com/opentensor/btwallet/pull/6
diff --git a/README.md b/README.md
index 8c012e3..eb1b5eb 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@
 
 ## Bittensor Wallet SDK
 
-The Bittensor Wallet SDK is a Python interface for a powerful Rust-based Bittensor wallet functionality. You do not need to know Rust to use this Wallet SDK. However, if you want to contribute to the Rust components of this Wallet SDK, the Rust source is located in the [src](./src) directory.
+The Bittensor Wallet SDK is a Python interface for a powerful Rust-based Bittensor wallet functionality. You do not need to know Rust to use this Wallet SDK. However, if you want to contribute to the Rust components of this Wallet SDK, the Rust source is located in the [src](./src) directory. You can look at the development history by inspecting the Changelog.md.
 
 ## Documentation
 
diff --git a/RUST-README.md b/RUST-README.md
index 8b7396b..1798983 100644
--- a/RUST-README.md
+++ b/RUST-README.md
@@ -64,14 +64,12 @@ kpw = WKeypair.create_from_mnemonic("stool feel open east woman high can denial
 assert kps.ss58_address == kpw.ss58_address
 assert kps.seed_hex == kpw.seed_hex
 assert kps.public_key == kpw.public_key
-assert kps.private_key == kpw.private_key
 
 kps = SKeypair.create_from_seed("0x023d5fbd7981676587a9f7232aeae1087ac7c265f9658fb643b6f5e61961dfbf")
 kpw = WKeypair.create_from_seed("0x023d5fbd7981676587a9f7232aeae1087ac7c265f9658fb643b6f5e61961dfbf")
 assert kps.ss58_address == kpw.ss58_address
 assert kps.seed_hex == kpw.seed_hex
 assert kps.public_key == kpw.public_key
-assert kps.private_key == kpw.private_key
 
 # substrateinterface has a bug -> can't create the KP without `ss58_format` passed
 kps = SKeypair.create_from_private_key("0x2b400f61c21cbaad4d5cb2dcbb4ef4fcdc238b98d04d48c6d2a451ebfd306c0eed845edcc69b0a19a6905afed0dd84c16ebd0f458928f2e91a6b67b95fc0b42f", ss58_format=42)
@@ -79,14 +77,12 @@ kpw = WKeypair.create_from_private_key("0x2b400f61c21cbaad4d5cb2dcbb4ef4fcdc238b
 assert kps.ss58_address == kpw.ss58_address
 assert kps.seed_hex == kpw.seed_hex
 assert kps.public_key == kpw.public_key
-assert kps.private_key == kpw.private_key
 
 kps = SKeypair.create_from_uri("//Alice")
 kpw = WKeypair.create_from_uri("//Alice")
 assert kps.ss58_address == kpw.ss58_address
 assert kps.seed_hex == kpw.seed_hex
 assert kps.public_key == kpw.public_key
-assert kps.private_key == kpw.private_key
 
 # substrateinterface has a bug -> can't create the KP without `ss58_format` passed
 from_private_key_new_kps = SKeypair(public_key=kps.public_key.hex(), ss58_format=42)
@@ -94,14 +90,13 @@ from_private_key_new_kpw = WKeypair(public_key=kps.public_key.hex())
 assert from_private_key_new_kps.ss58_address == from_private_key_new_kpw.ss58_address
 assert from_private_key_new_kps.seed_hex == from_private_key_new_kpw.seed_hex
 assert from_private_key_new_kps.public_key == from_private_key_new_kpw.public_key
-assert from_private_key_new_kps.private_key == from_private_key_new_kpw.private_key
+
 
 from_address_kps = SKeypair(ss58_address="5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY")
 from_address_kpw = WKeypair(ss58_address="5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY")
 assert from_address_kps.ss58_address == from_address_kpw.ss58_address
 assert from_address_kps.seed_hex == from_address_kpw.seed_hex
 assert from_address_kps.public_key == from_address_kpw.public_key
-assert from_address_kps.private_key == from_address_kpw.private_key
 
 # check signature
 assert kps.verify("asd", kpw.sign("asd")) == True
@@ -120,8 +115,6 @@ skp = S_Keypair.create_from_encrypted_json(data, passphrase)
 wkp = W_Keypair.create_from_encrypted_json(data, passphrase)
 
 assert skp.ss58_format == wkp.ss58_format
-assert skp.private_key == wkp.private_key
-assert skp.private_key.hex() == wkp.private_key.hex()
 assert skp.public_key == wkp.public_key
 assert skp.public_key.hex() == wkp.public_key.hex()
 ```
@@ -192,7 +185,6 @@ assert isinstance(kp_2, Keypair)
 assert kp.ss58_address == kp_2.ss58_address
 assert kp.seed_hex == kp_2.seed_hex
 assert kp.public_key == kp_2.public_key
-assert kp.private_key == kp_2.private_key
 ```
 
 #### Test Keyfile encrypt and decrypt
@@ -306,4 +298,29 @@ kf.make_dirs()
 assert kf.exists_on_device() == True
 assert kf.is_writable() == True
 assert kf.is_readable() == True
+```
+
+### Config parsing test
+```python
+import argparse
+import bittensor as bt
+
+parser = argparse.ArgumentParser(description='My parser')
+
+bt.wallet.add_args(parser)
+bt.subtensor.add_args(parser)
+bt.axon.add_args(parser)
+bt.logging.add_args(parser)
+
+config = bt.config(parser)
+
+config.wallet.name = "new_wallet_name"
+config.wallet.hotkey = "new_hotkey"
+config.wallet.path = "/some/not_default/path"
+
+wallet = bt.wallet(config=config)
+
+assert wallet.name == config.wallet.name
+assert wallet.hotkey_str == config.wallet.hotkey
+assert wallet.path == config.wallet.path
 ```
\ No newline at end of file
diff --git a/bittensor_wallet/__init__.py b/bittensor_wallet/__init__.py
index eb5c5fb..78d8376 100644
--- a/bittensor_wallet/__init__.py
+++ b/bittensor_wallet/__init__.py
@@ -1,14 +1,16 @@
 from bittensor_wallet.bittensor_wallet import (
-    Config as Config,
-    Keyfile as Keyfile,
-    Keypair as Keypair,
-    Wallet as Wallet,
-    utils as utils,
-    config as config,
-    errors as errors,
-    keypair as keypair,
-    keyfile as keyfile,
-    wallet as wallet,
+    # classes
+    Config,
+    Keyfile,
+    Keypair,
+    Wallet,
+    # modules
+    config,
+    errors,
+    keyfile,
+    keypair,
+    utils,
+    wallet,
 )
 
 __version__ = "2.0.0"
diff --git a/bittensor_wallet/keyfile/__init__.py b/bittensor_wallet/keyfile/__init__.py
index 522284e..9e5c0da 100644
--- a/bittensor_wallet/keyfile/__init__.py
+++ b/bittensor_wallet/keyfile/__init__.py
@@ -11,7 +11,9 @@
 keyfile_data_is_encrypted = _.keyfile_data_is_encrypted
 keyfile_data_encryption_method = _.keyfile_data_encryption_method
 legacy_encrypt_keyfile_data = _.legacy_encrypt_keyfile_data
-get_coldkey_password_from_environment = _.get_coldkey_password_from_environment
+get_coldkey_password_from_environment = _.get_password_from_environment
+# backwards compatibility
+get_password_from_environment = get_coldkey_password_from_environment
 encrypt_keyfile_data = _.encrypt_keyfile_data
 decrypt_keyfile_data = _.decrypt_keyfile_data
 Keyfile = _.Keyfile
diff --git a/pyproject.toml b/pyproject.toml
index e043022..310ab6e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "bittensor-wallet"
-version = "2.0.1"
+version = "2.0.2"
 description = ""
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -20,11 +20,10 @@ dependencies = [
 requires-python = ">= 3.9"
 
 authors = [
-  {name = "Cameron Fairchild", email = "cameron@opentensor.dev"},
-  {name = "Opentensor Foundation", email = "hello@bittensor.com"},
+  {name = "Roman Chkhaidze", email = "roman@opentensor.dev"},
 ]
 maintainers = [
-  {name = "Cameron Fairchild", email = "cameron@opentensor.dev"}
+  {name = "Cortex Team", email = "cortex@opentensor.dev"},
 ]
 classifiers = [
         "Development Status :: 5 - Production/Stable",
diff --git a/src/keyfile.rs b/src/keyfile.rs
index d9031e6..90a7a27 100644
--- a/src/keyfile.rs
+++ b/src/keyfile.rs
@@ -23,6 +23,9 @@ use crate::errors::KeyFileError;
 use crate::keypair::Keypair;
 use crate::utils;
 
+use sodiumoxide::crypto::pwhash;
+use sodiumoxide::crypto::secretbox;
+
 type PyFileNotFoundError = KeyFileError;
 type PyIOError = KeyFileError;
 type PyOSError = KeyFileError;
@@ -31,9 +34,6 @@ type PyRuntimeError = KeyFileError;
 type PyUnicodeDecodeError = KeyFileError;
 type PyValueError = KeyFileError;
 
-use sodiumoxide::crypto::pwhash;
-use sodiumoxide::crypto::secretbox;
-
 const NACL_SALT: &[u8] = b"\x13q\x83\xdf\xf1Z\t\xbc\x9c\x90\xb5Q\x879\xe9\xb1";
 const LEGACY_SALT: &[u8] = b"Iguesscyborgslikemyselfhaveatendencytobeparanoidaboutourorigins";
 
@@ -174,7 +174,7 @@ pub fn validate_password(_py: Python, password: &str) -> PyResult<bool> {
             Ok(false)
         }
     } else {
-        utils::print("Password not strong enough. Try increasing the length of the password or the password complexity.".to_string());
+        utils::print("Password not strong enough. Try increasing the length of the password or the password complexity.\n".to_string());
         Ok(false)
     }
 }
@@ -185,7 +185,6 @@ pub fn validate_password(_py: Python, password: &str) -> PyResult<bool> {
 ///         password (str): The valid password entered by the user.
 #[pyfunction]
 pub fn ask_password(py: Python, validation_required: bool) -> PyResult<String> {
-
     let mut valid = false;
     let password = utils::prompt_password("Enter your password: ".to_string());
 
@@ -304,7 +303,9 @@ pub fn legacy_encrypt_keyfile_data(
         // function to get password from user
         ask_password(py, true).unwrap());
 
-    utils::print(":exclamation_mark: Encrypting key with legacy encryption method...".to_string());
+    utils::print(
+        ":exclamation_mark: Encrypting key with legacy encryption method...\n".to_string(),
+    );
 
     // Encrypting key with legacy encryption method
     let encrypted_data = encrypt_vault(keyfile_data, password.as_str())
@@ -321,16 +322,18 @@ pub fn legacy_encrypt_keyfile_data(
 /// # Returns
 /// * `Option<String>` - The password retrieved from the environment variables, or `None` if not found.
 #[pyfunction]
-#[pyo3(signature = (coldkey_name))]
-pub fn get_coldkey_password_from_environment(
+#[pyo3(signature = (env_var_name))]
+pub fn get_password_from_environment(
     _py: Python,
-    coldkey_name: String,
+    env_var_name: String,
 ) -> PyResult<Option<String>> {
-    let env_key: String = format!(
-        "BT_COLD_PW_{}",
-        coldkey_name.to_uppercase().replace('-', "_")
-    );
-    Ok(env::var(env_key).ok())
+    match env::var(&env_var_name) {
+        Ok(encrypted_password) => {
+            let decrypted_password = decrypt_password(encrypted_password, env_var_name);
+            Ok(Some(decrypted_password))
+        }
+        Err(_) => Ok(None),
+    }
 }
 
 // decrypt of keyfile_data with secretbox
@@ -344,7 +347,7 @@ fn derive_key(password: &[u8]) -> secretbox::Key {
         pwhash::argon2i13::OPSLIMIT_SENSITIVE,
         pwhash::argon2i13::MEMLIMIT_SENSITIVE,
     )
-        .expect("Failed to derive key for NaCl decryption.");
+    .expect("Failed to derive key for NaCl decryption.");
     key
 }
 
@@ -369,7 +372,7 @@ pub fn encrypt_keyfile_data(
         None => ask_password(py, true)?,
     };
 
-    utils::print("Encrypting...".to_string());
+    utils::print("Encrypting...\n".to_string());
 
     // crate the key with pwhash Argon2i
     let key = derive_key(password.as_bytes());
@@ -397,12 +400,12 @@ pub fn encrypt_keyfile_data(
 ///     Returns
 ///         decrypted_data (bytes): The decrypted data.
 #[pyfunction]
-#[pyo3(signature = (keyfile_data, password = None, coldkey_name = None))]
+#[pyo3(signature = (keyfile_data, password = None, password_env_var = None))]
 pub fn decrypt_keyfile_data(
     py: Python,
     keyfile_data: &[u8],
     password: Option<String>,
-    coldkey_name: Option<String>,
+    password_env_var: Option<String>,
 ) -> PyResult<PyObject> {
     // decrypt of keyfile_data with secretbox
     fn nacl_decrypt(keyfile_data: &[u8], key: &secretbox::Key) -> PyResult<Vec<u8>> {
@@ -430,10 +433,10 @@ pub fn decrypt_keyfile_data(
 
     let mut password = password;
 
-    // Retrieve password from environment variable if coldkey_name is provided
-    if let Some(coldkey) = coldkey_name {
+    // Retrieve password from environment variable if env_var_name is provided
+    if let Some(env_var_name_) = password_env_var {
         if password.is_none() {
-            password = get_coldkey_password_from_environment(py, coldkey)?;
+            password = get_password_from_environment(py, env_var_name_)?;
         }
     }
 
@@ -444,7 +447,7 @@ pub fn decrypt_keyfile_data(
 
     let password = password.unwrap();
 
-    utils::print("Decrypting...".to_string());
+    utils::print("Decrypting...\n".to_string());
 
     // NaCl decryption
     if keyfile_data_is_encrypted_nacl(py, keyfile_data)? {
@@ -488,21 +491,50 @@ fn expand_tilde(path: &str) -> String {
     path.to_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);
+    }
+    println!(
+        ">>> key: {}, value: {}, encrypted: {}",
+        key, value, encrypted
+    );
+    encrypted
+}
+
+// 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
+}
+
 #[derive(Clone)]
 #[pyclass(subclass)]
 pub struct Keyfile {
     path: String,
     name: String,
+    should_save_to_env: bool,
 }
 
 #[pymethods]
 impl Keyfile {
     #[new]
-    #[pyo3(signature = (path, name = None))]
-    pub fn new(path: String, name: Option<String>) -> PyResult<Self> {
+    #[pyo3(signature = (path, name=None, should_save_to_env=false))]
+    pub fn new(path: String, name: Option<String>, should_save_to_env: bool) -> PyResult<Self> {
         let path = expand_tilde(&path);
         let name = name.unwrap_or_else(|| "Keyfile".to_string());
-        Ok(Keyfile { path, name })
+        Ok(Keyfile {
+            path,
+            name,
+            should_save_to_env,
+        })
     }
 
     #[allow(clippy::bool_comparison)]
@@ -526,7 +558,7 @@ impl Keyfile {
 
     /// Returns the keypair from path, decrypts data if the file is encrypted.
     #[getter(keypair)]
-    pub fn keypair_py(&self, py: Python) -> PyResult<Keypair>{
+    pub fn keypair_py(&self, py: Python) -> PyResult<Keypair> {
         self.get_keypair(None, py)
     }
 
@@ -540,7 +572,7 @@ impl Keyfile {
 
         // check if encrypted
         let decrypted_keyfile_data = if keyfile_data_is_encrypted(py, keyfile_data_bytes)? {
-            decrypt_keyfile_data(py, keyfile_data_bytes, password, Some(self.name.clone()))?
+            decrypt_keyfile_data(py, keyfile_data_bytes, password, Some(self.env_var_name()?))?
         } else {
             keyfile_data
         };
@@ -576,6 +608,16 @@ impl Keyfile {
         self._read_keyfile_data_from_file(py)
     }
 
+    /// Returns local environment variable key name based on Keyfile path.
+    #[getter]
+    fn env_var_name(&self) -> PyResult<String> {
+        let path = &self
+            .path
+            .replace(std::path::MAIN_SEPARATOR, "_")
+            .replace('.', "_");
+        Ok(format!("BT_PW_{}", path.to_uppercase()))
+    }
+
     /// Writes the keypair to the file and optionally encrypts data.
     #[pyo3(signature = (keypair, encrypt = true, overwrite = false, password = None))]
     pub fn set_keypair(
@@ -591,7 +633,14 @@ impl Keyfile {
         let keyfile_data = serialized_keypair_to_keyfile_data(py, &keypair)?;
 
         let final_keyfile_data = if encrypt {
-            let encrypted_data = encrypt_keyfile_data(py, keyfile_data.extract(py)?, password)?;
+            let encrypted_data =
+                encrypt_keyfile_data(py, keyfile_data.extract(py)?, password.clone())?;
+
+            // store password to local env
+            if self.should_save_to_env {
+                self.save_password_to_env(password.clone(), py)?;
+            }
+
             encrypted_data.extract::<&[u8]>(py)?
         } else {
             keyfile_data.extract::<&[u8]>(py)?
@@ -695,7 +744,7 @@ impl Keyfile {
             "File {} already exists. Overwrite? (y/N) ",
             self.path
         ))
-            .expect("Failed to read input.");
+        .expect("Failed to read input.");
 
         choice.trim().to_lowercase() == "y"
     }
@@ -710,21 +759,21 @@ impl Keyfile {
     ) -> PyResult<bool> {
         if !self.exists_on_device()? {
             if print_result {
-                utils::print(format!("Keyfile does not exist. {}", self.path));
+                utils::print(format!("Keyfile '{}' does not exist.\n", self.path));
             }
             return Ok(false);
         }
 
         if !self.is_readable()? {
             if print_result {
-                utils::print(format!("Keyfile is not readable. {}", self.path));
+                utils::print(format!("Keyfile '{}' is not readable.\n", self.path));
             }
             return Ok(false);
         }
 
         if !self.is_writable()? {
             if print_result {
-                utils::print(format!("Keyfile is not writable. {}", self.path));
+                utils::print(format!("Keyfile '{}' is not writable.\n", self.path));
             }
             return Ok(false);
         }
@@ -739,7 +788,7 @@ impl Keyfile {
             if keyfile_data_is_encrypted(py, keyfile_data_bytes)?
                 && !keyfile_data_is_encrypted_nacl(py, keyfile_data_bytes)?
             {
-                utils::print("You may update the keyfile to improve security...".to_string());
+                utils::print("You may update the keyfile to improve security...\n".to_string());
 
                 // ask user for the confirmation for updating
                 if update_keyfile == confirm_prompt("Update keyfile?") {
@@ -748,7 +797,7 @@ impl Keyfile {
                     // check mnemonic if saved
                     while !stored_mnemonic {
                         utils::print(
-                            "Please store your mnemonic in case an error occurs...".to_string(),
+                            "Please store your mnemonic in case an error occurs...\n".to_string(),
                         );
                         if confirm_prompt("Have you stored the mnemonic?") {
                             stored_mnemonic = true;
@@ -768,7 +817,7 @@ impl Keyfile {
                             py,
                             keyfile_data_bytes,
                             Some(pwd),
-                            Some(self.name.clone()),
+                            Some(self.env_var_name()?),
                         ) {
                             Ok(decrypted_data) => {
                                 let data: Vec<u8> = decrypted_data.extract(py)?;
@@ -804,17 +853,17 @@ impl Keyfile {
 
             return if !keyfile_data_is_encrypted(py, keyfile_data_bytes)? {
                 if print_result {
-                    utils::print("Keyfile is not encrypted.".to_string());
+                    utils::print("Keyfile is not encrypted.\n".to_string());
                 }
                 Ok(false)
             } else if keyfile_data_is_encrypted_nacl(py, keyfile_data_bytes)? {
                 if print_result {
-                    utils::print("Keyfile is updated.".to_string());
+                    utils::print("Keyfile is updated.\n".to_string());
                 }
                 Ok(true)
             } else {
                 if print_result {
-                    utils::print("Keyfile is outdated, please update using 'btcli'.".to_string());
+                    utils::print("Keyfile is outdated, please update using 'btcli'.\n".to_string());
                 }
                 Ok(false)
             };
@@ -824,7 +873,7 @@ impl Keyfile {
 
     /// Encrypts the file under the path.
     #[pyo3(signature = (password = None))]
-    pub fn encrypt(&self, password: Option<String>, py: Python) -> PyResult<()> {
+    pub fn encrypt(&self, mut password: Option<String>, py: Python) -> PyResult<()> {
         // checkers
         if !self.exists_on_device()? {
             return Err(PyErr::new::<PyValueError, _>(format!(
@@ -855,7 +904,19 @@ impl Keyfile {
             let as_keypair = deserialize_keypair_from_keyfile_data(keyfile_data_bytes, py)?;
             let serialized_data = serialized_keypair_to_keyfile_data(py, &as_keypair)?;
 
-            encrypt_keyfile_data(py, serialized_data.extract(py)?, password)?
+            // get password from local env if exist
+            if password.is_none() {
+                password = get_password_from_environment(py, self.env_var_name()?)?;
+            }
+
+            let encrypted_keyfile_data =
+                encrypt_keyfile_data(py, serialized_data.extract(py)?, password.clone())?;
+
+            if self.should_save_to_env {
+                self.save_password_to_env(password.clone(), py)?;
+            }
+
+            encrypted_keyfile_data
         } else {
             keyfile_data
         };
@@ -894,7 +955,7 @@ impl Keyfile {
         let keyfile_data_bytes: &[u8] = keyfile_data.extract(py)?;
 
         let decrypted_data = if keyfile_data_is_encrypted(py, keyfile_data_bytes)? {
-            decrypt_keyfile_data(py, keyfile_data_bytes, password, Some(self.name.clone()))?
+            decrypt_keyfile_data(py, keyfile_data_bytes, password, Some(self.env_var_name()?))?
         } else {
             keyfile_data
         };
@@ -984,4 +1045,59 @@ impl Keyfile {
         })?;
         Ok(())
     }
+
+    /// Saves the key's password to the associated local environment variable.
+    #[pyo3(signature = (password=None))]
+    fn save_password_to_env(&self, password: Option<String>, py: Python) -> PyResult<bool> {
+        // checking the password
+        let password = match password {
+            Some(pwd) => pwd,
+            None => match ask_password(py, true) {
+                Ok(pwd) => pwd,
+                Err(e) => {
+                    utils::print(format!("Error asking password: {:?}.\n", e));
+                    return Ok(false);
+                }
+            },
+        };
+        // 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(true)
+            }
+            Err(e) => {
+                utils::print(format!(
+                    "Error saving environment variable name: {:?}.\n",
+                    e
+                ));
+                Ok(false)
+            }
+        }
+    }
+
+    /// Removes the password associated with the Keyfile from the local environment.
+    fn remove_password_from_env(&self) -> PyResult<bool> {
+        let env_var_name = self.env_var_name()?;
+
+        if env::var(&env_var_name).is_ok() {
+            env::remove_var(&env_var_name);
+            let message = format!("Environment variable '{}' removed.\n", env_var_name);
+            utils::print(message);
+            Ok(true)
+        } else {
+            let message = format!("Environment variable '{}' does not exist.\n", env_var_name);
+            utils::print(message);
+            Ok(false)
+        }
+    }
 }
diff --git a/src/keypair.rs b/src/keypair.rs
index 095f6de..6090d49 100644
--- a/src/keypair.rs
+++ b/src/keypair.rs
@@ -587,7 +587,6 @@ impl Keypair {
     }
 }
 
-
 // Default values for Keypair
 impl Default for Keypair {
     fn default() -> Self {
diff --git a/src/lib.rs b/src/lib.rs
index 4a90cab..98db472 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -80,7 +80,7 @@ fn register_keyfile_module(main_module: &Bound<'_, PyModule>) -> PyResult<()> {
         &keyfile_module
     )?)?;
     keyfile_module.add_function(wrap_pyfunction!(
-        keyfile::get_coldkey_password_from_environment,
+        keyfile::get_password_from_environment,
         &keyfile_module
     )?)?;
     keyfile_module.add_function(wrap_pyfunction!(
diff --git a/src/utils.rs b/src/utils.rs
index 5d1aea2..2694f11 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -115,8 +115,16 @@ pub fn print(s: String) {
     Python::with_gil(|py| {
         let locals = PyDict::new_bound(py);
         locals.set_item("s", s).unwrap();
-        py.run_bound("print(s, end='')", None, Some(&locals))
-            .unwrap();
+        py.run_bound(
+            r#"
+import sys
+print(s, end='')
+sys.stdout.flush()
+"#,
+            None,
+            Some(&locals),
+        )
+        .unwrap();
     });
 }
 
diff --git a/src/wallet.rs b/src/wallet.rs
index 6537123..69169bb 100644
--- a/src/wallet.rs
+++ b/src/wallet.rs
@@ -20,7 +20,7 @@ type PyRuntimeError = KeyFileError;
 #[pyfunction]
 #[pyo3(signature = (mnemonic, key_type))]
 pub fn display_mnemonic_msg(mnemonic: String, key_type: &str) {
-    utils::print(format!("{}", "\nIMPORTANT: Store this mnemonic in a secure (preferable offline place), as anyone who has possession of this mnemonic can use it to regenerate the key and access your tokens.".red()));
+    utils::print(format!("{}", "\nIMPORTANT: Store this mnemonic in a secure (preferable offline place), as anyone who has possession of this mnemonic can use it to regenerate the key and access your tokens.\n".red()));
 
     utils::print(format!(
         "\nThe mnemonic to the new {} is: {}",
@@ -33,13 +33,33 @@ pub fn display_mnemonic_msg(mnemonic: String, key_type: &str) {
     ));
 }
 
+// Function to safely retrieve attribute as Option<String> from passed python object
+fn get_attribute_string(py: Python, obj: &Bound<PyAny>, attr_name: &str) -> PyResult<Option<String>> {
+    match obj.getattr(attr_name) {
+        Ok(attr) => {
+            if attr.is_none() {
+                Ok(None)
+            } else {
+                let value: String = attr.extract()?;
+                Ok(Some(value))
+            }
+        }
+        Err(e) => {
+            if e.is_instance_of::<pyo3::exceptions::PyAttributeError>(py) {
+                Ok(None)
+            } else {
+                Err(e)
+            }
+        }
+    }
+}
+
 #[derive(Clone)]
 #[pyclass(subclass)]
 pub struct Wallet {
     pub name: String,
     pub path: String,
     pub hotkey: String,
-    pub config: Option<Config>,
 
     _coldkey: Option<Keypair>,
     _coldkeypub: Option<Keypair>,
@@ -61,23 +81,63 @@ impl Wallet {
         name: Option<String>,
         hotkey: Option<String>,
         path: Option<String>,
-        config: Option<Config>,
+        config: Option<PyObject>,
+        py: Python
     ) -> PyResult<Wallet> {
-        let final_name = name
-            .or_else(|| Some(config.clone()?.wallet.name.clone()))
-            .unwrap_or_else(|| BT_WALLET_NAME.to_string());
-        let final_hotkey = hotkey
-            .or_else(|| Some(config.clone()?.wallet.hotkey.clone()))
-            .unwrap_or_else(|| BT_WALLET_HOTKEY.to_string());
-        let final_path = path
-            .or_else(|| Some(config.clone()?.wallet.path.clone()))
-            .unwrap_or_else(|| BT_WALLET_PATH.to_string());
+
+        // default config's values if config and config.wallet exist
+        let mut conf_name: Option<String> = None;
+        let mut conf_hotkey: Option<String> = None;
+        let mut conf_path: Option<String> = None;
+
+        // parse python config object if passed
+        if let Some(config_obj) = config {
+            let config_ref = config_obj.bind(py);
+
+            // parse python config.wallet object if exist in config object
+            if config_ref.hasattr("wallet")? {
+                let wallet_obj = config_ref.getattr("wallet")?;
+
+                if !wallet_obj.is_none() {
+                    let wallet_ref = wallet_obj.as_ref();
+
+                    // assign values instead of default ones
+                    conf_name = get_attribute_string(py, wallet_ref, "name")?;
+                    conf_hotkey = get_attribute_string(py, wallet_ref, "hotkey")?;
+                    conf_path = get_attribute_string(py, wallet_ref, "path")?;
+                }
+            }
+        }
+
+        let final_name = if let Some(name) = name {
+            name
+        } else if let Some(conf_name) = conf_name {
+            conf_name
+        } else {
+            BT_WALLET_NAME.to_string()
+        };
+
+        let final_hotkey = if let Some(hotkey) = hotkey {
+            hotkey
+        } else if let Some(conf_hotkey) = conf_hotkey {
+            conf_hotkey
+        } else {
+            BT_WALLET_HOTKEY.to_string()
+        };
+
+        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()
+        } else {
+            BT_WALLET_PATH.to_string()
+        };
 
         Ok(Wallet {
             name: final_name,
             hotkey: final_hotkey,
             path: final_path,
-            config: config.or(None),
+
             _coldkey: None,
             _coldkeypub: None,
             _hotkey: None,
@@ -124,7 +184,7 @@ impl Wallet {
         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(|_| BT_WALLET_PATH.to_string());
+            env::var("BT_WALLET_PATH").unwrap_or_else(|_| format!("~/{}", BT_WALLET_PATH.to_string()));
 
         let prefix_str = if let Some(value) = prefix {
             format!("\"{}\"", value)
@@ -172,41 +232,92 @@ except argparse.ArgumentError:
     }
 
     /// Checks for existing coldkeypub and hotkeys, and creates them if non-existent.
-    #[pyo3(signature = (coldkey_use_password=true, hotkey_use_password=false))]
+    #[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_if_non_existent(
         &mut self,
         coldkey_use_password: bool,
         hotkey_use_password: bool,
+        save_coldkey_to_env: bool,
+        save_hotkey_to_env: bool,
+        coldkey_password: Option<String>,
+        hotkey_password: Option<String>,
+        overwrite: bool,
+        suppress: bool,
         py: Python,
     ) -> PyResult<Wallet> {
-        self.create(coldkey_use_password, hotkey_use_password, py)
+        self.create(
+            coldkey_use_password,
+            hotkey_use_password,
+            save_coldkey_to_env,
+            save_hotkey_to_env,
+            coldkey_password,
+            hotkey_password,
+            overwrite,
+            suppress,
+            py,
+        )
     }
 
     /// 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.
     #[allow(clippy::bool_comparison)]
-    #[pyo3(signature = (coldkey_use_password=true, hotkey_use_password=false))]
+    #[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: bool,
         hotkey_use_password: bool,
+        save_coldkey_to_env: bool,
+        save_hotkey_to_env: bool,
+        coldkey_password: Option<String>,
+        hotkey_password: Option<String>,
+        overwrite: bool,
+        suppress: bool,
         py: Python,
     ) -> PyResult<Self> {
-        if self.coldkey_file()?.exists_on_device()? == false
-            && self.coldkeypub_file()?.exists_on_device()? == false
+        if overwrite
+            || self.coldkey_file()?.exists_on_device()? == false
+                && self.coldkeypub_file()?.exists_on_device()? == false
         {
-            self.create_new_coldkey(12, coldkey_use_password, false, false, py)?;
+            self.create_new_coldkey(
+                12,
+                coldkey_use_password,
+                overwrite,
+                suppress,
+                save_coldkey_to_env,
+                coldkey_password,
+                py,
+            )?;
         } else {
             utils::print(format!(
-                "ColdKey for the wallet '{}' already exists.",
+                "ColdKey for the wallet '{}' already exists.\n",
                 self.name
             ));
         }
 
-        if !self.hotkey_file()?.exists_on_device()? {
-            self.create_new_hotkey(12, hotkey_use_password, false, false, py)?;
+        if overwrite || !self.hotkey_file()?.exists_on_device()? {
+            self.create_new_hotkey(
+                12,
+                hotkey_use_password,
+                overwrite,
+                suppress,
+                save_hotkey_to_env,
+                hotkey_password,
+                py,
+            )?;
         } else {
             utils::print(format!(
-                "HotKey for the wallet '{}' already exists.",
+                "HotKey for the wallet '{}' already exists.\n",
                 self.name
             ));
         }
@@ -214,16 +325,51 @@ except argparse.ArgumentError:
         Ok(self.clone())
     }
 
-    /// Checks for existing coldkeypub and hotkeys and creates them if non-existent.
-    #[pyo3(signature = (coldkey_use_password=true, hotkey_use_password=false))]
+    /// Checks for existing coldkeypub and hotkeys, and recreates 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 recreate(
         &mut self,
         coldkey_use_password: bool,
         hotkey_use_password: bool,
+        save_coldkey_to_env: bool,
+        save_hotkey_to_env: bool,
+        coldkey_password: Option<String>,
+        hotkey_password: Option<String>,
+        overwrite: bool,
+        suppress: bool,
         py: Python,
     ) -> PyResult<Wallet> {
-        self.create_new_coldkey(12, coldkey_use_password, false, false, py)?;
-        self.create_new_hotkey(12, hotkey_use_password, false, false, py)?;
+        self.create_new_coldkey(
+            12,
+            coldkey_use_password,
+            overwrite,
+            suppress,
+            save_coldkey_to_env,
+            coldkey_password,
+            py,
+        )?;
+        self.create_new_hotkey(
+            12,
+            hotkey_use_password,
+            overwrite,
+            suppress,
+            save_hotkey_to_env,
+            hotkey_password,
+            py,
+        )?;
 
         Ok(self.clone())
     }
@@ -231,6 +377,12 @@ except argparse.ArgumentError:
     /// Property that returns the hotkey file.
     #[getter]
     pub fn hotkey_file(&self) -> PyResult<Keyfile> {
+        self.create_hotkey_file(false)
+    }
+
+    /// Created Hot Keyfile for Keypair
+    #[pyo3(signature = (save_hotkey_to_env=false))]
+    pub fn create_hotkey_file(&self, save_hotkey_to_env: bool) -> PyResult<Keyfile> {
         // get home dir
         let home = home_dir()
             .ok_or_else(|| PyErr::new::<PyRuntimeError, _>("Failed to get user home directory."))?;
@@ -243,13 +395,20 @@ except argparse.ArgumentError:
 
         Keyfile::new(
             hotkey_path.to_string_lossy().into_owned(),
-            Some(self.name.clone()),
+            Some(self.hotkey.clone()),
+            save_hotkey_to_env,
         )
     }
 
     /// Property that returns the coldkey file.
     #[getter]
     pub fn coldkey_file(&self) -> PyResult<Keyfile> {
+        self.create_coldkey_file(false)
+    }
+
+    /// Created Cold Keyfile for Keypair
+    #[pyo3(signature = (save_coldkey_to_env=false))]
+    pub fn create_coldkey_file(&self, save_coldkey_to_env: bool) -> PyResult<Keyfile> {
         // get home dir
         let home = home_dir()
             .ok_or_else(|| PyErr::new::<PyRuntimeError, _>("Failed to get user home directory."))?;
@@ -259,10 +418,10 @@ except argparse.ArgumentError:
 
         // concatenate coldkey path
         let coldkey_path = wallet_path.join("coldkey");
-
         Keyfile::new(
             coldkey_path.to_string_lossy().into_owned(),
-            Some(self.name.clone()),
+            Some("coldkey".to_string()),
+            save_coldkey_to_env,
         )
     }
 
@@ -281,7 +440,8 @@ except argparse.ArgumentError:
 
         Keyfile::new(
             coldkeypub_path.to_string_lossy().into_owned(),
-            Some(self.name.clone()),
+            Some("coldkeypub.txt".parse()?),
+            false,
         )
     }
 
@@ -330,17 +490,24 @@ except argparse.ArgumentError:
     }
 
     /// Sets the hotkey for the wallet.
-    #[pyo3(signature = (keypair, encrypt=true, overwrite=false))]
+    #[pyo3(signature = (keypair, encrypt=true, overwrite=false, save_coldkey_to_env=false, coldkey_password=None))]
     pub fn set_coldkey(
         &mut self,
         keypair: Keypair,
         encrypt: bool,
         overwrite: bool,
+        save_coldkey_to_env: bool,
+        coldkey_password: Option<String>,
         py: Python,
     ) -> PyResult<()> {
         self._coldkey = Some(keypair.clone());
-        self.coldkey_file()?
-            .set_keypair(keypair, encrypt, overwrite, None, py)
+        self.create_coldkey_file(save_coldkey_to_env)?.set_keypair(
+            keypair,
+            encrypt,
+            overwrite,
+            coldkey_password,
+            py,
+        )
     }
 
     /// Sets the coldkeypub for the wallet.
@@ -366,17 +533,24 @@ except argparse.ArgumentError:
     }
 
     /// Sets the hotkey for the wallet.
-    #[pyo3(signature = (keypair, encrypt=false, overwrite=false))]
+    #[pyo3(signature = (keypair, encrypt=false, overwrite=false, save_hotkey_to_env=false, hotkey_password=None))]
     pub fn set_hotkey(
         &mut self,
         keypair: Keypair,
         encrypt: bool,
         overwrite: bool,
+        save_hotkey_to_env: bool,
+        hotkey_password: Option<String>,
         py: Python,
     ) -> PyResult<()> {
         self._hotkey = Some(keypair.clone());
-        self.hotkey_file()?
-            .set_keypair(keypair.clone(), encrypt, overwrite, None, py)
+        self.create_hotkey_file(save_hotkey_to_env)?.set_keypair(
+            keypair.clone(),
+            encrypt,
+            overwrite,
+            hotkey_password,
+            py,
+        )
     }
 
     /// Gets the coldkey from the wallet.
@@ -398,13 +572,15 @@ except argparse.ArgumentError:
     }
 
     /// Creates coldkey from uri string, optionally encrypts it with the user-provided password.
-    #[pyo3(signature = (uri, use_password = true, overwrite = false, suppress = false))]
+    #[pyo3(signature = (uri, use_password = true, overwrite = false, suppress = false, save_coldkey_to_env=false, coldkey_password=None))]
     pub fn create_coldkey_from_uri(
         &mut self,
         uri: String,
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_coldkey_to_env: bool,
+        coldkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
         let keypair = Keypair::create_from_uri(uri.as_str())?;
@@ -415,19 +591,28 @@ except argparse.ArgumentError:
             }
         }
 
-        self.set_coldkey(keypair.clone(), use_password, overwrite, py)?;
+        self.set_coldkey(
+            keypair.clone(),
+            use_password,
+            overwrite,
+            save_coldkey_to_env,
+            coldkey_password,
+            py,
+        )?;
         self.set_coldkeypub(keypair.clone(), false, overwrite, py)?;
         Ok(self.clone())
     }
 
     /// Creates hotkey from uri string, optionally encrypts it with the user-provided password.
-    #[pyo3(signature = (uri, use_password = true, overwrite = false, suppress = false))]
+    #[pyo3(signature = (uri, use_password = true, overwrite = false, suppress = false, save_hotkey_to_env=false, hotkey_password=None))]
     pub fn create_hotkey_from_uri(
         &mut self,
         uri: String,
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_hotkey_to_env: bool,
+        hotkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
         let keypair = Keypair::create_from_uri(uri.as_str())?;
@@ -438,7 +623,14 @@ except argparse.ArgumentError:
             }
         }
 
-        self.set_hotkey(keypair.clone(), use_password, overwrite, py)?;
+        self.set_hotkey(
+            keypair.clone(),
+            use_password,
+            overwrite,
+            save_hotkey_to_env,
+            hotkey_password,
+            py,
+        )?;
         Ok(self.clone())
     }
 
@@ -482,26 +674,38 @@ except argparse.ArgumentError:
     }
 
     /// Creates a new coldkey, optionally encrypts it with the user-provided password and saves to disk.
-    #[pyo3(signature = (n_words = 12, use_password = true, overwrite = false, suppress = false))]
+    #[pyo3(signature = (n_words = 12, use_password=true, overwrite=false, suppress=false, save_coldkey_to_env=false, coldkey_password=None))]
     pub fn new_coldkey(
         &mut self,
         n_words: usize,
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_coldkey_to_env: bool,
+        coldkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
-        self.create_new_coldkey(n_words, use_password, overwrite, suppress, py)
+        self.create_new_coldkey(
+            n_words,
+            use_password,
+            overwrite,
+            suppress,
+            save_coldkey_to_env,
+            coldkey_password,
+            py,
+        )
     }
 
     /// Creates a new coldkey, optionally encrypts it with the user-provided password and saves to disk.
-    #[pyo3(signature = (n_words=12, use_password=true, overwrite=false, suppress=false))]
+    #[pyo3(signature = (n_words=12, use_password=true, overwrite=false, suppress=false, save_coldkey_to_env=false, coldkey_password=None))]
     fn create_new_coldkey(
         &mut self,
         n_words: usize,
-        use_password: bool,
+        mut use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_coldkey_to_env: bool,
+        coldkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
         let mnemonic = Keypair::generate_mnemonic(n_words)?;
@@ -511,33 +715,59 @@ except argparse.ArgumentError:
             display_mnemonic_msg(mnemonic, "coldkey");
         }
 
-        self.set_coldkey(keypair.clone(), use_password, overwrite, py)?;
+        // if coldkey_password is passed then coldkey_use_password always is true
+        use_password = if coldkey_password.is_some() {
+            true
+        } else {
+            use_password
+        };
+
+        self.set_coldkey(
+            keypair.clone(),
+            use_password,
+            overwrite,
+            save_coldkey_to_env,
+            coldkey_password,
+            py,
+        )?;
         self.set_coldkeypub(keypair.clone(), false, overwrite, py)?;
 
         Ok(self.clone())
     }
 
     /// Creates a new hotkey, optionally encrypts it with the user-provided password and saves to disk.
-    #[pyo3(signature = (n_words = 12, use_password = false, overwrite = false, suppress = false))]
+    #[pyo3(signature = (n_words = 12, use_password = false, overwrite = false, suppress = false, save_hotkey_to_env=false, hotkey_password=None))]
     pub fn new_hotkey(
         &mut self,
         n_words: usize,
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_hotkey_to_env: bool,
+        hotkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
-        self.create_new_hotkey(n_words, use_password, overwrite, suppress, py)
+        self.create_new_hotkey(
+            n_words,
+            use_password,
+            overwrite,
+            suppress,
+            save_hotkey_to_env,
+            hotkey_password,
+            py,
+        )
     }
 
     /// Creates a new hotkey, optionally encrypts it with the user-provided password and saves to disk.
-    #[pyo3(signature = (n_words=12, use_password=false, overwrite=false, suppress=false))]
+    #[pyo3(signature = (n_words=12, use_password=false, overwrite=false, suppress=false, save_hotkey_to_env=false, hotkey_password=None))]
     pub fn create_new_hotkey(
         &mut self,
         n_words: usize,
-        use_password: bool,
+        mut use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_hotkey_to_env: bool,
+        hotkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Wallet> {
         let mnemonic = Keypair::generate_mnemonic(n_words)?;
@@ -547,7 +777,21 @@ except argparse.ArgumentError:
             display_mnemonic_msg(mnemonic, "hotkey");
         }
 
-        self.set_hotkey(keypair.clone(), use_password, overwrite, py)?;
+        // if hotkey_password is passed then hotkey_use_password always is true
+        use_password = if hotkey_password.is_some() {
+            true
+        } else {
+            use_password
+        };
+
+        self.set_hotkey(
+            keypair.clone(),
+            use_password,
+            overwrite,
+            save_hotkey_to_env,
+            hotkey_password,
+            py,
+        )?;
         Ok(self.clone())
     }
 
@@ -590,7 +834,7 @@ except argparse.ArgumentError:
 
     /// Regenerates the coldkey from the passed mnemonic or seed, or JSON encrypts it with the user's password and saves the file.
     #[allow(clippy::too_many_arguments)]
-    #[pyo3(signature = (mnemonic=None, seed=None, json=None, use_password=true, overwrite=false, suppress=false))]
+    #[pyo3(signature = (mnemonic=None, seed=None, json=None, use_password=true, overwrite=false, suppress=false, save_coldkey_to_env=false, coldkey_password=None))]
     pub fn regenerate_coldkey(
         &mut self,
         mnemonic: Option<String>,
@@ -599,6 +843,8 @@ except argparse.ArgumentError:
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_coldkey_to_env: bool,
+        coldkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Self> {
         let keypair = if let Some(mnemonic) = mnemonic {
@@ -621,14 +867,21 @@ except argparse.ArgumentError:
             ));
         };
 
-        self.set_coldkey(keypair.clone(), use_password, overwrite, py)?;
+        self.set_coldkey(
+            keypair.clone(),
+            use_password,
+            overwrite,
+            save_coldkey_to_env,
+            coldkey_password,
+            py,
+        )?;
         self.set_coldkeypub(keypair.clone(), false, overwrite, py)?;
         Ok(self.clone())
     }
 
     /// Regenerates the hotkey from passed mnemonic or seed, encrypts it with the user's password and saves the file.
     #[allow(clippy::too_many_arguments)]
-    #[pyo3(signature = (mnemonic=None, seed=None, json=None, use_password=true, overwrite=false, suppress=false))]
+    #[pyo3(signature = (mnemonic=None, seed=None, json=None, use_password=true, overwrite=false, suppress=false, save_hotkey_to_env=false, hotkey_password=None))]
     pub fn regenerate_hotkey(
         &mut self,
         mnemonic: Option<String>,
@@ -637,6 +890,8 @@ except argparse.ArgumentError:
         use_password: bool,
         overwrite: bool,
         suppress: bool,
+        save_hotkey_to_env: bool,
+        hotkey_password: Option<String>,
         py: Python,
     ) -> PyResult<Self> {
         let keypair = if let Some(mnemonic) = mnemonic {
@@ -659,7 +914,14 @@ except argparse.ArgumentError:
             ));
         };
 
-        self.set_hotkey(keypair, use_password, overwrite, py)?;
+        self.set_hotkey(
+            keypair,
+            use_password,
+            overwrite,
+            save_hotkey_to_env,
+            hotkey_password,
+            py,
+        )?;
 
         Ok(self.clone())
     }


Test patch:
diff --git a/tests/test_keyfile.py b/tests/test_keyfile.py
index b8c33aa..182e9f1 100644
--- a/tests/test_keyfile.py
+++ b/tests/test_keyfile.py
@@ -15,17 +15,19 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-import os
 import json
-import time
-import pytest
+import os
 import shutil
+import time
+from unittest import mock
 
+import pytest
 from bip39 import bip39_validate
+from jedi.common import monkeypatch
 
-from bittensor_wallet.keyfile import get_coldkey_password_from_environment
+from bittensor_wallet.errors import ConfigurationError, KeyFileError
 from bittensor_wallet.keyfile import Keyfile
-from bittensor_wallet.errors import ConfigurationError
+from bittensor_wallet.keyfile import get_coldkey_password_from_environment
 from bittensor_wallet.keypair import Keypair
 
 
@@ -328,10 +330,6 @@ def test_create(keyfile_setup_teardown):
         keyfile.get_keypair(password="thisisafakepassword").ss58_address
         == alice.ss58_address
     )
-    assert (
-        keyfile.get_keypair(password="thisisafakepassword").private_key
-        == alice.private_key
-    )
     assert (
         keyfile.get_keypair(password="thisisafakepassword").public_key
         == alice.public_key
@@ -352,19 +350,20 @@ def test_create(keyfile_setup_teardown):
     repr(keyfile)
 
 
-def test_validate_password():
-    """
-    Test case for the validate_password function.
-
-    This function tests the behavior of the validate_password function from the bittensor.keyfile module.
-    It checks various scenarios to ensure that the function correctly validates passwords.
-    """
-    from bittensor_wallet.keyfile import validate_password
-
-    with pytest.raises(TypeError):
-        validate_password(None)
-    assert validate_password("passw0rd") is False
-    assert validate_password("123456789") is False
+# we can't mock rust methods
+# def test_validate_password():
+#     """
+#     Test case for the validate_password function.
+#
+#     This function tests the behavior of the validate_password function from the bittensor.keyfile module.
+#     It checks various scenarios to ensure that the function correctly validates passwords.
+#     """
+#     from bittensor_wallet.keyfile import validate_password
+#
+#     with pytest.raises(TypeError):
+#         validate_password(None)
+#     assert validate_password("passw0rd") is False
+#     assert validate_password("123456789") is False
 
 
 # def test_user_interface():
@@ -384,23 +383,23 @@ def test_validate_password():
 #         assert ask_password_to_encrypt() == "asdury3294y"
 
 
-# def test_overwriting(keyfile_setup_teardown):
-#     """
-#     Test case for overwriting a keypair in the keyfile.
-#     """
-#     root_path = keyfile_setup_teardown
-#     keyfile = Keyfile(path=os.path.join(root_path, "keyfile"))
-#     alice = Keypair.create_from_uri("/Alice")
-#     keyfile.set_keypair(
-#         alice, encrypt=True, overwrite=True, password="thisisafakepassword"
-#     )
-#     bob = Keypair.create_from_uri("/Bob")
-#
-#     with pytest.raises(KeyFileError):
-#         with mock.patch("builtins.input", return_value="n"):
-#             keyfile.set_keypair(
-#                 bob, encrypt=True, overwrite=False, password="thisisafakepassword"
-#             )
+def test_overwriting(keyfile_setup_teardown):
+    """
+    Test case for overwriting a keypair in the keyfile.
+    """
+    root_path = keyfile_setup_teardown
+    keyfile = Keyfile(path=os.path.join(root_path, "keyfile"))
+    alice = Keypair.create_from_uri("/Alice")
+    keyfile.set_keypair(
+        alice, encrypt=True, overwrite=True, password="thisisafakepassword"
+    )
+    bob = Keypair.create_from_uri("/Bob")
+
+    with pytest.raises(KeyFileError):
+        with mock.patch("builtins.input", return_value="n"):
+            keyfile.set_keypair(
+                bob, encrypt=True, overwrite=False, password="thisisafakepassword"
+            )
 
 
 def test_serialized_keypair_to_keyfile_data(keyfile_setup_teardown):
@@ -457,20 +456,21 @@ def test_deserialize_keypair_from_keyfile_data(keyfile_setup_teardown):
 
     assert deserialized_keypair.ss58_address == keypair.ss58_address
     assert deserialized_keypair.public_key == keypair.public_key
-    assert deserialized_keypair.private_key == keypair.private_key
-
-
-def test_get_coldkey_password_from_environment(monkeypatch):
-    password_by_wallet = {
-        "WALLET": "password",
-        "my_wallet": "password2",
-        "my-wallet": "password2",
-    }
-
-    monkeypatch.setenv("bt_cold_pw_wallet".upper(), password_by_wallet["WALLET"])
-    monkeypatch.setenv("BT_COLD_PW_My_Wallet".upper(), password_by_wallet["my_wallet"])
 
-    for wallet, password in password_by_wallet.items():
-        assert get_coldkey_password_from_environment(wallet) == password
 
-    assert get_coldkey_password_from_environment("non_existent_wallet") is None
+@pytest.mark.parametrize(
+    "env_name,encrypted,decrypted",
+    [
+        ("BT_PW_COLD_WALLET", "61,$>18", "testin{"),
+        ("BT_PW_COLD_WALLET", " =+$21,:!t``", "bittenoum0?7"),
+    ],
+)
+def test_get_coldkey_password_from_environment(
+    monkeypatch, env_name, encrypted, decrypted
+):
+    # Preps
+    monkeypatch.setenv(env_name, encrypted)
+
+    # Calls + Assertions
+    assert get_coldkey_password_from_environment(env_name) == decrypted
+    assert get_coldkey_password_from_environment("non_existent_env_variable") is None
diff --git a/tests/test_wallet.py b/tests/test_wallet.py
index 02acadb..b371683 100644
--- a/tests/test_wallet.py
+++ b/tests/test_wallet.py
@@ -17,10 +17,12 @@
 
 import json
 import time
+from unittest.mock import patch
+
 import pytest
-from bittensor_wallet import Wallet, keyfile
 from ansible_vault import Vault
-from unittest.mock import patch
+
+from bittensor_wallet import Wallet, keyfile
 
 
 def legacy_encrypt_keyfile_data(keyfile_data: bytes, password: str = None) -> bytes: