neo3/neo_crypto/
wif.rs

1use sha2::{Digest, Sha256};
2
3use crate::crypto::{CryptoError, Secp256r1PrivateKey};
4use neo3::prelude::*;
5
6/// Converts a WIF (Wallet Import Format) string into a `Secp256r1PrivateKey`.
7///
8/// This function decodes a WIF string, verifies its format and checksum,
9/// and then constructs a `Secp256r1PrivateKey` from it.
10///
11/// # Arguments
12/// * `wif` - A string slice representing the WIF to be converted.
13///
14/// # Returns
15/// A `Result` which is `Ok` with `Secp256r1PrivateKey` if the WIF string is valid,
16/// or an `Err` with `CryptoError` if the WIF string is invalid.
17///
18/// # Errors
19/// This function returns an error if:
20/// * The WIF string is not properly base58 encoded.
21/// * The decoded data does not have the correct length, prefix, or suffix expected for a WIF.
22/// * The checksum of the WIF does not match the expected value.
23pub 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
39/// Converts a `Secp256r1PrivateKey` into a WIF (Wallet Import Format) string.
40///
41/// This function takes a `Secp256r1PrivateKey`, converts it to its raw byte representation,
42/// adds the appropriate prefix and suffix, calculates the checksum, and then encodes it into WIF format.
43///
44/// # Arguments
45/// * `private_key` - A reference to the `Secp256r1PrivateKey` to be converted.
46///
47/// # Returns
48/// A `String` containing the WIF representation of the provided private key.
49pub 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		// wif_from_private_key(&
123		assert!(Secp256r1PrivateKey::from_slice(&invalid_len).is_err());
124	}
125}