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_merkle_tree::merkle_proof,
21
    sp_runtime::{traits::Keccak256, AccountId32},
22
    std::{
23
        collections::BTreeMap,
24
        path::{Path, PathBuf},
25
    },
26
    xcm_payload::PayloadGeneratorCmd,
27
};
28

            
29
pub mod xcm_payload;
30

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

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

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

            
51
#[derive(Debug, Subcommand)]
52
#[command(rename_all = "kebab-case")]
53
pub enum TanssiUtilsCmd {
54
    RewardClaimGenerator(RewardClaimGeneratorCmd),
55
    PayloadGenerator(PayloadGeneratorCmd),
56
}
57

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

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

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

            
96
/// Extract a set of rewards information from a JSON file.
97
fn extract_rewards_data_from_file(reward_path: &Path) -> Result<RewardClaimInput, String> {
98
    let reader = std::fs::File::open(reward_path).expect("Can open file");
99
    let reward_input = serde_json::from_reader(&reader).expect("Cant parse reward input from JSON");
100
    Ok(reward_input)
101
}
102

            
103
fn generate_reward_utils(reward_input: RewardClaimInput) {
104
    let era_index = reward_input.era;
105
    let mut total_points = 0;
106
    let individual_rewards: BTreeMap<_, _> = reward_input
107
        .operator_rewards
108
        .clone()
109
        .into_iter()
110
        .map(|data| {
111
            total_points += data.amount;
112
            (data.account, data.amount)
113
        })
114
        .collect();
115
    let era_rewards = pallet_external_validators_rewards::EraRewardPoints::<AccountId32> {
116
        total: total_points,
117
        individual: individual_rewards,
118
    };
119

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

            
137
                println!("\n=== Merkle Proofs ===");
138
            }
139

            
140
            let merkle_proof = account_utils
141
                .leaf_index
142
                .map(|index| merkle_proof::<Keccak256, _>(account_utils.leaves.into_iter(), index));
143

            
144
            if let Some(proof) = merkle_proof {
145
                println!(
146
                    "\nMerkle proof for account {:?} in era {:?}: \n",
147
                    reward.account, era_index
148
                );
149
                println!("   - Root: {:?}", proof.root);
150
                println!("   - Proof: {:?}", proof.proof);
151
                println!("   - Number of leaves: {:?}", proof.number_of_leaves);
152
                println!("   - Leaf index: {:?}", proof.leaf_index);
153
                println!("   - Leaf: {:?}", proof.leaf);
154
            } else {
155
                println!("No proof generated for account {:?}", reward.account);
156
            }
157
        } else {
158
            println!("No utils generated for account {:?}", reward.account);
159
        };
160
    });
161
}
162

            
163
fn main() {
164
    // Parses the options
165
    let cmd = TanssiUtils::parse();
166
    cmd.run();
167
}