1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use base58::ToBase58;
use bitcoin_hashes::sha256d::Hash as SHA256d;
use bitcoin_hashes::Hash;
use strum_macros::{Display, EnumString};

use super::{AddressOptions, Error, Result};

/// Options for formatting a private key as WIF.
///
/// Currently, this just specifies the blockchain network for which the WIF is intended to be used
/// because for each blockchain network there is exactly one supported way to serialize the private
/// key as a string.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, Display)]
pub enum WIFOptions {
    Bitcoin = 0x80,
    Litecoin = 0xB0,
    Dash = 0xCC,
    Dogecoin = 0x9E,
    Europecoin = 0xA8,
    Peercoin = 0xB7,
}

impl From<&AddressOptions> for WIFOptions {
    /// Obtain a `WIFOptions` using the same blockchain network as an AddressOptions
    fn from(opts: &AddressOptions) -> Self {
        match opts {
            AddressOptions::Bitcoin => Self::Bitcoin,
            AddressOptions::BitcoinGold => Self::Bitcoin,
            AddressOptions::BitcoinCash(_) => Self::Bitcoin,
            AddressOptions::Litecoin => Self::Litecoin,
            AddressOptions::Dogecoin => Self::Dogecoin,
            AddressOptions::Dash => Self::Dash,
            AddressOptions::Europecoin => Self::Europecoin,
            AddressOptions::Peercoin => Self::Peercoin,
            _ => panic!("could not Convert {:?} to Network", opts),
        }
    }
}

impl WIFOptions {
    /// Encode private key bytes as a WIF string
    pub(crate) fn encode(&self, private_key_bytes: impl AsRef<[u8]>) -> Result<String> {
        // First check that the input slice is exactly 32 bytes long;
        let private_key_slice = private_key_bytes.as_ref();
        if private_key_slice.len() != 32 {
            return Err(Error::InvalidLength {
                received: private_key_slice.len(),
                expected: 32,
            });
        }
        // allocated on the stack, no biggie
        let mut buf = [0; 38];
        buf[0] = *self as u8;
        buf[1..33].copy_from_slice(private_key_slice);
        buf[33] = 1;

        let checksum = SHA256d::hash(&buf[..34]);
        buf[34..].copy_from_slice(&checksum[..4]);
        Ok(buf.to_base58())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_vectors::wif::*;

    fn testfun(vectors: &str, opts: WIFOptions) {
        for tc in vectors.lines().map(TestCase::from) {
            let wif = opts.encode(tc.prv);
            assert_eq!(wif.as_deref(), Ok(tc.wif));
        }
    }

    #[test]
    fn btc_wif() {
        testfun(BTC_TEST_VECTORS, WIFOptions::Bitcoin)
    }
    #[test]
    fn ltc_wif() {
        testfun(LTC_TEST_VECTORS, WIFOptions::Litecoin)
    }
    #[test]
    fn doge_wif() {
        testfun(DOGE_TEST_VECTORS, WIFOptions::Dogecoin)
    }
    #[test]
    fn dash_wif() {
        testfun(DASH_TEST_VECTORS, WIFOptions::Dash)
    }
    #[test]
    fn erc_wif() {
        testfun(ERC_TEST_VECTORS, WIFOptions::Europecoin)
    }

    #[test]
    fn from_addr_opts() {
        assert_eq!(
            WIFOptions::from(&AddressOptions::Bitcoin),
            WIFOptions::Bitcoin
        );
        assert_eq!(
            WIFOptions::from(&AddressOptions::BitcoinCash(None)),
            WIFOptions::Bitcoin
        );
        assert_eq!(
            WIFOptions::from(&AddressOptions::BitcoinCash(Some("ligma".to_owned()))),
            WIFOptions::Bitcoin
        );
        assert_eq!(
            WIFOptions::from(&AddressOptions::Litecoin),
            WIFOptions::Litecoin
        );
        assert_eq!(
            WIFOptions::from(&AddressOptions::Dogecoin),
            WIFOptions::Dogecoin
        );
        assert_eq!(WIFOptions::from(&AddressOptions::Dash), WIFOptions::Dash);
        assert_eq!(
            WIFOptions::from(&AddressOptions::Europecoin),
            WIFOptions::Europecoin
        );
        assert_eq!(
            WIFOptions::from(&AddressOptions::Peercoin),
            WIFOptions::Peercoin
        );
    }
}