1use crate::error::BittensorError;
11use serde::{Deserialize, Serialize};
12use std::path::Path;
13use sp_core::{sr25519, Pair};
14use thiserror::Error;
15
16#[derive(Debug, Error)]
18pub enum KeyfileError {
19 #[error("Failed to read keyfile: {0}")]
21 ReadError(#[from] std::io::Error),
22
23 #[error("Failed to parse keyfile JSON: {0}")]
25 ParseError(String),
26
27 #[error("Decryption failed: {0}")]
29 DecryptionError(String),
30
31 #[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#[derive(Debug, Clone)]
46pub struct KeyfileData {
47 pub secret: String,
49 pub is_mnemonic: bool,
51 pub format: KeyfileFormat,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum KeyfileFormat {
58 JsonSecretPhrase,
60 JsonSecretSeed,
62 PlainMnemonic,
64 PlainHexSeed,
66 Encrypted,
68}
69
70#[derive(Debug, Deserialize, Serialize)]
72#[serde(rename_all = "camelCase")]
73struct JsonKeyfile {
74 secret_phrase: Option<String>,
76 secret_seed: Option<String>,
78 public_key: Option<String>,
80 account_id: Option<String>,
82 ss58_address: Option<String>,
84}
85
86impl KeyfileData {
87 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 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
115pub fn load_keyfile(path: &Path) -> Result<KeyfileData, KeyfileError> {
127 let content = std::fs::read_to_string(path)?;
128 parse_keyfile_content(&content)
129}
130
131pub 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
148fn parse_keyfile_content(content: &str) -> Result<KeyfileData, KeyfileError> {
150 let trimmed = content.trim();
151
152 if let Ok(json_keyfile) = serde_json::from_str::<JsonKeyfile>(trimmed) {
154 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 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 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 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 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 Ok(KeyfileData {
208 secret: trimmed.to_string(),
209 is_mnemonic: true,
210 format: KeyfileFormat::PlainMnemonic,
211 })
212}
213
214fn decrypt_keyfile(data: &[u8], _password: &str) -> Result<KeyfileData, KeyfileError> {
221 if data.len() < 57 {
223 return Err(KeyfileError::DecryptionError(
224 "Encrypted keyfile too short".to_string(),
225 ));
226 }
227
228 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(), 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}