Skip to main content

bittensor_wallet/
utils.rs

1use sp_core::crypto::{AccountId32, Ss58Codec};
2use std::str;
3
4use crate::keypair::Keypair;
5
6pub(crate) const SS58_FORMAT: u8 = 42;
7
8/// Returns the SS58 format of the given address string.
9///
10///     Arguments:
11///         ss58_address (str): The SS58 address to extract the format from.
12///     Returns:
13///         format (u16): The SS58 format number.
14pub fn get_ss58_format(ss58_address: &str) -> Result<u16, &'static str> {
15    match <AccountId32 as Ss58Codec>::from_ss58check_with_version(ss58_address) {
16        Ok((_, format)) => Ok(u16::from(format)),
17        Err(_) => Err("Invalid SS58 address."),
18    }
19}
20
21/// Checks if the given address is a valid ss58 address.
22///
23///     Arguments:
24///         address (str): The address to check.
25///     Returns:
26///         ``True`` if the address is a valid ss58 address for Bittensor, ``False`` otherwise.
27pub fn is_valid_ss58_address(address: &str) -> bool {
28    if address.is_empty() {
29        // Possibly there could be a debug log, but not a print
30        // utils::print(format!("The given address is empty"));
31        return false;
32    }
33
34    sp_core::sr25519::Public::from_ss58check(address).is_ok()
35}
36
37///    Checks if the given public_key is a valid ed25519 key.
38///
39///     Arguments:
40///         public_key (str): The public_key to check as string.
41///     Returns:
42///         valid (bool): ``True`` if the public_key is a valid ed25519 key, ``False`` otherwise.
43pub fn is_string_valid_ed25519_pubkey(public_key: &str) -> bool {
44    if public_key.len() != 64 && public_key.len() != 66 {
45        return false;
46    }
47
48    let pub_key_var = Some(public_key.to_string());
49    let keypair_result = Keypair::new(None, pub_key_var, None, SS58_FORMAT, None, 1);
50
51    match keypair_result {
52        Ok(keypair) => keypair.ss58_address().is_some(),
53        Err(_) => false,
54    }
55}
56
57///    Checks if the given public_key is a valid ed25519 key.
58///
59///     Arguments:
60///         public_key (bytes): The public_key to check as bytes.
61///     Returns:
62///         valid (bool): ``True`` if the public_key is a valid ed25519 key, ``False`` otherwise.
63pub fn are_bytes_valid_ed25519_pubkey(public_key: &[u8]) -> bool {
64    if public_key.len() != 32 {
65        return false;
66    }
67
68    let pub_key_var = Some(hex::encode(public_key));
69    let keypair_result = Keypair::new(None, pub_key_var, None, SS58_FORMAT, None, 1);
70
71    match keypair_result {
72        Ok(keypair) => keypair.ss58_address().is_some(),
73        Err(_) => false,
74    }
75}
76
77///    Checks if the given address is a valid destination address.
78///
79///     Arguments:
80///         address (str): The address to check.
81///     Returns:
82///         valid (bool): ``True`` if the address is a valid destination address, ``False`` otherwise.
83pub fn is_valid_bittensor_address_or_public_key(address: &str) -> bool {
84    if address.starts_with("0x") {
85        // Convert hex string to bytes
86        if let Ok(bytes) = hex::decode(&address[2..]) {
87            are_bytes_valid_ed25519_pubkey(&bytes)
88        } else {
89            is_valid_ss58_address(address)
90        }
91    } else {
92        is_valid_ss58_address(address)
93    }
94}
95
96/// When running inside a Python process, output must go through Python's `sys.stdout`
97/// to stay synchronized with Python-level I/O (e.g. Jupyter, logging redirects).
98/// In pure Rust mode, standard stdout is used directly.
99#[cfg(not(feature = "python-bindings"))]
100pub fn print(s: String) {
101    use std::io::{self, Write};
102    print!("{}", s);
103    io::stdout().flush().unwrap();
104}
105
106#[cfg(feature = "python-bindings")]
107pub fn print(s: String) {
108    use pyo3::types::{PyDict, PyDictMethods};
109    use std::ffi::CString;
110    pyo3::Python::attach(|py| {
111        let locals = PyDict::new(py);
112        locals.set_item("s", s).unwrap();
113        let code = CString::new(
114            r#"
115import sys
116print(s, end='')
117sys.stdout.flush()
118"#,
119        )
120        .unwrap();
121        py.run(&code, None, Some(&locals)).unwrap();
122    });
123}
124
125/// Prompts the user and returns the response, if any.
126///    
127///     Arguments:
128///         prompt: String
129///     Returns:
130///         response: Option<String>
131pub fn prompt(prompt: String) -> Option<String> {
132    use std::io::{self, Write};
133
134    print!("{}", prompt);
135    io::stdout().flush().ok()?;
136
137    let mut input = String::new();
138    match io::stdin().read_line(&mut input) {
139        Ok(_) => Some(input.trim().to_string()),
140        Err(_) => None,
141    }
142}
143
144/// Prompts the user with a password entry and returns the response, if any.
145///    
146///     Arguments:
147///         prompt (String): the prompt to ask the user with.
148///     Returns:
149///         response: Option<String>
150pub fn prompt_password(prompt: String) -> Option<String> {
151    use rpassword::read_password;
152    use std::io::{self, Write};
153
154    print!("{}", prompt);
155    io::stdout().flush().ok()?;
156
157    match read_password() {
158        Ok(password) => Some(password.trim().to_string()),
159        Err(_) => None,
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn test_get_ss58_format_success() {
169        let test_address = "5FHneW46xGXgs5mUiveU4sbTyGBzmstUspZC92UhjJM694ty";
170        assert!(is_valid_ss58_address(test_address));
171    }
172}