feat(contract-manager): add script to fetch account balances

Author: ali-behjati

contract_manager/package.json

@@ -26,6 +26,7 @@
     "@pythnetwork/price-service-client": "*",
     "@pythnetwork/pyth-sui-js": "*",
     "@injectivelabs/networks": "1.0.68",
+    "aptos": "^1.5.0",
     "bs58": "^5.0.0",
     "ts-node": "^10.9.1",
     "typescript": "^4.9.3"

contract_manager/scripts/fetch_account_balance.ts

@@ -0,0 +1,59 @@
+import yargs from "yargs";
+import { hideBin } from "yargs/helpers";
+import { DefaultStore, PrivateKey, toPrivateKey } from "../src";
+
+const parser = yargs(hideBin(process.argv))
+  .usage("Usage: $0 --private-key <private-key> [--chain <chain>]")
+  .options({
+    "private-key": {
+      type: "string",
+      demandOption: true,
+      desc: "Private key to use to sign transaction",
+    },
+    chain: {
+      type: "string",
+      desc: "Chain to get the balance for. If not provided the balance for all chains is returned.",
+    },
+  });
+
+type AccountBalance = {
+  chain: string;
+  address: string | undefined;
+  balance: number | undefined;
+};
+
+async function getBalance(
+  chain: string,
+  privateKey: PrivateKey
+): Promise<AccountBalance | undefined> {
+  const address = await DefaultStore.chains[chain].getAccountAddress(
+    privateKey
+  );
+
+  try {
+    const balance = await DefaultStore.chains[chain].getAccountBalance(
+      privateKey
+    );
+    return { chain, address, balance };
+  } catch (e) {
+    console.error(`Error fetching balance for ${chain}`, e);
+  }
+  return { chain, address, balance: undefined };
+}
+
+async function main() {
+  const argv = await parser.argv;
+  const chains = argv.chain
+    ? [argv.chain]
+    : Object.keys(DefaultStore.chains).filter((chain) => chain !== "global");
+
+  const privateKey = toPrivateKey(argv["private-key"]);
+
+  const balances = await Promise.all(
+    chains.map((chain) => getBalance(chain, privateKey))
+  );
+
+  console.table(balances);
+}
+
+main();

contract_manager/src/chains.ts

@@ -12,14 +12,20 @@ import {
   DataSource,
   EvmSetWormholeAddress,
 } from "xc_admin_common";
-import { AptosClient } from "aptos";
+import { AptosClient, AptosAccount, CoinClient } from "aptos";
 import Web3 from "web3";
 import {
   CosmwasmExecutor,
   CosmwasmQuerier,
   InjectiveExecutor,
 } from "@pythnetwork/cosmwasm-deploy-tools";
 import { Network } from "@injectivelabs/networks";
+import {
+  Connection,
+  Ed25519Keypair,
+  JsonRpcProvider,
+  RawSigner,
+} from "@mysten/sui.js";
 
 export type ChainConfig = Record<string, string> & {
   mainnet: boolean;
@@ -96,19 +102,44 @@ export abstract class Chain extends Storable {
    * @param upgradeInfo based on the contract type, this can be a contract address, codeId, package digest, etc.
    */
   abstract generateGovernanceUpgradePayload(upgradeInfo: unknown): Buffer;
+
+  /**
+   * Returns the account address associated with the given private key.
+   * @param privateKey the account private key
+   */
+  abstract getAccountAddress(privateKey: PrivateKey): Promise<string>;
+
+  /**
+   * Returns the balance of the account associated with the given private key.
+   * @param privateKey the account private key
+   */
+  abstract getAccountBalance(privateKey: PrivateKey): Promise<number>;
 }
 
+/**
+ * A Chain object that represents all chains. This is used for governance instructions that apply to all chains.
+ * For example, governance instructions to upgrade Pyth data sources.
+ */
 export class GlobalChain extends Chain {
   static type = "GlobalChain";
   constructor() {
     super("global", true, "unset");
   }
+
   generateGovernanceUpgradePayload(): Buffer {
     throw new Error(
       "Can not create a governance message for upgrading contracts on all chains!"
     );
   }
 
+  async getAccountAddress(_privateKey: PrivateKey): Promise<string> {
+    throw new Error("Can not get account for GlobalChain.");
+  }
+
+  async getAccountBalance(_privateKey: PrivateKey): Promise<number> {
+    throw new Error("Can not get account balance for GlobalChain.");
+  }
+
   getType(): string {
     return GlobalChain.type;
   }
@@ -177,7 +208,9 @@ export class CosmWasmChain extends Chain {
     return new CosmosUpgradeContract(this.wormholeChainName, codeId).encode();
   }
 
-  async getExecutor(privateKey: PrivateKey) {
+  async getExecutor(
+    privateKey: PrivateKey
+  ): Promise<CosmwasmExecutor | InjectiveExecutor> {
     if (this.getId().indexOf("injective") > -1) {
       return InjectiveExecutor.fromPrivateKey(
         this.isMainnet() ? Network.Mainnet : Network.Testnet,
@@ -190,6 +223,20 @@ export class CosmWasmChain extends Chain {
       this.gasPrice + this.feeDenom
     );
   }
+
+  async getAccountAddress(privateKey: PrivateKey): Promise<string> {
+    const executor = await this.getExecutor(privateKey);
+    if (executor instanceof InjectiveExecutor) {
+      return executor.getAddress();
+    } else {
+      return await executor.getAddress();
+    }
+  }
+
+  async getAccountBalance(privateKey: PrivateKey): Promise<number> {
+    const executor = await this.getExecutor(privateKey);
+    return await executor.getBalance();
+  }
 }
 
 export class SuiChain extends Chain {
@@ -238,6 +285,27 @@ export class SuiChain extends Chain {
       digest
     ).encode();
   }
+
+  getProvider(): JsonRpcProvider {
+    return new JsonRpcProvider(new Connection({ fullnode: this.rpcUrl }));
+  }
+
+  async getAccountAddress(privateKey: PrivateKey): Promise<string> {
+    const provider = this.getProvider();
+    const keypair = Ed25519Keypair.fromSecretKey(
+      Buffer.from(privateKey, "hex")
+    );
+    const wallet = new RawSigner(keypair, provider);
+    return await wallet.getAddress();
+  }
+
+  async getAccountBalance(privateKey: PrivateKey): Promise<number> {
+    const provider = this.getProvider();
+    const balance = await provider.getBalance({
+      owner: await this.getAccountAddress(privateKey),
+    });
+    return Number(balance.totalBalance) / 10 ** 9;
+  }
 }
 
 export class EvmChain extends Chain {
@@ -361,6 +429,20 @@ export class EvmChain extends Chain {
     });
     return deployedContract.options.address;
   }
+
+  async getAccountAddress(privateKey: PrivateKey): Promise<string> {
+    const web3 = new Web3(this.getRpcUrl());
+    const signer = web3.eth.accounts.privateKeyToAccount(privateKey);
+    return signer.address;
+  }
+
+  async getAccountBalance(privateKey: PrivateKey): Promise<number> {
+    const web3 = new Web3(this.getRpcUrl());
+    const balance = await web3.eth.getBalance(
+      await this.getAccountAddress(privateKey)
+    );
+    return Number(balance) / 10 ** 18;
+  }
 }
 
 export class AptosChain extends Chain {
@@ -413,4 +495,20 @@ export class AptosChain extends Chain {
       parsed.rpcUrl
     );
   }
+
+  async getAccountAddress(privateKey: PrivateKey): Promise<string> {
+    const account = new AptosAccount(
+      new Uint8Array(Buffer.from(privateKey, "hex"))
+    );
+    return account.address().toString();
+  }
+
+  async getAccountBalance(privateKey: PrivateKey): Promise<number> {
+    const client = this.getClient();
+    const account = new AptosAccount(
+      new Uint8Array(Buffer.from(privateKey, "hex"))
+    );
+    const coinClient = new CoinClient(client);
+    return Number(await coinClient.checkBalance(account)) / 10 ** 8;
+  }
 }

contract_manager/src/contracts/sui.ts

@@ -1,7 +1,5 @@
 import {
-  Connection,
   Ed25519Keypair,
-  JsonRpcProvider,
   ObjectId,
   RawSigner,
   SUI_CLOCK_OBJECT_ID,
@@ -377,7 +375,7 @@ export class SuiContract extends Contract {
   }
 
   getProvider() {
-    return new JsonRpcProvider(new Connection({ fullnode: this.chain.rpcUrl }));
+    return this.chain.getProvider();
   }
 
   private async getStateFields() {

target_chains/cosmwasm/tools/src/chains-manager/cosmwasm.ts

@@ -18,6 +18,7 @@ import {
   UpdateContractAdminResponse,
 } from "./chain-executor";
 import {
+  CosmWasmClient,
   DeliverTxResponse,
   MsgExecuteContractEncodeObject,
   MsgInstantiateContractEncodeObject,
@@ -63,7 +64,20 @@ export class CosmwasmExecutor implements ChainExecutor {
     );
   }
 
-  private async getAddress(): Promise<string> {
+  async getBalance(): Promise<number> {
+    const address = (await this.signer.getAccounts())[0].address;
+    const cosmwasmClient = await CosmWasmClient.connect(this.endpoint);
+
+    // We are interested only in the coin that we pay gas fees in.
+    const denom = GasPrice.fromString(this.gasPrice).denom;
+    const balance = await cosmwasmClient.getBalance(address, denom);
+
+    // By default the coins have 6 decimal places in CosmWasm
+    // and the denom is usually `u<chain>`.
+    return Number(balance.amount) / 10 ** 6;
+  }
+
+  async getAddress(): Promise<string> {
     return (await this.signer.getAccounts())[0].address;
   }
 

target_chains/cosmwasm/tools/src/chains-manager/injective.ts

@@ -10,6 +10,8 @@ import {
   TxGrpcClient,
   TxResponse,
   createTransactionFromMsg,
+  GrpcAccountPortfolio,
+  ChainGrpcBankApi,
 } from "@injectivelabs/sdk-ts";
 import {
   ChainExecutor,
@@ -53,10 +55,23 @@ export class InjectiveExecutor implements ChainExecutor {
     return new InjectiveExecutor(network, wallet);
   }
 
-  private getAddress(): string {
+  getAddress(): string {
     return this.wallet.toBech32();
   }
 
+  async getBalance(): Promise<number> {
+    const endpoints = getNetworkEndpoints(this.network);
+
+    const chainGrpcAuthApi = new ChainGrpcBankApi(endpoints.grpc);
+
+    const balance = await chainGrpcAuthApi.fetchBalance({
+      accountAddress: this.getAddress(),
+      denom: "inj",
+    });
+
+    return Number(balance.amount) / 10 ** 18;
+  }
+
   private async signAndBroadcastMsg(msg: Msgs): Promise<TxResponse> {
     const networkInfo = getNetworkInfo(this.network);
     const endpoints = getNetworkEndpoints(this.network);