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
130
131
132
133
134
135
136
137
138
use base58::ToBase58;
use bitcoin_hashes::sha256d::Hash as SHA256d;
use bitcoin_hashes::Hash;
use strum_macros::{Display, EnumString};

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

/// Format specification for formatting a private key as
/// [WIF](https://en.bitcoin.it/wiki/Wallet_import_format).
///
/// 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 WIFFormat {
    Bitcoin = 0x80,
    Litecoin = 0xB0,
    Dash = 0xCC,
    Dogecoin = 0x9E,
    Europecoin = 0xA8,
    Peercoin = 0xB7,
    #[cfg(feature = "eth")]
    Ethereum,
}

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

impl WIFFormat {
    /// Encode private key bytes as a WIF string
    pub(crate) fn encode(&self, private_key_bytes: impl AsRef<[u8]>) -> Result<String> {
        #[cfg(feature = "eth")]
        if let Self::Ethereum = self {
            return Ok(hex::encode(private_key_bytes));
        }
        // 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: WIFFormat) {
        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, WIFFormat::Bitcoin)
    }
    #[test]
    fn ltc_wif() {
        testfun(LTC_TEST_VECTORS, WIFFormat::Litecoin)
    }
    #[test]
    fn doge_wif() {
        testfun(DOGE_TEST_VECTORS, WIFFormat::Dogecoin)
    }
    #[test]
    fn dash_wif() {
        testfun(DASH_TEST_VECTORS, WIFFormat::Dash)
    }
    #[test]
    fn erc_wif() {
        testfun(ERC_TEST_VECTORS, WIFFormat::Europecoin)
    }

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