From 88da28d6d057f7baacf169047573e75a17b25caa Mon Sep 17 00:00:00 2001 From: Ben Hauser Date: Tue, 30 Jan 2024 03:12:02 +0400 Subject: [PATCH] feat: allow_retry when loading account --- brownie/network/account.py | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/brownie/network/account.py b/brownie/network/account.py index f8dfb6c7e..c1f70f708 100644 --- a/brownie/network/account.py +++ b/brownie/network/account.py @@ -182,7 +182,9 @@ def from_mnemonic( return new_accounts[0] return new_accounts - def load(self, filename: str = None, password: str = None) -> Union[List, "LocalAccount"]: + def load( + self, filename: str = None, password: str = None, allow_retry: bool = False + ) -> Union[List, "LocalAccount"]: """ Load a local account from a keystore file. @@ -193,6 +195,8 @@ def load(self, filename: str = None, password: str = None) -> Union[List, "Local password: str Password to unlock the keystore. If `None`, password is entered via a getpass prompt. + allow_retry: bool + If True, allows re-attempt when the given password is incorrect. Returns ------- @@ -218,10 +222,22 @@ def load(self, filename: str = None, password: str = None) -> Union[List, "Local raise FileNotFoundError(f"Cannot find {json_file}") with json_file.open() as fp: - priv_key = web3.eth.account.decrypt( - json.load(fp), - password or getpass(f'Enter password for "{json_file.stem}": '), - ) + encrypted = json.load(fp) + + prompt = f'Enter password for "{json_file.stem}": ' + while True: + if password is None: + password = getpass(prompt) + try: + priv_key = web3.eth.account.decrypt(encrypted, password) + break + except ValueError as e: + if allow_retry: + prompt = f"Incorrect password, try again: " + password = None + continue + raise e + return self.add(priv_key) def at(self, address: str, force: bool = False) -> "LocalAccount":