Merge pull request #502 from trezor/fix/import-tool

Fix/import tool

src/actions/DiscoveryActions.js

@@ -16,7 +16,6 @@ import type {
     Account,
 } from 'flowtype';
 import type { Discovery, State } from 'reducers/DiscoveryReducer';
-import * as LocalStorageActions from 'actions/LocalStorageActions';
 import * as BlockchainActions from './BlockchainActions';
 import * as EthereumDiscoveryActions from './ethereum/DiscoveryActions';
 import * as RippleDiscoveryActions from './ripple/DiscoveryActions';
@@ -121,7 +120,6 @@ const start = (device: TrezorDevice, network: string, ignoreCompleted?: boolean)
     }
 
     if (!discoveryProcess) {
-        dispatch(addImportedAccounts());
         dispatch(begin(device, network));
     } else if (discoveryProcess.completed && !ignoreCompleted) {
         dispatch({
@@ -382,17 +380,3 @@ export const addAccount = (): ThunkAction => (dispatch: Dispatch, getState: GetS
     if (!selected) return;
     dispatch(start(selected, getState().router.location.state.network, true));
 };
-
-export const addImportedAccounts = (): ThunkAction => (dispatch: Dispatch): void => {
-    // get imported accounts from local storage
-    const importedAccounts = LocalStorageActions.getImportedAccounts();
-    if (importedAccounts) {
-        // create each account
-        importedAccounts.forEach(account => {
-            dispatch({
-                type: ACCOUNT.CREATE,
-                payload: account,
-            });
-        });
-    }
-};

src/actions/ImportAccountActions.js

@@ -3,7 +3,7 @@
 import * as ACCOUNT from 'actions/constants/account';
 import * as IMPORT from 'actions/constants/importAccount';
 import * as NOTIFICATION from 'actions/constants/notification';
-import type { AsyncAction, TrezorDevice, Network, Dispatch, GetState } from 'flowtype';
+import type { AsyncAction, Account, TrezorDevice, Network, Dispatch, GetState } from 'flowtype';
 import * as BlockchainActions from 'actions/ethereum/BlockchainActions';
 import * as LocalStorageActions from 'actions/LocalStorageActions';
 import TrezorConnect from 'trezor-connect';
@@ -21,6 +21,15 @@ export type ImportAccountAction =
           error: ?string,
       };
 
+const findIndex = (accounts: Array<Account>, network: Network, device: TrezorDevice): number => {
+    return accounts.filter(
+        a =>
+            a.imported === true &&
+            a.network === network.shortcut &&
+            a.deviceID === (device.features || {}).device_id
+    ).length;
+};
+
 export const importAddress = (
     address: string,
     network: Network,
@@ -33,20 +42,14 @@ export const importAddress = (
     });
 
     let payload;
-    const index = getState().accounts.filter(
-        a =>
-            a.imported === true &&
-            a.network === network.shortcut &&
-            device &&
-            a.deviceState === device.state
-    ).length;
-
     try {
         if (network.type === 'ethereum') {
             const account = await dispatch(
                 BlockchainActions.discoverAccount(device, address, network.shortcut)
             );
 
+            const index = findIndex(getState().accounts, network, device);
+
             const empty = account.nonce <= 0 && account.balance === '0';
             payload = {
                 imported: true,
@@ -97,6 +100,7 @@ export const importAddress = (
 
             const account = response.payload;
             const empty = account.sequence <= 0 && account.balance === '0';
+            const index = findIndex(getState().accounts, network, device);
 
             payload = {
                 imported: true,

src/actions/LocalStorageActions.js

@@ -237,6 +237,16 @@ const loadStorageData = (): ThunkAction => (dispatch: Dispatch): void => {
         });
     }
 
+    const importedAccounts = getImportedAccounts();
+    if (importedAccounts) {
+        importedAccounts.forEach(account => {
+            dispatch({
+                type: ACCOUNT.CREATE,
+                payload: account,
+            });
+        });
+    }
+
     const userTokens: ?string = storageUtils.get(TYPE, KEY_TOKENS);
     if (userTokens) {
         dispatch({

src/reducers/AccountsReducer.js

@@ -42,23 +42,24 @@ export type State = Array<Account>;
 
 const initialState: State = [];
 
-export const findAccount = (
-    state: State,
-    index: number,
-    deviceState: string,
-    network: string
-): ?Account =>
-    state.find(a => a.deviceState === deviceState && a.index === index && a.network === network);
-
 export const findDeviceAccounts = (
     state: State,
     device: TrezorDevice,
     network: string
 ): Array<Account> => {
     if (network) {
-        return state.filter(addr => addr.deviceState === device.state && addr.network === network);
+        return state.filter(
+            addr =>
+                (addr.deviceState === device.state ||
+                    (addr.imported && addr.deviceID === (device.features || {}).device_id)) &&
+                addr.network === network
+        );
     }
-    return state.filter(addr => addr.deviceState === device.state);
+    return state.filter(
+        addr =>
+            addr.deviceState === device.state ||
+            (addr.imported && addr.deviceID === (device.features || {}).device_id)
+    );
 };
 
 const createAccount = (state: State, account: Account): State => {
@@ -83,8 +84,16 @@ const createAccount = (state: State, account: Account): State => {
     return newState;
 };
 
-const removeAccounts = (state: State, device: TrezorDevice): State =>
+const removeAccounts = (
-    state.filter(account => account.deviceState !== device.state);
+    state: State,
+    device: TrezorDevice,
+    keepImportedAccounts = false
+): State => {
+    if (keepImportedAccounts) {
+        return state.filter(account => account.deviceState !== device.state || account.imported);
+    }
+    return state.filter(account => account.deviceState !== device.state);
+};
 
 const clear = (state: State, devices: Array<TrezorDevice>): State => {
     let newState: State = [...state];
@@ -114,9 +123,11 @@ export default (state: State = initialState, action: Action): State => {
         case CONNECT.FORGET:
         case CONNECT.FORGET_SINGLE:
         case CONNECT.FORGET_SILENT:
-        case CONNECT.RECEIVE_WALLET_TYPE:
             return removeAccounts(state, action.device);
 
+        case CONNECT.RECEIVE_WALLET_TYPE:
+            return removeAccounts(state, action.device, true); // removes all accounts except imported ones
+
         case WALLET.CLEAR_UNAVAILABLE_DEVICE_DATA:
             return clear(state, action.devices);
 

src/reducers/utils/index.js

@@ -92,7 +92,8 @@ export const getSelectedAccount = (state: State): ?Account => {
     return state.accounts.find(
         a =>
             a.imported === isImported &&
-            a.deviceState === device.state &&
+            (a.deviceState === device.state ||
+                (a.imported && a.deviceID === (device.features || {}).device_id)) &&
             a.index === index &&
             a.network === locationState.network
     );