Skip to content

Commit

Permalink
feat: allow_retry when loading account
Browse files Browse the repository at this point in the history
  • Loading branch information
iamdefinitelyahuman committed Jan 29, 2024
1 parent a60b232 commit 88da28d
Showing 1 changed file with 21 additions and 5 deletions.
26 changes: 21 additions & 5 deletions brownie/network/account.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
-------
Expand All @@ -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":
Expand Down

0 comments on commit 88da28d

Please sign in to comment.