phantom/bitcoin
1
0
Fork 0
mirror of https://github.com/bitcoin/bitcoin.git synced 2025-04-29 23:09:44 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

test: Add Wallet Unlock Context Manager

Browse source 
Add Context Manager to manage wallet locking/unlocking with passphrase
This commit is contained in:
Brandon Odiwuor 2023-10-09 15:30:40 +03:00
parent db283a6b6f
commit 004903ebad
7 changed files with 140 additions and 125 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
test/functional/test_framework/wallet_util.py
View file

@ -122,3 +122,22 @@ def generate_keypair(compressed=True, wif=False):
if wif:
privkey = bytes_to_wif(privkey.get_bytes(), compressed)
return privkey, pubkey
class WalletUnlock():
"""
A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.
"""
MAXIMUM_TIMEOUT = 999000
def __init__(self, wallet, passphrase, timeout=MAXIMUM_TIMEOUT):
self.wallet = wallet
self.passphrase = passphrase
self.timeout = timeout
def __enter__(self):
self.wallet.walletpassphrase(self.passphrase, self.timeout)
def __exit__(self, *args):
_ = args
self.wallet.walletlock()

8
test/functional/wallet_createwallet.py
View file

@ -12,7 +12,7 @@ from test_framework.util import (
assert_equal,
assert_raises_rpc_error,
)
from test_framework.wallet_util import generate_keypair
from test_framework.wallet_util import generate_keypair, WalletUnlock
EMPTY_PASSPHRASE_MSG = "Empty string given as passphrase, wallet will not be encrypted."
@ -108,8 +108,8 @@ class CreateWalletTest(BitcoinTestFramework):
w4.encryptwallet('pass')
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", w4.getnewaddress)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", w4.getrawchangeaddress)
with WalletUnlock(w4, "pass"):
# Now set a seed and it should work. Wallet should also be encrypted
w4.walletpassphrase("pass", 999000)
if self.options.descriptors:
w4.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPcwuZGKp8TeWppSuLMiLe2d9PupB14QpPeQsqoj3LneJLhGHH13xESfvASyd4EFLJvLrG8b7DrLxEuV7hpF9uUc6XruKA1Wq/0h/*)'),
@ -142,7 +142,7 @@ class CreateWalletTest(BitcoinTestFramework):
self.nodes[0].createwallet(wallet_name='wblank', disable_private_keys=False, blank=True, passphrase='thisisapassphrase')
wblank = node.get_wallet_rpc('wblank')
assert_raises_rpc_error(-13, "Error: Please enter the wallet passphrase with walletpassphrase first.", wblank.signmessage, "needanargument", "test")
wblank.walletpassphrase("thisisapassphrase", 999000)
with WalletUnlock(wblank, "thisisapassphrase"):
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getnewaddress)
assert_raises_rpc_error(-4, "Error: This wallet has no available keys", wblank.getrawchangeaddress)
@ -151,7 +151,7 @@ class CreateWalletTest(BitcoinTestFramework):
self.nodes[0].createwallet(wallet_name='w6', disable_private_keys=False, blank=False, passphrase='thisisapassphrase')
w6 = node.get_wallet_rpc('w6')
assert_raises_rpc_error(-13, "Error: Please enter the wallet passphrase with walletpassphrase first.", w6.signmessage, "needanargument", "test")
w6.walletpassphrase("thisisapassphrase", 999000)
with WalletUnlock(w6, "thisisapassphrase"):
w6.signmessage(w6.getnewaddress('', 'legacy'), "test")
w6.keypoolrefill(1)
# There should only be 1 key for legacy, 3 for descriptors

7
test/functional/wallet_descriptor.py
View file

@ -16,6 +16,7 @@ from test_framework.util import (
assert_equal,
assert_raises_rpc_error
)
from test_framework.wallet_util import WalletUnlock
class WalletDescriptorTest(BitcoinTestFramework):
@ -129,11 +130,10 @@ class WalletDescriptorTest(BitcoinTestFramework):
# Encrypt wallet 0
send_wrpc.encryptwallet('pass')
send_wrpc.walletpassphrase("pass", 999000)
with WalletUnlock(send_wrpc, "pass"):
addr = send_wrpc.getnewaddress()
info2 = send_wrpc.getaddressinfo(addr)
assert info1['hdmasterfingerprint'] != info2['hdmasterfingerprint']
send_wrpc.walletlock()
assert 'hdmasterfingerprint' in send_wrpc.getaddressinfo(send_wrpc.getnewaddress())
info3 = send_wrpc.getaddressinfo(addr)
assert_equal(info2['desc'], info3['desc'])
@ -143,14 +143,13 @@ class WalletDescriptorTest(BitcoinTestFramework):
send_wrpc.getnewaddress()
self.log.info("Test that unlock is needed when deriving only hardened keys in an encrypted wallet")
send_wrpc.walletpassphrase("pass", 999000)
with WalletUnlock(send_wrpc, "pass"):
send_wrpc.importdescriptors([{
"desc": "wpkh(tprv8ZgxMBicQKsPd7Uf69XL1XwhmjHopUGep8GuEiJDZmbQz6o58LninorQAfcKZWARbtRtfnLcJ5MQ2AtHcQJCCRUcMRvmDUjyEmNUWwx8UbK/0h/*h)#y4dfsj7n",
"timestamp": "now",
"range": [0,10],
"active": True
}])
send_wrpc.walletlock()
# Exhaust keypool of 100
for _ in range(100):
send_wrpc.getnewaddress(address_type='bech32')

3
test/functional/wallet_dump.py
View file

@ -12,6 +12,7 @@ from test_framework.util import (
assert_equal,
assert_raises_rpc_error,
)
from test_framework.wallet_util import WalletUnlock
def read_dump(file_name, addrs, script_addrs, hd_master_addr_old):
@ -173,7 +174,7 @@ class WalletDumpTest(BitcoinTestFramework):
# encrypt wallet, restart, unlock and dump
self.nodes[0].encryptwallet('test')
self.nodes[0].walletpassphrase("test", 999000)
with WalletUnlock(self.nodes[0], "test"):
# Should be a no-op:
self.nodes[0].keypoolrefill()
self.nodes[0].dumpwallet(wallet_enc_dump)

10
test/functional/wallet_encryption.py
View file

@ -11,6 +11,7 @@ from test_framework.util import (
assert_raises_rpc_error,
assert_equal,
)
from test_framework.wallet_util import WalletUnlock
class WalletEncryptionTest(BitcoinTestFramework):
@ -59,19 +60,17 @@ class WalletEncryptionTest(BitcoinTestFramework):
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase + "wrong", 10)
# Test walletlock
self.nodes[0].walletpassphrase(passphrase, 999000)
with WalletUnlock(self.nodes[0], passphrase):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
assert_raises_rpc_error(-13, "Please enter the wallet passphrase with walletpassphrase first", self.nodes[0].signmessage, address, msg)
# Test passphrase changes
self.nodes[0].walletpassphrasechange(passphrase, passphrase2)
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase, 10)
self.nodes[0].walletpassphrase(passphrase2, 999000)
with WalletUnlock(self.nodes[0], passphrase2):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
# Test timeout bounds
assert_raises_rpc_error(-8, "Timeout cannot be negative.", self.nodes[0].walletpassphrase, passphrase2, -10)
@ -97,10 +96,9 @@ class WalletEncryptionTest(BitcoinTestFramework):
self.nodes[0].walletpassphrasechange(passphrase2, passphrase_with_nulls)
# walletpassphrasechange should not stop at null characters
assert_raises_rpc_error(-14, "wallet passphrase entered was incorrect", self.nodes[0].walletpassphrase, passphrase_with_nulls.partition("\0")[0], 10)
self.nodes[0].walletpassphrase(passphrase_with_nulls, 999000)
with WalletUnlock(self.nodes[0], passphrase_with_nulls):
sig = self.nodes[0].signmessage(address, msg)
assert self.nodes[0].verifymessage(address, sig, msg)
self.nodes[0].walletlock()
if __name__ == '__main__':

10
test/functional/wallet_fundrawtransaction.py
View file

@ -25,7 +25,7 @@ from test_framework.util import (
find_vout_for_address,
get_fee,
)
from test_framework.wallet_util import generate_keypair
from test_framework.wallet_util import generate_keypair, WalletUnlock
ERR_NOT_ENOUGH_PRESET_INPUTS = "The preselected coins total amount does not cover the transaction target. " \
"Please allow other inputs to be automatically selected or include more coins manually"
@ -581,7 +581,7 @@ class RawTransactionsTest(BitcoinTestFramework):
wallet.encryptwallet("test")
if self.options.descriptors:
wallet.walletpassphrase("test", 999000)
with WalletUnlock(wallet, "test"):
wallet.importdescriptors([{
'desc': descsum_create('wpkh(tprv8ZgxMBicQKsPdYeeZbPSKd2KYLmeVKtcFA7kqCxDvDR13MQ6us8HopUR2wLcS2ZKPhLyKsqpDL2FtL73LMHcgoCL7DXsciA8eX8nbjCR2eG/0h/*h)'),
'timestamp': 'now',
@ -593,7 +593,6 @@ class RawTransactionsTest(BitcoinTestFramework):
'active': True,
'internal': True
}])
wallet.walletlock()
# Drain the keypool.
wallet.getnewaddress()
@ -619,9 +618,8 @@ class RawTransactionsTest(BitcoinTestFramework):
assert_raises_rpc_error(-4, "Transaction needs a change address, but we can't generate it.", wallet.fundrawtransaction, rawtx)
# Refill the keypool.
wallet.walletpassphrase("test", 999000)
with WalletUnlock(wallet, "test"):
wallet.keypoolrefill(8) #need to refill the keypool to get an internal change address
wallet.walletlock()
assert_raises_rpc_error(-13, "walletpassphrase", wallet.sendtoaddress, self.nodes[0].getnewaddress(), 1.2)
@ -634,7 +632,7 @@ class RawTransactionsTest(BitcoinTestFramework):
assert fundedTx["changepos"] != -1
# Now we need to unlock.
wallet.walletpassphrase("test", 999000)
with WalletUnlock(wallet, "test"):
signedTx = wallet.signrawtransactionwithwallet(fundedTx['hex'])
wallet.sendrawtransaction(signedTx['hex'])
self.generate(self.nodes[1], 1)

8
test/functional/wallet_keypool.py
View file

@ -9,6 +9,7 @@ from decimal import Decimal
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal, assert_raises_rpc_error
from test_framework.wallet_util import WalletUnlock
class KeyPoolTest(BitcoinTestFramework):
def add_options(self, parser):
@ -85,9 +86,8 @@ class KeyPoolTest(BitcoinTestFramework):
assert_raises_rpc_error(-12, "Error: Keypool ran out, please call keypoolrefill first", nodes[0].getnewaddress)
# put six (plus 2) new keys in the keypool (100% external-, +100% internal-keys, 1 in min)
nodes[0].walletpassphrase("test", 999000)
with WalletUnlock(nodes[0], 'test'):
nodes[0].keypoolrefill(6)
nodes[0].walletlock()
wi = nodes[0].getwalletinfo()
if self.options.descriptors:
assert_equal(wi['keypoolsize_hd_internal'], 24)
@ -131,7 +131,7 @@ class KeyPoolTest(BitcoinTestFramework):
nodes[0].getnewaddress()
assert_raises_rpc_error(-12, "Keypool ran out", nodes[0].getnewaddress)
nodes[0].walletpassphrase("test", 999000)
with WalletUnlock(nodes[0], 'test'):
nodes[0].keypoolrefill(100)
wi = nodes[0].getwalletinfo()
if self.options.descriptors:
@ -170,8 +170,8 @@ class KeyPoolTest(BitcoinTestFramework):
else:
res = w2.importmulti([{'desc': desc, 'timestamp': 'now'}])
assert_equal(res[0]['success'], True)
w1.walletpassphrase("test", 999000)
with WalletUnlock(w1, 'test'):
res = w1.sendtoaddress(address=address, amount=0.00010000)
self.generate(nodes[0], 1)
destination = addr.pop()