From a721d30425744461e09bdab501e5aa7f124ed5cf Mon Sep 17 00:00:00 2001
From: Noah Saso <noahsaso@gmail.com>
Date: Fri, 10 May 2024 10:51:28 -0400
Subject: [PATCH] get all keys from wallet

---
 packages/graz/src/actions/account.ts          | 33 +++++------------
 packages/graz/src/actions/wallet/capsule.ts   | 37 ++++++++++---------
 .../actions/wallet/wallet-connect/index.ts    | 14 ++-----
 3 files changed, 33 insertions(+), 51 deletions(-)

diff --git a/packages/graz/src/actions/account.ts b/packages/graz/src/actions/account.ts
index 9853b37..ce6f76b 100644
--- a/packages/graz/src/actions/account.ts
+++ b/packages/graz/src/actions/account.ts
@@ -1,4 +1,3 @@
-import { fromBech32, toBech32 } from "@cosmjs/encoding";
 import type { OfflineDirectSigner } from "@cosmjs/proto-signing";
 import type { ChainInfo, Key, OfflineAminoSigner } from "@keplr-wallet/types";
 
@@ -67,17 +66,11 @@ export const connect = async (args?: ConnectArgs): Promise<ConnectResult> => {
       const _resAcc = useGrazSessionStore.getState().accounts;
       useGrazSessionStore.setState({ status: "connecting" });
 
-      const key = await wallet.getKey(chainIds[0]!);
-      const resultAcccounts: Record<string, Key> = {};
-      chainIds.forEach((chainId) => {
-        resultAcccounts[chainId] = {
-          ...key,
-          bech32Address: toBech32(
-            chains!.find((x) => x.chainId === chainId)!.bech32Config.bech32PrefixAccAddr,
-            fromBech32(key.bech32Address).data,
-          ),
-        };
-      });
+      const resultAcccounts = Object.fromEntries(
+        await Promise.all(
+          chainIds.map(async (chainId): Promise<[string, Key]> => [chainId, await wallet.getKey(chainId)]),
+        ),
+      );
       useGrazSessionStore.setState((prev) => ({
         accounts: { ...(prev.accounts || {}), ...resultAcccounts },
       }));
@@ -105,17 +98,11 @@ export const connect = async (args?: ConnectArgs): Promise<ConnectResult> => {
       return { accounts: _resAcc!, walletType: currentWalletType, chains: connectedChains };
     }
     if (!isWalletConnect(currentWalletType)) {
-      const key = await wallet.getKey(chainIds[0]!);
-      const resultAcccounts: Record<string, Key> = {};
-      chainIds.forEach((chainId) => {
-        resultAcccounts[chainId] = {
-          ...key,
-          bech32Address: toBech32(
-            chains!.find((x) => x.chainId === chainId)!.bech32Config.bech32PrefixAccAddr,
-            fromBech32(key.bech32Address).data,
-          ),
-        };
-      });
+      const resultAcccounts = Object.fromEntries(
+        await Promise.all(
+          chainIds.map(async (chainId): Promise<[string, Key]> => [chainId, await wallet.getKey(chainId)]),
+        ),
+      );
       useGrazSessionStore.setState((prev) => ({
         accounts: { ...(prev.accounts || {}), ...resultAcccounts },
       }));
diff --git a/packages/graz/src/actions/wallet/capsule.ts b/packages/graz/src/actions/wallet/capsule.ts
index d0f015a..010f52d 100644
--- a/packages/graz/src/actions/wallet/capsule.ts
+++ b/packages/graz/src/actions/wallet/capsule.ts
@@ -1,5 +1,5 @@
 import type { AminoSignResponse } from "@cosmjs/amino";
-import { fromBech32, toBech32 } from "@cosmjs/encoding";
+import { fromBech32 } from "@cosmjs/encoding";
 import type { DirectSignResponse } from "@cosmjs/proto-signing";
 import type { Keplr, Key } from "@keplr-wallet/types";
 
@@ -40,22 +40,25 @@ export const getCapsule = (): Wallet => {
     await client.enable();
     const chainIds = useGrazInternalStore.getState().capsuleState?.chainId;
     if (!chainIds) throw new Error("Chain ids are not set");
-    const key = await client.getAccount(chainIds[0]!);
-    const resultAcccounts: Record<string, Key> = {};
-    chainIds.forEach((chainId) => {
-      resultAcccounts[chainId] = {
-        address: fromBech32(key.address).data,
-        bech32Address: toBech32(
-          chains.find((x) => x.chainId === chainId)!.bech32Config.bech32PrefixAccAddr,
-          fromBech32(key.address).data,
-        ),
-        algo: key.algo,
-        name: key.username || "",
-        pubKey: key.pubkey,
-        isKeystone: false,
-        isNanoLedger: false,
-      };
-    });
+    const resultAcccounts = Object.fromEntries(
+      await Promise.all(
+        chainIds.map(async (chainId): Promise<[string, Key]> => {
+          const account = await client.getAccount(chainId);
+          return [
+            chainId,
+            {
+              address: fromBech32(account.address).data,
+              bech32Address: account.address,
+              algo: account.algo,
+              name: account.username || "",
+              pubKey: account.pubkey,
+              isKeystone: false,
+              isNanoLedger: false,
+            },
+          ];
+        }),
+      ),
+    );
     useGrazSessionStore.setState((prev) => ({
       accounts: { ...(prev.accounts || {}), ...resultAcccounts },
     }));
diff --git a/packages/graz/src/actions/wallet/wallet-connect/index.ts b/packages/graz/src/actions/wallet/wallet-connect/index.ts
index 86e87f3..3a74e06 100644
--- a/packages/graz/src/actions/wallet/wallet-connect/index.ts
+++ b/packages/graz/src/actions/wallet/wallet-connect/index.ts
@@ -297,17 +297,9 @@ export const getWalletConnect = (params?: GetWalletConnectParams): Wallet => {
     try {
       await promiseWithTimeout(
         (async () => {
-          const wcAccount = await getKey(chainId[0]!);
-          const resultAcccounts: Record<string, Key> = {};
-          chainId.forEach((x) => {
-            resultAcccounts[x] = {
-              ...wcAccount,
-              bech32Address: toBech32(
-                chains!.find((y) => y.chainId === x)!.bech32Config.bech32PrefixAccAddr,
-                fromBech32(wcAccount.bech32Address).data,
-              ),
-            };
-          });
+          const resultAcccounts = Object.fromEntries(
+            await Promise.all(chainId.map(async (c): Promise<[string, Key]> => [c, await getKey(c)])),
+          );
           useGrazSessionStore.setState({
             accounts: resultAcccounts,
           });