bittensor_rs/wallet/
keyfile.rs

1//! # Keyfile Handling
2//!
3//! Utilities for loading and parsing Bittensor keyfiles.
4//!
5//! Bittensor keyfiles can be in several formats:
6//! - JSON format with `secretPhrase` or `secretSeed` fields
7//! - Plain text mnemonic phrase
8//! - Encrypted keyfiles (for coldkeys)
9
10use crate::error::BittensorError;
11use serde::{Deserialize, Serialize};
12use std::path::Path;
13use sp_core::{sr25519, Pair};
14use thiserror::Error;
15
16/// Errors that can occur when loading keyfiles
17#[derive(Debug, Error)]
18pub enum KeyfileError {
19    /// File could not be read
20    #[error("Failed to read keyfile: {0}")]
21    ReadError(#[from] std::io::Error),
22
23    /// JSON parsing failed
24    #[error("Failed to parse keyfile JSON: {0}")]
25    ParseError(String),
26
27    /// Decryption failed
28    #[error("Decryption failed: {0}")]
29    DecryptionError(String),
30
31    /// Invalid key format
32    #[error("Invalid key format: {0}")]
33    InvalidFormat(String),
34}
35
36impl From<KeyfileError> for BittensorError {
37    fn from(err: KeyfileError) -> Self {
38        BittensorError::WalletError {
39            message: err.to_string(),
40        }
41    }
42}
43
44/// Parsed keyfile data
45#[derive(Debug, Clone)]
46pub struct KeyfileData {
47    /// The secret seed or mnemonic phrase
48    pub secret: String,
49    /// Whether this is a mnemonic phrase
50    pub is_mnemonic: bool,
51    /// Original format of the keyfile
52    pub format: KeyfileFormat,
53}
54
55/// Format of the keyfile
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum KeyfileFormat {
58    /// JSON with secretPhrase field
59    JsonSecretPhrase,
60    /// JSON with secretSeed field (hex)
61    JsonSecretSeed,
62    /// Plain text mnemonic
63    PlainMnemonic,
64    /// Plain text hex seed
65    PlainHexSeed,
66    /// Encrypted keyfile
67    Encrypted,
68}
69
70/// JSON structure for Bittensor keyfiles
71#[derive(Debug, Deserialize, Serialize)]
72#[serde(rename_all = "camelCase")]
73struct JsonKeyfile {
74    /// Mnemonic phrase
75    secret_phrase: Option<String>,
76    /// Hex-encoded seed
77    secret_seed: Option<String>,
78    /// Public key (SS58)
79    public_key: Option<String>,
80    /// Account ID
81    account_id: Option<String>,
82    /// SS58 address
83    ss58_address: Option<String>,
84}
85
86impl KeyfileData {
87    /// Convert the keyfile data to an sr25519 keypair
88    pub fn to_keypair(&self) -> Result<sr25519::Pair, BittensorError> {
89        if self.is_mnemonic {
90            sr25519::Pair::from_string(&self.secret, None).map_err(|e| {
91                BittensorError::WalletError {
92                    message: format!("Invalid mnemonic phrase: {e:?}"),
93                }
94            })
95        } else {
96            // It's a hex seed
97            let hex_str = self.secret.strip_prefix("0x").unwrap_or(&self.secret);
98            let seed_bytes = hex::decode(hex_str).map_err(|e| BittensorError::WalletError {
99                message: format!("Invalid hex seed: {e}"),
100            })?;
101
102            if seed_bytes.len() != 32 {
103                return Err(BittensorError::WalletError {
104                    message: format!("Seed must be 32 bytes, got {} bytes", seed_bytes.len()),
105                });
106            }
107
108            let mut seed_array = [0u8; 32];
109            seed_array.copy_from_slice(&seed_bytes);
110            Ok(sr25519::Pair::from_seed(&seed_array))
111        }
112    }
113}
114
115/// Load a keyfile from disk
116///
117/// Supports both JSON and plain text formats.
118///
119/// # Arguments
120///
121/// * `path` - Path to the keyfile
122///
123/// # Returns
124///
125/// Parsed keyfile data
126pub fn load_keyfile(path: &Path) -> Result<KeyfileData, KeyfileError> {
127    let content = std::fs::read_to_string(path)?;
128    parse_keyfile_content(&content)
129}
130
131/// Load an encrypted keyfile
132///
133/// Currently only supports nacl-encrypted keyfiles as used by the Python SDK.
134///
135/// # Arguments
136///
137/// * `path` - Path to the encrypted keyfile
138/// * `password` - Password for decryption
139///
140/// # Returns
141///
142/// Parsed keyfile data after decryption
143pub fn load_encrypted_keyfile(path: &Path, password: &str) -> Result<KeyfileData, KeyfileError> {
144    let content = std::fs::read(path)?;
145    decrypt_keyfile(&content, password)
146}
147
148/// Parse keyfile content
149fn parse_keyfile_content(content: &str) -> Result<KeyfileData, KeyfileError> {
150    let trimmed = content.trim();
151
152    // Try to parse as JSON first
153    if let Ok(json_keyfile) = serde_json::from_str::<JsonKeyfile>(trimmed) {
154        // Check for secretPhrase (mnemonic)
155        if let Some(phrase) = json_keyfile.secret_phrase {
156            return Ok(KeyfileData {
157                secret: phrase,
158                is_mnemonic: true,
159                format: KeyfileFormat::JsonSecretPhrase,
160            });
161        }
162
163        // Check for secretSeed (hex)
164        if let Some(seed) = json_keyfile.secret_seed {
165            return Ok(KeyfileData {
166                secret: seed,
167                is_mnemonic: false,
168                format: KeyfileFormat::JsonSecretSeed,
169            });
170        }
171
172        return Err(KeyfileError::InvalidFormat(
173            "JSON keyfile missing secretPhrase or secretSeed".to_string(),
174        ));
175    }
176
177    // Not JSON - try to parse as plain text
178    // Check if it's a hex seed (starts with 0x and is 66 chars)
179    if trimmed.starts_with("0x") && trimmed.len() == 66 {
180        return Ok(KeyfileData {
181            secret: trimmed.to_string(),
182            is_mnemonic: false,
183            format: KeyfileFormat::PlainHexSeed,
184        });
185    }
186
187    // Check if it looks like a hex seed without 0x prefix
188    if trimmed.len() == 64 && trimmed.chars().all(|c| c.is_ascii_hexdigit()) {
189        return Ok(KeyfileData {
190            secret: format!("0x{}", trimmed),
191            is_mnemonic: false,
192            format: KeyfileFormat::PlainHexSeed,
193        });
194    }
195
196    // Assume it's a mnemonic phrase
197    let word_count = trimmed.split_whitespace().count();
198    if word_count == 12 || word_count == 24 {
199        return Ok(KeyfileData {
200            secret: trimmed.to_string(),
201            is_mnemonic: true,
202            format: KeyfileFormat::PlainMnemonic,
203        });
204    }
205
206    // Still treat as mnemonic but warn
207    Ok(KeyfileData {
208        secret: trimmed.to_string(),
209        is_mnemonic: true,
210        format: KeyfileFormat::PlainMnemonic,
211    })
212}
213
214/// Decrypt an encrypted keyfile
215///
216/// Bittensor uses NaCl secretbox for encryption with:
217/// - Salt: first 16 bytes
218/// - Nonce: next 24 bytes
219/// - Ciphertext: remaining bytes
220fn decrypt_keyfile(data: &[u8], _password: &str) -> Result<KeyfileData, KeyfileError> {
221    // Check minimum length: 16 (salt) + 24 (nonce) + 16 (auth tag) + 1 (min data)
222    if data.len() < 57 {
223        return Err(KeyfileError::DecryptionError(
224            "Encrypted keyfile too short".to_string(),
225        ));
226    }
227
228    // For now, we return an error indicating encrypted keyfiles need the Python SDK
229    // TODO: Implement nacl decryption or use sodiumoxide crate
230    Err(KeyfileError::DecryptionError(
231        "Encrypted coldkey decryption not yet implemented. \
232         Please use `btcli wallet regen_coldkey` to create an unencrypted coldkey, \
233         or decrypt using the Python bittensor SDK."
234            .to_string(),
235    ))
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_parse_json_mnemonic() {
244        let content = r#"{"secretPhrase": "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"}"#;
245        let result = parse_keyfile_content(content).unwrap();
246        assert!(result.is_mnemonic);
247        assert_eq!(result.format, KeyfileFormat::JsonSecretPhrase);
248        assert!(result.secret.contains("abandon"));
249    }
250
251    #[test]
252    fn test_parse_json_seed() {
253        let content = r#"{"secretSeed": "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"}"#;
254        let result = parse_keyfile_content(content).unwrap();
255        assert!(!result.is_mnemonic);
256        assert_eq!(result.format, KeyfileFormat::JsonSecretSeed);
257    }
258
259    #[test]
260    fn test_parse_plain_mnemonic() {
261        let content = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about";
262        let result = parse_keyfile_content(content).unwrap();
263        assert!(result.is_mnemonic);
264        assert_eq!(result.format, KeyfileFormat::PlainMnemonic);
265    }
266
267    #[test]
268    fn test_parse_plain_hex_with_prefix() {
269        let content = "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
270        let result = parse_keyfile_content(content).unwrap();
271        assert!(!result.is_mnemonic);
272        assert_eq!(result.format, KeyfileFormat::PlainHexSeed);
273    }
274
275    #[test]
276    fn test_parse_plain_hex_without_prefix() {
277        let content = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
278        let result = parse_keyfile_content(content).unwrap();
279        assert!(!result.is_mnemonic);
280        assert_eq!(result.format, KeyfileFormat::PlainHexSeed);
281        assert!(result.secret.starts_with("0x"));
282    }
283
284    #[test]
285    fn test_to_keypair_from_mnemonic() {
286        let data = KeyfileData {
287            secret: "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about".to_string(),
288            is_mnemonic: true,
289            format: KeyfileFormat::PlainMnemonic,
290        };
291        let result = data.to_keypair();
292        assert!(result.is_ok());
293    }
294
295    #[test]
296    fn test_to_keypair_from_seed() {
297        let data = KeyfileData {
298            secret: "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
299                .to_string(),
300            is_mnemonic: false,
301            format: KeyfileFormat::PlainHexSeed,
302        };
303        let result = data.to_keypair();
304        assert!(result.is_ok());
305    }
306
307    #[test]
308    fn test_to_keypair_invalid_mnemonic() {
309        let data = KeyfileData {
310            secret: "invalid mnemonic phrase".to_string(),
311            is_mnemonic: true,
312            format: KeyfileFormat::PlainMnemonic,
313        };
314        let result = data.to_keypair();
315        assert!(result.is_err());
316    }
317
318    #[test]
319    fn test_to_keypair_invalid_hex() {
320        let data = KeyfileData {
321            secret: "0xNOTHEX".to_string(),
322            is_mnemonic: false,
323            format: KeyfileFormat::PlainHexSeed,
324        };
325        let result = data.to_keypair();
326        assert!(result.is_err());
327    }
328
329    #[test]
330    fn test_to_keypair_wrong_seed_length() {
331        let data = KeyfileData {
332            secret: "0x0123456789abcdef".to_string(), // 16 bytes instead of 32
333            is_mnemonic: false,
334            format: KeyfileFormat::PlainHexSeed,
335        };
336        let result = data.to_keypair();
337        assert!(result.is_err());
338        if let Err(BittensorError::WalletError { message }) = result {
339            assert!(message.contains("32 bytes"));
340        }
341    }
342
343    #[test]
344    fn test_json_missing_secret() {
345        let content = r#"{"publicKey": "something"}"#;
346        let result = parse_keyfile_content(content);
347        assert!(result.is_err());
348    }
349
350    #[test]
351    fn test_decrypt_too_short() {
352        let data = vec![0u8; 10];
353        let result = decrypt_keyfile(&data, "password");
354        assert!(result.is_err());
355    }
356
357    #[test]
358    fn test_parse_24_word_mnemonic() {
359        let content = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";
360        let result = parse_keyfile_content(content).unwrap();
361        assert!(result.is_mnemonic);
362        assert_eq!(result.format, KeyfileFormat::PlainMnemonic);
363    }
364
365    #[test]
366    fn test_keyfile_error_display() {
367        let err = KeyfileError::ParseError("test".to_string());
368        assert!(err.to_string().contains("parse"));
369
370        let err = KeyfileError::DecryptionError("failed".to_string());
371        assert!(err.to_string().contains("failed"));
372
373        let err = KeyfileError::InvalidFormat("bad".to_string());
374        assert!(err.to_string().contains("bad"));
375    }
376
377    #[test]
378    fn test_keyfile_error_to_bittensor_error() {
379        let err: BittensorError = KeyfileError::ParseError("test".to_string()).into();
380        if let BittensorError::WalletError { message } = err {
381            assert!(message.contains("parse"));
382        } else {
383            panic!("Expected WalletError");
384        }
385    }
386
387    #[test]
388    fn test_keyfile_data_clone() {
389        let data = KeyfileData {
390            secret: "test".to_string(),
391            is_mnemonic: true,
392            format: KeyfileFormat::PlainMnemonic,
393        };
394        let cloned = data.clone();
395        assert_eq!(data.secret, cloned.secret);
396        assert_eq!(data.is_mnemonic, cloned.is_mnemonic);
397    }
398
399    #[test]
400    fn test_keyfile_format_equality() {
401        assert_eq!(KeyfileFormat::PlainMnemonic, KeyfileFormat::PlainMnemonic);
402        assert_ne!(KeyfileFormat::PlainMnemonic, KeyfileFormat::PlainHexSeed);
403    }
404
405    #[test]
406    fn test_parse_whitespace_content() {
407        let content = "  \n  abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about  \n  ";
408        let result = parse_keyfile_content(content).unwrap();
409        assert!(result.is_mnemonic);
410    }
411}