dfinity · NikolaMilosa · Aug 19, 2024 · Aug 19, 2024 · Aug 19, 2024 · Aug 19, 2024
@@ -11,6 +11,7 @@ use get::Get;
 use heal::Heal;
 use hostos::HostOs;
 use ic_management_types::{MinNakamotoCoefficients, Network, NodeFeature};
+use neuron::Neuron;
 use node_metrics::NodeMetrics;
 use nodes::Nodes;
 use proposals::Proposals;
@@ -24,7 +25,7 @@ use url::Url;
 use version::Version;
 use vote::Vote;
 
-use crate::auth::Neuron;
+use crate::auth::Neuron as AuthNeuron;
 
 mod api_boundary_nodes;
 mod completions;
@@ -33,6 +34,7 @@ mod firewall;
 pub mod get;
 mod heal;
 pub mod hostos;
+mod neuron;
 mod node_metrics;
 mod nodes;
 mod proposals;
@@ -144,7 +146,7 @@ macro_rules! impl_executable_command_for_enums {
 }
 pub(crate) use impl_executable_command_for_enums;
 
-impl_executable_command_for_enums! { DerToPrincipal, Heal, Subnet, Get, Propose, UpdateUnassignedNodes, Version, NodeMetrics, HostOs, Nodes, ApiBoundaryNodes, Vote, Registry, Firewall, Upgrade, Proposals, Completions, Qualify, UpdateAuthorizedSubnets }
+impl_executable_command_for_enums! { DerToPrincipal, Heal, Subnet, Get, Propose, UpdateUnassignedNodes, Version, NodeMetrics, HostOs, Nodes, ApiBoundaryNodes, Vote, Registry, Firewall, Upgrade, Proposals, Completions, Qualify, UpdateAuthorizedSubnets, Neuron }
 
 pub trait ExecutableCommand {
     fn require_ic_admin(&self) -> IcAdminRequirement;
@@ -224,9 +226,9 @@ pub trait ExecutableCommand {
 
 pub enum IcAdminRequirement {
     None,
-    Anonymous,                                          // for get commands
-    Detect,                                             // detect the neuron
-    OverridableBy { network: Network, neuron: Neuron }, // eg automation which we know where is placed
+    Anonymous,                                              // for get commands
+    Detect,                                                 // detect the neuron
+    OverridableBy { network: Network, neuron: AuthNeuron }, // eg automation which we know where is placed
 }
 
 impl ExecutableCommand for Args {

@@ -0,0 +1,32 @@
+use clap::Args;
+use ic_canisters::governance::GovernanceCanisterWrapper;
+
+use crate::commands::ExecutableCommand;
+
+#[derive(Args, Debug)]
+pub struct Balance {
+    /// Neuron to query, by default will use the one from configured identity
+    #[clap(long)]
+    neuron: Option<u64>,
+}
+
+impl ExecutableCommand for Balance {
+    fn require_ic_admin(&self) -> crate::commands::IcAdminRequirement {
+        match &self.neuron {
+            Some(_) => crate::commands::IcAdminRequirement::None,
+            None => crate::commands::IcAdminRequirement::Detect,
+        }
+    }
+
+    fn validate(&self, _cmd: &mut clap::Command) {}
+
+    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
+        let governance = GovernanceCanisterWrapper::from(ctx.create_canister_client()?);
+        let neuron_info = governance
+            .get_neuron_info(self.neuron.unwrap_or_else(|| ctx.ic_admin().neuron.neuron_id))
+            .await?;
+
+        println!("{}", neuron_info.stake_e8s / 10_u64.pow(8));
+        Ok(())
+    }
+}
@@ -0,0 +1,32 @@
+use clap::Args;
+
+use super::{impl_executable_command_for_enums, ExecutableCommand, IcAdminRequirement};
+use crate::commands::neuron::balance::Balance;
+use crate::commands::neuron::refresh::Refresh;
+use crate::commands::neuron::top_up::TopUp;
+
+mod balance;
+mod refresh;
+mod top_up;
+
+#[derive(Args, Debug)]
+pub struct Neuron {
+    #[clap(subcommand)]
+    pub subcommand: Subcommands,
+}
+
+impl_executable_command_for_enums! { Balance, TopUp, Refresh }
+
+impl ExecutableCommand for Neuron {
+    fn require_ic_admin(&self) -> IcAdminRequirement {
+        self.subcommand.require_ic_admin()
+    }
+
+    fn validate(&self, cmd: &mut Command) {
+        self.subcommand.validate(cmd)
+    }
+
+    async fn execute(&self, ctx: DreContext) -> anyhow::Result<()> {
+        self.subcommand.execute(ctx).await
+    }
+}
@@ -0,0 +1,24 @@
+use clap::Args;
+use ic_canisters::governance::GovernanceCanisterWrapper;
+
+use crate::commands::ExecutableCommand;
+
+#[derive(Args, Debug)]
+pub struct Refresh {}
+
+impl ExecutableCommand for Refresh {
+    fn require_ic_admin(&self) -> crate::commands::IcAdminRequirement {
+        crate::commands::IcAdminRequirement::Detect
+    }
+
+    fn validate(&self, _cmd: &mut clap::Command) {}
+
+    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
+        let governance_canister = GovernanceCanisterWrapper::from(ctx.create_canister_client()?);
+
+        let resp = governance_canister.refresh_neuron(ctx.ic_admin().neuron.neuron_id).await?;
+        println!("{:?}", resp);
+
+        Ok(())
+    }
+}
@@ -0,0 +1,32 @@
+use clap::Args;
+use ic_canisters::governance::GovernanceCanisterWrapper;
+use itertools::Itertools;
+
+use crate::commands::ExecutableCommand;
+
+#[derive(Args, Debug)]
+pub struct TopUp {}
+
+impl ExecutableCommand for TopUp {
+    fn require_ic_admin(&self) -> crate::commands::IcAdminRequirement {
+        crate::commands::IcAdminRequirement::Detect
+    }
+
+    fn validate(&self, _cmd: &mut clap::Command) {}
+
+    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
+        let governance = GovernanceCanisterWrapper::from(ctx.create_canister_client()?);
+        let full_neuron = governance.get_full_neuron(ctx.ic_admin().neuron.neuron_id).await?;
+        let account_hex = full_neuron.account.iter().map(|byte| format!("{:02x}", byte)).join("");
+
+        println!("Please request ICP in the #icp-to-go slack channel:");
+        println!(
+            "> Hi! Can I please get XX ICPs on the account address `{}` for neuron ID {} in order to be able to submit more NNS proposals. Thank you\n",
+            account_hex,
+            ctx.ic_admin().neuron.neuron_id
+        );
+        println!("You can check balance by running `dre neuron balance`");
+
+        Ok(())
+    }
+}
@@ -4,11 +4,17 @@ use ic_agent::Agent;
 use ic_nns_common::pb::v1::NeuronId;
 use ic_nns_common::pb::v1::ProposalId;
 use ic_nns_constants::GOVERNANCE_CANISTER_ID;
+use ic_nns_governance::pb::v1::manage_neuron::claim_or_refresh::By;
+use ic_nns_governance::pb::v1::manage_neuron::ClaimOrRefresh;
+use ic_nns_governance::pb::v1::manage_neuron::Command::ClaimOrRefresh as CoR;
 use ic_nns_governance::pb::v1::manage_neuron::RegisterVote;
+use ic_nns_governance::pb::v1::GovernanceError;
 use ic_nns_governance::pb::v1::ListProposalInfo;
 use ic_nns_governance::pb::v1::ListProposalInfoResponse;
 use ic_nns_governance::pb::v1::ManageNeuron;
 use ic_nns_governance::pb::v1::ManageNeuronResponse;
+use ic_nns_governance::pb::v1::Neuron;
+use ic_nns_governance::pb::v1::NeuronInfo;
 use ic_nns_governance::pb::v1::ProposalInfo;
 use log::warn;
 use serde::{self, Serialize};
@@ -162,6 +168,17 @@ impl GovernanceCanisterWrapper {
         }
     }
 
+    pub async fn refresh_neuron(&self, neuron_id: u64) -> anyhow::Result<ManageNeuronResponse> {
+        self.manage_neuron(&ManageNeuron {
+            id: Some(NeuronId { id: neuron_id }),
+            neuron_id_or_subaccount: None,
+            command: Some(CoR(ClaimOrRefresh {
+                by: Some(By::NeuronIdOrSubaccount(ic_nns_governance::pb::v1::Empty {})),
+            })),
+        })
+        .await
+    }
+
     async fn manage_neuron(&self, manage_neuron: &ManageNeuron) -> anyhow::Result<ManageNeuronResponse> {
         match self
             .client
@@ -195,4 +212,32 @@ impl GovernanceCanisterWrapper {
             Err(e) => Err(anyhow::anyhow!("Error executing query: {}", e)),
         }
     }
+
+    pub async fn get_neuron_info(&self, neuron_id: u64) -> anyhow::Result<NeuronInfo> {
+        let args = candid::encode_one(neuron_id)?;
+        match self
+            .client
+            .agent
+            .execute_query(&GOVERNANCE_CANISTER_ID, "get_neuron_info", args)
+            .await
+            .map_err(|e| anyhow::anyhow!(e))?
+        {
+            Some(response) => Ok(Decode!(response.as_slice(), Result<NeuronInfo, GovernanceError>)?.map_err(|e| anyhow::anyhow!(e))?),
+            None => Err(anyhow::anyhow!("Didn't find neuron with id {}", neuron_id)),
+        }
+    }
+
+    pub async fn get_full_neuron(&self, neuron_id: u64) -> anyhow::Result<Neuron> {
+        let args = candid::encode_one(neuron_id)?;
+        match self
+            .client
+            .agent
+            .execute_query(&GOVERNANCE_CANISTER_ID, "get_full_neuron", args)
+            .await
+            .map_err(|e| anyhow::anyhow!(e))?
+        {
+            Some(response) => Ok(Decode!(response.as_slice(), Result<Neuron, GovernanceError>)?.map_err(|e| anyhow::anyhow!(e))?),
+            None => Err(anyhow::anyhow!("Didn't find neuron with id {}", neuron_id)),
+        }
+    }
 }