fix: Fix exporting private key for a select account (#431)

packages/adena-extension/src/common/constants/storage.constant.ts

@@ -1,2 +1,3 @@
 export const WALLET_EXPORT_TYPE_STORAGE_KEY = 'WALLET_EXPORT_TYPE' as const;
+export const WALLET_EXPORT_ACCOUNT_ID = 'WALLET_EXPORT_ACCOUNT_ID' as const;
 export const QUESTIONNAIRE_EXPIRATION_MIN = 30 * 24 * 60;

packages/adena-extension/src/hooks/web/wallet-export/use-wallet-export-screen.ts

@@ -2,15 +2,19 @@ import { useCallback, useEffect, useState } from 'react';
 import { Account } from 'adena-module';
 
 import useAppNavigate from '@hooks/use-app-navigate';
-import { useAdenaContext } from '@hooks/use-context';
+import { useAdenaContext, useWalletContext } from '@hooks/use-context';
 import { useCurrentAccount } from '@hooks/use-current-account';
 import { RoutePath } from '@types';
 import { AdenaStorage } from '@common/storage';
-import { WALLET_EXPORT_TYPE_STORAGE_KEY } from '@common/constants/storage.constant';
+import {
+  WALLET_EXPORT_ACCOUNT_ID,
+  WALLET_EXPORT_TYPE_STORAGE_KEY,
+} from '@common/constants/storage.constant';
 import useQuestionnaire from '../use-questionnaire';
 import useIndicatorStep, {
   UseIndicatorStepReturn,
 } from '@hooks/wallet/broadcast-transaction/use-indicator-step';
+import { useQuery } from '@tanstack/react-query';
 
 export type UseWalletExportReturn = {
   currentAccount: Account | null;
@@ -56,33 +60,51 @@ export const walletExportStepNo: Record<WalletExportStateType, number> = {
 };
 
 const useWalletExportScreen = (): UseWalletExportReturn => {
-  const { currentAccount } = useCurrentAccount();
+  const { wallet } = useWalletContext();
   const { walletService } = useAdenaContext();
+  const { currentAccount } = useCurrentAccount();
   const { ableToSkipQuestionnaire } = useQuestionnaire();
   const [exportType, setExportType] = useState<ExportType>('NONE');
   const { navigate, params } = useAppNavigate<RoutePath.WebWalletExport>();
   const [walletExportState, setWalletExportState] = useState<WalletExportStateType>(
     params?.doneQuestionnaire ? 'CHECK_PASSWORD' : 'INIT',
   );
   const [exportData, setExportData] = useState<string | null>(null);
+  const [exportAccountId, setExportAccountId] = useState<string | null>(null);
   const indicatorInfo = useIndicatorStep({
     stepMap: walletExportStepNo,
     currentState: walletExportState,
     hasQuestionnaire: true,
   });
 
+  const { data: account = null } = useQuery(
+    ['walletExportScreen/account', exportData, wallet, exportAccountId, currentAccount],
+    async () => {
+      if (exportData === 'SEED_PHRASE') {
+        return currentAccount;
+      }
+      return wallet?.accounts.find((account) => account.id === exportAccountId) || currentAccount;
+    },
+    {},
+  );
+
   const _initExportType = useCallback(async () => {
     const sessionStorage = AdenaStorage.session();
     const exportType = await sessionStorage.get(WALLET_EXPORT_TYPE_STORAGE_KEY);
+    const exportAccountId = await sessionStorage.get(WALLET_EXPORT_ACCOUNT_ID);
+    console.log('???', exportType);
+    await sessionStorage.set(WALLET_EXPORT_TYPE_STORAGE_KEY, '');
+    await sessionStorage.set(WALLET_EXPORT_ACCOUNT_ID, '');
     switch (exportType) {
       case 'PRIVATE_KEY':
+        setExportAccountId(exportAccountId);
         setExportType('PRIVATE_KEY');
         break;
       case 'SEED_PHRASE':
         setExportType('SEED_PHRASE');
         break;
       default:
-        navigate(RoutePath.Home);
+        location.replace('/register.html');
         break;
     }
   }, []);
@@ -121,12 +143,12 @@ const useWalletExportScreen = (): UseWalletExportReturn => {
 
   const moveExport = useCallback(
     async (password: string) => {
-      if (exportType === 'NONE' || !currentAccount) {
+      if (exportType === 'NONE' || !account) {
         return;
       }
       const wallet = await walletService.loadWalletWithPassword(password);
       const instance = wallet.clone();
-      instance.currentAccountId = currentAccount.id;
+      instance.currentAccountId = account.id;
       if (exportType === 'PRIVATE_KEY') {
         const privateKey = await instance.getPrivateKeyStr();
         setExportData(privateKey);
@@ -136,7 +158,7 @@ const useWalletExportScreen = (): UseWalletExportReturn => {
       }
       setWalletExportState('RESULT');
     },
-    [exportType, currentAccount, walletService],
+    [exportType, account, walletService],
   );
 
   useEffect(() => {
@@ -145,7 +167,7 @@ const useWalletExportScreen = (): UseWalletExportReturn => {
 
   return {
     indicatorInfo,
-    currentAccount,
+    currentAccount: account,
     exportType,
     walletExportState,
     exportData,

packages/adena-extension/src/pages/popup/certify/security-privacy/index.tsx

@@ -11,7 +11,10 @@ import mixins from '@styles/mixins';
 import useAppNavigate from '@hooks/use-app-navigate';
 import useLink from '@hooks/use-link';
 import { AdenaStorage } from '@common/storage';
-import { WALLET_EXPORT_TYPE_STORAGE_KEY } from '@common/constants/storage.constant';
+import {
+  WALLET_EXPORT_ACCOUNT_ID,
+  WALLET_EXPORT_TYPE_STORAGE_KEY,
+} from '@common/constants/storage.constant';
 
 type MenuType =
   | RoutePath.SettingChangePassword
@@ -25,41 +28,41 @@ const getMenuMakerInfo = (
   availRemove: boolean,
 ): {
   title: string;
-  navigatePath: MenuType
+  navigatePath: MenuType;
   mode: string;
   disabled: boolean;
 }[] => [
-    {
-      title: 'Change Password',
-      navigatePath: RoutePath.SettingChangePassword,
-      mode: 'DEFAULT',
-      disabled: false,
-    },
-    {
-      title: 'Reveal Seed Phrase',
-      navigatePath: 'EXPORT_SEED_PHRASE',
-      mode: 'DEFAULT',
-      disabled: !account || !isSeedAccount(account),
-    },
-    {
-      title: 'Export Private Key',
-      navigatePath: 'EXPORT_PRIVATE_KEY',
-      mode: 'DEFAULT',
-      disabled: !account || !hasPrivateKeyAccount(account),
-    },
-    {
-      title: 'Remove Account',
-      navigatePath: RoutePath.RemoveAccount,
-      mode: 'DANGER',
-      disabled: !availRemove,
-    },
-    {
-      title: 'Reset Wallet',
-      navigatePath: RoutePath.ResetWallet,
-      mode: 'DANGER',
-      disabled: false,
-    },
-  ];
+  {
+    title: 'Change Password',
+    navigatePath: RoutePath.SettingChangePassword,
+    mode: 'DEFAULT',
+    disabled: false,
+  },
+  {
+    title: 'Reveal Seed Phrase',
+    navigatePath: 'EXPORT_SEED_PHRASE',
+    mode: 'DEFAULT',
+    disabled: !account || !isSeedAccount(account),
+  },
+  {
+    title: 'Export Private Key',
+    navigatePath: 'EXPORT_PRIVATE_KEY',
+    mode: 'DEFAULT',
+    disabled: !account || !hasPrivateKeyAccount(account),
+  },
+  {
+    title: 'Remove Account',
+    navigatePath: RoutePath.RemoveAccount,
+    mode: 'DANGER',
+    disabled: !availRemove,
+  },
+  {
+    title: 'Reset Wallet',
+    navigatePath: RoutePath.ResetWallet,
+    mode: 'DANGER',
+    disabled: false,
+  },
+];
 
 export const SecurityPrivacy = (): JSX.Element => {
   const { navigate, goBack } = useAppNavigate();
@@ -68,22 +71,20 @@ export const SecurityPrivacy = (): JSX.Element => {
   const { availRemoveAccount } = useRemoveAccount();
   const [availRemove, setAvailRemove] = useState(true);
 
-  const onClickMenuItem = useCallback((navigatePath: MenuType) => {
-    if (
-      navigatePath === 'EXPORT_SEED_PHRASE' ||
-      navigatePath === 'EXPORT_PRIVATE_KEY'
-    ) {
-      const exportType = navigatePath ===
-        'EXPORT_SEED_PHRASE' ?
-        'SEED_PHRASE' :
-        'PRIVATE_KEY';
-      AdenaStorage.session().set(WALLET_EXPORT_TYPE_STORAGE_KEY, exportType).then(() => {
-        openSecurity();
-      });
-      return;
-    }
-    navigate(navigatePath);
-  }, [navigate])
+  const onClickMenuItem = useCallback(
+    (navigatePath: MenuType) => {
+      if (navigatePath === 'EXPORT_SEED_PHRASE' || navigatePath === 'EXPORT_PRIVATE_KEY') {
+        const exportType = navigatePath === 'EXPORT_SEED_PHRASE' ? 'SEED_PHRASE' : 'PRIVATE_KEY';
+        Promise.all([
+          AdenaStorage.session().set(WALLET_EXPORT_TYPE_STORAGE_KEY, exportType),
+          AdenaStorage.session().set(WALLET_EXPORT_ACCOUNT_ID, currentAccount?.id || ''),
+        ]).then(() => openSecurity());
+        return;
+      }
+      navigate(navigatePath);
+    },
+    [navigate],
+  );
 
   useEffect(() => {
     availRemoveAccount().then(setAvailRemove);

packages/adena-extension/src/pages/popup/wallet/account-details/index.tsx

@@ -9,7 +9,10 @@ import { useAccountName } from '@hooks/use-account-name';
 import { CommonFullContentLayout } from '@components/atoms';
 import useLink from '@hooks/use-link';
 import { AdenaStorage } from '@common/storage';
-import { WALLET_EXPORT_TYPE_STORAGE_KEY } from '@common/constants/storage.constant';
+import {
+  WALLET_EXPORT_ACCOUNT_ID,
+  WALLET_EXPORT_TYPE_STORAGE_KEY,
+} from '@common/constants/storage.constant';
 
 const ACCOUNT_NAME_LENGTH_LIMIT = 23;
 
@@ -55,6 +58,7 @@ const AccountDetailsContainer: React.FC = () => {
   const moveExportPrivateKey = useCallback(async () => {
     const sessionStorage = AdenaStorage.session();
     await sessionStorage.set(WALLET_EXPORT_TYPE_STORAGE_KEY, 'PRIVATE_KEY');
+    await sessionStorage.set(WALLET_EXPORT_ACCOUNT_ID, accountId || '');
     openSecurity();
   }, [account]);