1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
16

            
17
use {
18
    clap::{Parser, Subcommand},
19
    serde::{Deserialize, Deserializer},
20
    snowbridge_outbound_queue_merkle_tree::merkle_proof,
21
    sp_runtime::{traits::Keccak256, AccountId32},
22
    std::{
23
        collections::BTreeMap,
24
        path::{Path, PathBuf},
25
    },
26
};
27

            
28
#[derive(Deserialize, Debug, Clone)]
29
pub struct RewardData {
30
    #[serde(deserialize_with = "hex_to_account_id32")]
31
    account: AccountId32,
32
    amount: u32,
33
}
34

            
35
#[derive(Deserialize, Debug, Clone)]
36
pub struct RewardClaimInput {
37
    pub(crate) operator_rewards: Vec<RewardData>,
38
    pub(crate) era: u32,
39
}
40

            
41
#[derive(Debug, Parser)]
42
#[command(rename_all = "kebab-case", version, about)]
43
pub struct TanssiUtils {
44
    #[command(subcommand)]
45
    pub command: TanssiUtilsCmd,
46
}
47

            
48
#[derive(Debug, Subcommand)]
49
#[command(rename_all = "kebab-case")]
50
pub enum TanssiUtilsCmd {
51
    RewardClaimGenerator(RewardClaimGeneratorCmd),
52
}
53

            
54
#[derive(Parser, Debug)]
55
pub struct RewardClaimGeneratorCmd {
56
    /// The path where the json containing the values is located.
57
    #[arg(long, short)]
58
    pub input_path: PathBuf,
59
}
60

            
61
impl TanssiUtils {
62
    /// Executes the internal command.
63
    pub fn run(&self) {
64
        match &self.command {
65
            TanssiUtilsCmd::RewardClaimGenerator(cmd) => {
66
                println!("\nInput path is: {:?}\n", cmd.input_path);
67
                let rewards =
68
                    extract_rewards_data_from_file(&cmd.input_path).expect("command fail");
69
                generate_reward_utils(rewards)
70
            }
71
        }
72
    }
73
}
74

            
75
// Helper function to deserialize hex strings into AccountId32.
76
// Example: 0x040404...
77
fn hex_to_account_id32<'de, D>(deserializer: D) -> Result<AccountId32, D::Error>
78
where
79
    D: Deserializer<'de>,
80
{
81
    let hex_str: String = Deserialize::deserialize(deserializer)?;
82
    let hex_trimmed = hex_str.strip_prefix("0x").unwrap_or(hex_str.as_str());
83
    let bytes = hex::decode(hex_trimmed).map_err(serde::de::Error::custom)?;
84
    let mut array = [0u8; 32];
85
    array.copy_from_slice(&bytes);
86
    Ok(AccountId32::from(array))
87
}
88

            
89
/// Extract a set of rewards information from a JSON file.
90
fn extract_rewards_data_from_file(reward_path: &Path) -> Result<RewardClaimInput, String> {
91
    let reader = std::fs::File::open(reward_path).expect("Can open file");
92
    let reward_input = serde_json::from_reader(&reader).expect("Cant parse reward input from JSON");
93
    Ok(reward_input)
94
}
95

            
96
fn generate_reward_utils(reward_input: RewardClaimInput) {
97
    let era_index = reward_input.era;
98
    let mut total_points = 0;
99
    let individual_rewards: BTreeMap<_, _> = reward_input
100
        .operator_rewards
101
        .clone()
102
        .into_iter()
103
        .map(|data| {
104
            total_points += data.amount;
105
            (data.account, data.amount)
106
        })
107
        .collect();
108
    let era_rewards = pallet_external_validators_rewards::EraRewardPoints::<AccountId32> {
109
        total: total_points,
110
        individual: individual_rewards,
111
    };
112

            
113
    let mut show_general_info = true;
114
    reward_input.operator_rewards.iter().for_each(|reward| {
115
        if let Some(account_utils) = era_rewards
116
            .generate_era_rewards_utils::<Keccak256>(era_index, Some(reward.account.clone()))
117
        {
118
            // Only show the general info once
119
            if show_general_info {
120
                println!("=== Era Rewards Utils: Overall info ===\n");
121
                println!("Era index       : {:?}", era_index);
122
                println!("Merkle Root     : {:?}", account_utils.rewards_merkle_root);
123
                println!("Total Points    : {}", account_utils.total_points);
124
                println!("Leaves:");
125
                for (i, leaf) in account_utils.leaves.iter().enumerate() {
126
                    println!("  [{}] {:?}", i, leaf);
127
                }
128
                show_general_info = false;
129

            
130
                println!("\n=== Merkle Proofs ===");
131
            }
132

            
133
            let merkle_proof = account_utils
134
                .leaf_index
135
                .map(|index| merkle_proof::<Keccak256, _>(account_utils.leaves.into_iter(), index));
136

            
137
            if let Some(proof) = merkle_proof {
138
                println!(
139
                    "\nMerkle proof for account {:?} in era {:?}: \n",
140
                    reward.account, era_index
141
                );
142
                println!("   - Root: {:?}", proof.root);
143
                println!("   - Proof: {:?}", proof.proof);
144
                println!("   - Number of leaves: {:?}", proof.number_of_leaves);
145
                println!("   - Leaf index: {:?}", proof.leaf_index);
146
                println!("   - Leaf: {:?}", proof.leaf);
147
            } else {
148
                println!("No proof generated for account {:?}", reward.account);
149
            }
150
        } else {
151
            println!("No utils generated for account {:?}", reward.account);
152
        };
153
    });
154
}
155

            
156
fn main() {
157
    // Parses the options
158
    let cmd = TanssiUtils::parse();
159
    cmd.run();
160
}