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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

use crate::fragment::FragmentBuilderError;
use chain_addr::Discrimination;
use chain_impl_mockchain::{
    account::SpendingCounter,
    accounting::account::SpendingCounterIncreasing,
    fee::{FeeAlgorithm, LinearFee},
    transaction::{
        Balance, Input, InputOutputBuilder, Payload, PayloadSlice, TransactionSignDataHash,
        UnspecifiedAccountIdentifier, Witness,
    },
};
use jormungandr_lib::{
    crypto::{
        account::{Identifier, SigningKey},
        hash::Hash,
    },
    interfaces::{Address, Value},
};
use rand_core::{CryptoRng, RngCore};

/// wallet for an account
#[derive(Debug, Clone)]
pub struct Wallet {
    /// the spending key
    signing_key: SigningKey,

    /// the counter as we know of this value needs to be in sync
    /// with what is in the blockchain
    internal_counters: SpendingCounterIncreasing,

    discrimination: Discrimination,
}

impl Wallet {
    pub fn generate<RNG>(rng: &mut RNG, discrimination: Discrimination) -> Self
    where
        RNG: CryptoRng + RngCore,
    {
        let signing_key = SigningKey::generate_extended(rng);
        Wallet {
            signing_key,
            internal_counters: SpendingCounterIncreasing::default(),
            discrimination,
        }
    }

    pub fn from_secret_key(
        signing_key: SigningKey,
        internal_counters: SpendingCounterIncreasing,
        discrimination: Discrimination,
    ) -> Self {
        Wallet {
            signing_key,
            internal_counters,
            discrimination,
        }
    }

    pub fn from_existing_account(
        bech32_str: &str,
        spending_counter: Option<SpendingCounter>,
        discrimination: Discrimination,
    ) -> Self {
        let signing_key = SigningKey::from_bech32_str(bech32_str).expect("bad bech32");
        Wallet {
            signing_key,
            internal_counters: SpendingCounterIncreasing::new_from_counter(
                spending_counter
                    .map(Into::into)
                    .unwrap_or_else(SpendingCounter::zero),
            ),
            discrimination,
        }
    }

    pub fn save_to<W: std::io::Write>(&self, mut w: W) -> std::io::Result<()> {
        writeln!(w, "{}", self.signing_key().to_bech32_str())
    }

    pub fn address(&self) -> Address {
        self.identifier().to_address(self.discrimination).into()
    }

    pub fn set_counter(&mut self, counter: SpendingCounter) {
        let mut counters = self.internal_counters.get_valid_counters();
        counters[counter.lane()] = counter;
        self.internal_counters = SpendingCounterIncreasing::new_from_counters(counters).unwrap();
    }

    pub fn increment_counter(&mut self, lane: usize) {
        self.internal_counters
            .next_verify(self.internal_counters.get_valid_counters()[lane])
            .unwrap();
    }

    pub fn decrement_counter(&mut self, lane: usize) {
        self.set_counter(SpendingCounter::from(
            <u32>::from(self.internal_counters()[lane]) - 1,
        ))
    }

    pub fn spending_counter(&self) -> &SpendingCounterIncreasing {
        &self.internal_counters
    }

    /// Use the default counter
    pub fn internal_counter(&self) -> SpendingCounter {
        self.internal_counters.get_valid_counter()
    }

    pub fn internal_counters(&self) -> [SpendingCounter; SpendingCounterIncreasing::LANES] {
        self.internal_counters.get_valid_counters()
    }

    pub fn stake_key(&self) -> UnspecifiedAccountIdentifier {
        UnspecifiedAccountIdentifier::from_single_account(self.identifier().to_inner())
    }

    pub fn identifier(&self) -> Identifier {
        self.signing_key.identifier()
    }

    pub fn signing_key(&self) -> &SigningKey {
        &self.signing_key
    }

    pub fn mk_witness(
        &self,
        block0_hash: &Hash,
        signing_data: &TransactionSignDataHash,
    ) -> Witness {
        Witness::new_account(
            &(*block0_hash).into_hash(),
            signing_data,
            self.internal_counters.get_valid_counter(),
            |d| self.signing_key().as_ref().sign(d),
        )
    }

    pub fn add_input_with_value(&self, value: Value) -> Input {
        Input::from_account_single(self.identifier().to_inner(), value.into())
    }

    pub fn add_input<Extra: Payload>(
        &self,
        payload: PayloadSlice<'_, Extra>,
        iobuilder: &mut InputOutputBuilder,
        fees: &LinearFee,
    ) -> Result<(), FragmentBuilderError>
    where
        LinearFee: FeeAlgorithm,
    {
        let balance = iobuilder
            .get_balance_with_placeholders(payload, fees, 1, 0)
            .map_err(|_| FragmentBuilderError::CannotComputeBalance)?;
        let value = match balance {
            Balance::Negative(value) => value,
            Balance::Zero => return Err(FragmentBuilderError::TransactionAlreadyBalanced),
            Balance::Positive(value) => {
                return Err(FragmentBuilderError::TransactionAlreadyExtraValue(
                    value.into(),
                ))
            }
        };

        let input = Input::from_account_single(self.identifier().to_inner(), value);
        iobuilder.add_input(&input).unwrap();
        Ok(())
    }
}