1use sha2::{Digest, Sha256};
2
3use crate::crypto::{CryptoError, Secp256r1PrivateKey};
4use neo3::prelude::*;
5
6pub fn private_key_from_wif(wif: &str) -> Result<Secp256r1PrivateKey, CryptoError> {
24 let data = bs58::decode(wif)
25 .into_vec()
26 .map_err(|_| CryptoError::InvalidFormat("Incorrect WIF format.".to_string()))?;
27 if data.len() != 38 || data[0] != 0x80 || data[33] != 0x01 {
28 return Err(CryptoError::InvalidFormat("Incorrect WIF format.".to_string()));
29 }
30
31 let checksum_calculated = Sha256::digest(&Sha256::digest(&data[..34]));
32 if checksum_calculated[..4] != data[34..] {
33 return Err(CryptoError::InvalidFormat("Incorrect WIF checksum.".to_string()));
34 }
35
36 Secp256r1PrivateKey::from_bytes(&data[1..33].to_vec())
37}
38
39pub fn wif_from_private_key(private_key: &Secp256r1PrivateKey) -> String {
50 let mut extended_key: Vec<u8> = vec![0x80];
51 extended_key.extend(private_key.to_raw_bytes());
52 extended_key.push(0x01);
53
54 let hash = Sha256::digest(&Sha256::digest(&extended_key));
55 let checksum = &hash[0..4];
56 extended_key.extend_from_slice(checksum);
57
58 bs58::encode(extended_key).into_string()
59}
60
61#[cfg(test)]
62mod tests {
63 use crate::crypto::{
64 private_key_from_wif, wif_from_private_key, PrivateKeyExtension, Secp256r1PrivateKey,
65 };
66
67 #[test]
68 fn test_valid_wif_to_private_key() {
69 let wif = "L25kgAQJXNHnhc7Sx9bomxxwVSMsZdkaNQ3m2VfHrnLzKWMLP13A";
70 let expected_key =
71 hex::decode("9117f4bf9be717c9a90994326897f4243503accd06712162267e77f18b49c3a3")
72 .unwrap();
73
74 let key = private_key_from_wif(wif).unwrap().to_raw_bytes().to_vec();
75
76 assert_eq!(key, expected_key);
77 }
78
79 #[test]
80 fn test_invalid_wif_sizes() {
81 let too_long = "L25kgAQJXNHnhc7Sx9bomxxwVSMsZdkaNQ3m2VfHrnLzKWMLP13Ahc7S";
82 let too_short = "L25kgAQJXNHnhc7Sx9bomxxwVSMsZdkaNQ3m2VfHrnLzKWML";
83
84 assert!(private_key_from_wif(too_long).is_err());
85 assert!(private_key_from_wif(too_short).is_err());
86 }
87
88 #[test]
89 fn test_invalid_wif_bytes() {
90 let wif = "L25kgAQJXNHnhc7Sx9bomxxwVSMsZdkaNQ3m2VfHrnLzKWMLP13A";
91 let expected_key =
92 hex::decode("9117f4bf9be717c9a90994326897f4243503accd06712162267e77f18b49c3a3")
93 .unwrap();
94
95 let mut decoded = bs58::decode(wif).into_vec().unwrap();
96 decoded[0] = 0x81;
97 let invalid_first = bs58::encode(&decoded).into_string();
98
99 decoded[33] = 0;
100 let invalid_33rd = bs58::encode(&decoded).into_string();
101
102 assert!(private_key_from_wif(invalid_first.as_str()).is_err());
103 assert!(private_key_from_wif(invalid_33rd.as_str()).is_err());
104 }
105
106 #[test]
107 fn test_valid_private_key_to_wif() {
108 let pk = hex::decode("9117f4bf9be717c9a90994326897f4243503accd06712162267e77f18b49c3a3")
109 .unwrap();
110 let expected_wif = "L25kgAQJXNHnhc7Sx9bomxxwVSMsZdkaNQ3m2VfHrnLzKWMLP13A";
111
112 let wif = wif_from_private_key(&Secp256r1PrivateKey::from_slice(&pk).unwrap());
113
114 assert_eq!(wif, expected_wif);
115 }
116
117 #[test]
118 fn test_invalid_private_key_length() {
119 let invalid_len =
120 hex::decode("9117f4bf9be717c9a90994326897f4243503accd06712162267e77f18b49c3").unwrap();
121
122 assert!(Secp256r1PrivateKey::from_slice(&invalid_len).is_err());
124 }
125}