From 0e2c7ce0d673ef83cbf28d061101e8abf26ab866 Mon Sep 17 00:00:00 2001 From: "Peter D. Gray" Date: Thu, 5 Mar 2026 11:51:21 -0500 Subject: [PATCH] WIF tests, support paper wallet format --- shared/wif.py | 83 ++++++---- testing/conftest.py | 4 +- testing/test_ux.py | 335 +--------------------------------------- testing/test_wif.py | 368 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 426 insertions(+), 364 deletions(-) create mode 100644 testing/test_wif.py diff --git a/shared/wif.py b/shared/wif.py index 39eeba50..ca131cec 100644 --- a/shared/wif.py +++ b/shared/wif.py @@ -6,7 +6,7 @@ from ubinascii import unhexlify as a2b_hex from ux import ux_show_story, ux_confirm, the_ux, import_export_prompt, ux_input_text, show_qr_code from menu import MenuSystem, MenuItem from utils import problem_file_line, show_single_address, node_from_pubkey -from files import CardSlot +from files import CardSlot, CardMissingError, needs_microsd from glob import settings from charcodes import KEY_QR, KEY_NFC, KEY_CANCEL from public_constants import AF_P2WPKH @@ -14,16 +14,22 @@ from msgsign import msg_signing_done def decode_wif(wif): + # Decode base58 encoded WIF string, return keypair and metadata raw = ngu.codecs.b58_decode(wif) assert raw[0] in (0xef, 0x80) + testnet = True if raw[0] == 0xef else False + assert len(raw) in (33, 34) + compressed = False if len(raw) == 34: # compressed pubkey assert raw[33] == 0x01 compressed = True + sk = raw[1:33] kp = ngu.secp256k1.keypair(sk) # catches wrong private keys + return kp, testnet, compressed @@ -241,46 +247,47 @@ class WIFStore(MenuSystem): else: # pick a likely-looking file: just looking at size and extension - fn = await file_picker(suffix='.txt', min_size=51, max_size=2000, + # - kinda big so we can import paper wallet directly + fn = await file_picker(suffix=['.csv', '.txt'], min_size=51, max_size=11000, none_msg="Must contain WIF(s)", **ch) if not fn: return - with CardSlot(readonly=True, **ch) as card: - with open(fn, 'rt') as fd: - got = fd.read().strip() + try: + with CardSlot(readonly=True, **ch) as card: + with open(fn, 'rt') as fd: + got = fd.read() + except CardMissingError: + await needs_microsd() + return + except Exception as e: + await ux_show_story('Failed to read file!\n\n%s' % e) + return if not got: return dis.fullscreen("Wait...") - # allow both newlines and commas as separators - if "\n" in got: - wifs = got.split("\n") - elif "," in got: - wifs = got.split(",") - else: - # just one wif - wifs = [got] + # allow commas, spaces, and newlines as separators + got = got.replace(',', ' ').split() saved = settings.get("wifs", []) - len_saved = len(saved) - - if (len_saved + len(wifs)) > self.MAX_ITEMS: - await ux_show_story("Max %d items allowed in WIF Store.\n\nAttempted to import %d keys," - " while remaining WIF store capacity is only %d. Please, make room" - " first." % (self.MAX_ITEMS, len(wifs), self.MAX_ITEMS - len_saved), - title="Failure") - return try: - for wif in wifs: + new_wifs = [] + dups = 0 + + for here in got: + here = here.strip() + if not here: + continue + try: - wif = wif.strip() - kp, testnet, compressed = decode_wif(wif) + kp, testnet, compressed = decode_wif(here) except Exception: - raise ValueError("wif decode") + # ignore garbage text, headers, addresses, etc. + continue assert compressed, "compressed only" assert testnet == (chains.current_chain().ctype != "BTC"), "chain" @@ -290,13 +297,25 @@ class WIFStore(MenuSystem): item = (pk, sk) - if item not in saved: # ignore duplicates - saved.append(item) + if item not in saved: # ignore dups + new_wifs.append(item) + else: + dups += 1 - if len_saved < len(saved): - settings.set('wifs', saved) - settings.save() - self.update_contents() + assert new_wifs, 'no valid WIF found' if not dups else 'duplicate WIF(s)' + + len_saved = len(saved) + if (len_saved + len(new_wifs)) > self.MAX_ITEMS: + await ux_show_story("Max %d items allowed in WIF Store.\n\nAttempted to import %d keys," + " while remaining WIF store capacity is only %d. Please, make room" + " first." % (self.MAX_ITEMS, len(new_wifs), self.MAX_ITEMS - len_saved), + title="Failure") + return + + saved.extend(new_wifs) + settings.set('wifs', saved) + settings.save() + self.update_contents() except Exception as e: await ux_show_story('Failed to import WIF.\n\n%s\n%s' % (e, problem_file_line(e)), @@ -312,4 +331,4 @@ def init_wif_store(): res[a2b_hex(pk)] = a2b_hex(sk) return res -# EOF \ No newline at end of file +# EOF diff --git a/testing/conftest.py b/testing/conftest.py index f01a2b08..fc852bbe 100644 --- a/testing/conftest.py +++ b/testing/conftest.py @@ -2962,6 +2962,8 @@ def set_deltamode(sim_exec): def import_wif_to_store(goto_home, pick_menu_item, cap_story, press_select, cap_menu, microsd_path, virtdisk_path, is_q1, scan_a_qr, need_keypress, garbage_collector, press_nfc, nfc_write_text, enter_complex): + + # Import a list of WIF keys into the "WIF Store" def doit(wif_lst, way="sd", sep="\n", early_exit=False): home = True try: @@ -2998,7 +3000,7 @@ def import_wif_to_store(goto_home, pick_menu_item, cap_story, press_select, cap_ time.sleep(0.3) elif way == "qr": if not is_q1: - raise pytest.xfail("Mk4 no QR") + raise pytest.xfail("needs scanner") assert f"{KEY_QR} to scan QR code" in story need_keypress(KEY_QR) diff --git a/testing/test_ux.py b/testing/test_ux.py index 091176fa..25a41b1c 100644 --- a/testing/test_ux.py +++ b/testing/test_ux.py @@ -5,13 +5,11 @@ from helpers import xfp2str, prandom, addr_from_display_format from charcodes import KEY_DOWN, KEY_QR, KEY_NFC, KEY_DELETE, KEY_UP from constants import AF_CLASSIC, simulator_fixed_words, simulator_fixed_xfp from mnemonic import Mnemonic -from bip32 import BIP32Node, PrivateKey -from base58 import encode_base58_checksum +from bip32 import BIP32Node mnem = Mnemonic('english') wordlist = mnem.wordlist - @pytest.fixture def enable_hw_ux(pick_menu_item, cap_story, press_select, goto_home): def doit(way, disable=False): @@ -462,7 +460,7 @@ def test_import_prv(way, testnet, pick_menu_item, cap_story, need_keypress, unit skip_if_useless_way(way) - node = BIP32Node.from_master_secret(os.urandom(32), netcode=netcode) + node = BIP32Node.from_master_secret(prandom(32), netcode=netcode) prv = node.hwif(as_private=True)+'\n' if testnet: assert "tprv" in prv @@ -825,7 +823,7 @@ def test_sign_file_from_list_files(f_len, goto_home, cap_story, pick_menu_item, fname = "test_sign_listed.pdf" signame = "test_sign_listed.sig" fpath = microsd_path(fname) - contents = os.urandom(f_len) + contents = prandom(f_len) digest = hashlib.sha256(contents).digest().hex() with open(fpath, "wb") as f: f.write(contents) @@ -874,7 +872,7 @@ def test_rename_from_list_files(goto_home, cap_story, pick_menu_item, need_keypr fname = "file_to_rename.pdf" fpath = microsd_path(fname) - contents = os.urandom(64) + contents = prandom(64) digest = hashlib.sha256(contents).digest().hex() with open(fpath, "wb") as f: f.write(contents) @@ -1231,331 +1229,6 @@ def test_file_picker_suffixes(pick_menu_item, goto_home, cap_story, microsd_wipe microsd_wipe() -@pytest.mark.parametrize("num_wifs", [1, 11]) -@pytest.mark.parametrize("separator", ["\n", ',']) -@pytest.mark.parametrize("way", ["sd", "nfc", "qr", "vdisk"]) -def test_wif_store_import(num_wifs, separator, way, import_wif_to_store, skip_if_useless_way, - settings_remove, goto_home): - skip_if_useless_way(way) - settings_remove("wifs") - - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - for _ in range(num_wifs) - ] - - import_wif_to_store(wif_list, way=way, sep=separator) - goto_home() - - -def test_wif_store_import_manual(import_wif_to_store, settings_remove, goto_home): - settings_remove("wifs") - - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - ] - - import_wif_to_store(wif_list, way="input") - goto_home() - - -@pytest.mark.parametrize("wif,err,way", [ - ("cWALDjUu1tszsCBMjBjL4mhYj2wHUWYDR8Q8aSjLKzjkWaXMLRaY", "wif decode", "sd"), # curve order - ("cMahea7zqjxrtgAbB7LSGbcQUr1uX1ojuat9jZodMN87J7g8rY9t", "wif decode", "nfc"), # zero - ("Ky2BtsR8qRN91PjktxaTQWMgJZUWSBJLjwip642vvoNyH1PeEpUP", "chain", "qr"), # mainnet key on testnet - ("91zb4oYGEvwEroihAbkdeoBpLSKnZYMdD1CPhfQD76fxrfNSp5J", "compressed only", "sd"), # uncompressed pk - ("cPPBMnQzGV4QAqD2HNPamprjvnmv6dQ2oysHCUVSRv2yXkVvWVtX", "wif decode", "nfc"), # wrong csum - ("cPPBMnQzGV4QAqD2HNPamprjvnmv6dQ2oysHCUVSRv2yXkVvWVtX;cN7M6sNzn4LGBxAozsmphxjuxVNaHcLre7Nm163qM3DpY3BZog1v", "wif decode", "sd"), # wrong separator -]) -def test_wif_store_import_fail(way, wif, err, import_wif_to_store, skip_if_useless_way, - settings_remove, press_select, cap_story, use_testnet, settings_get): - skip_if_useless_way(way) - use_testnet() - settings_remove("wifs") - - import_wif_to_store([wif], way=way, early_exit=True) - time.sleep(.1) - title, story = cap_story() - assert "Failed to import WIF" in story - assert err in story - press_select() - assert not settings_get("wifs") - - -@pytest.mark.parametrize("netcode", ["XTN", "BTC"]) -def test_wif_store_detail(netcode, import_wif_to_store, use_mainnet, cap_menu, pick_menu_item, - cap_story, need_keypress, settings_remove, cap_screen_qr, is_q1, - press_cancel, nfc_is_enabled, press_nfc, nfc_read_text, goto_home): - goto_home() - if netcode == "BTC": - use_mainnet() - - settings_remove("wifs") - - prefix = bytes([128]) if netcode == "BTC" else bytes([239]) - privkeys = [PrivateKey.parse(os.urandom(32)) for _ in range(5)] - wif_list = [ - encode_base58_checksum(prefix + bytes(sk) + b'\x01') - for sk in privkeys - ] - - import_wif_to_store(wif_list) - - time.sleep(.1) - menu = cap_menu() - target_mi = [] - for mi in menu: - if "⋯" in mi: - target_mi.append(mi) - - assert len(target_mi) == len(wif_list) - for mi, wif, sk in zip(target_mi, wif_list, privkeys): - mi_split = mi.split(" ")[-1].split("⋯") - assert len(mi_split) == 2 - assert mi_split[0] in wif - assert mi_split[1] in wif - pick_menu_item(mi) - - time.sleep(.1) - menu = cap_menu() - assert menu[0] == "Detail" - assert menu[1] == "Addresses" - assert menu[2] == "Sign MSG" - assert menu[3] == "Delete" - - pick_menu_item("Detail") - - title, story = cap_story() - assert title == "WIF" - - split_story = story.split("\n\n") - story_wif = split_story[0] - story_sk = split_story[1].split("\n")[-1] - story_pk = split_story[2].split("\n")[-1] - - assert f'{KEY_QR if is_q1 else "(4)"} to show QR code' in story - - need_keypress(KEY_QR if is_q1 else "4") - time.sleep(.1) - wif_qr = cap_screen_qr().decode() - press_cancel() - - if nfc_is_enabled(): - assert f"{KEY_NFC if is_q1 else '(3)'} to share via NFC" in story - - press_nfc() - time.sleep(0.3) - nfc_wif = nfc_read_text() - time.sleep(0.3) - press_cancel() - assert nfc_wif == wif - - - assert story_wif == wif == wif_qr - assert story_sk == bytes(sk).hex() - assert story_pk == sk.K.sec().hex() - - press_cancel() # exit Detail - press_cancel() # exit WIF submenu - - -@pytest.mark.parametrize("netcode", ["XTN", "BTC"]) -def test_wif_store_addresses(netcode, import_wif_to_store, use_mainnet, cap_menu, pick_menu_item, - cap_story, need_keypress, settings_remove, cap_screen_qr, is_q1, - nfc_is_enabled, press_nfc, nfc_read_text, goto_home, press_cancel): - goto_home() - if netcode == "BTC": - use_mainnet() - - settings_remove("wifs") - - prefix = bytes([128]) if netcode == "BTC" else bytes([239]) - n = BIP32Node.from_master_secret(os.urandom(32)) - privkey = n.node.private_key - - wif_list = [ - encode_base58_checksum(prefix + bytes(privkey) + b'\x01') - ] - - import_wif_to_store(wif_list) - - time.sleep(.1) - menu = cap_menu() - assert len(menu) == 2 - pick_menu_item(menu[1]) - pick_menu_item("Addresses") - - for mi, af in [("P2SH-Segwit", "p2sh-p2wpkh"), ("Segwit P2WPKH", "p2wpkh"), ("Classic P2PKH", "p2pkh")]: - pick_menu_item(mi) - time.sleep(.1) - title, story = cap_story() - if is_q1: - # Q has title as it needs hint keys - assert title == mi - - target_addr = n.address(addr_fmt=af, netcode=netcode) - addr = addr_from_display_format(story.split("\n\n")[0]) - assert addr == target_addr - - if not is_q1: - assert "Press (1) to show address QR code." in story - - need_keypress(KEY_QR if is_q1 else "1") - time.sleep(.1) - qr_addr = cap_screen_qr().decode() - if af == "p2wpkh": - qr_addr = qr_addr.lower() - press_cancel() - assert qr_addr == target_addr - - if nfc_is_enabled(): - if not is_q1: - assert "(3) to share via NFC." in story - - press_nfc() - time.sleep(0.3) - nfc_addr = nfc_read_text() - time.sleep(0.3) - press_cancel() - assert nfc_addr == target_addr - - press_cancel() - press_cancel() - press_cancel() - - -def test_wif_store_clear_all(import_wif_to_store, press_select, cap_story, settings_get, - need_keypress, cap_menu, settings_remove, is_q1, goto_home): - - goto_home() - settings_remove("wifs") - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - for _ in range(30) # MAX - ] - import_wif_to_store(wif_list) - time.sleep(.1) - - menu = cap_menu() - assert "Import WIF" not in menu # WIF store is full - assert "Clear All" in menu - need_keypress(KEY_UP if is_q1 else "5") - time.sleep(.1) - press_select() # on Clear All - time.sleep(.1) - title, story = cap_story() - assert "Remove all saved WIF keys?" in story - assert "(4)" in story - press_select() # does not work & gets you back to menu - assert len(settings_get("wifs")) == 30 - - press_select() # on Clear All - time.sleep(.1) - need_keypress("4") - time.sleep(.1) - menu = cap_menu() - assert len(menu) == 2 - assert "(none yet)" in menu - assert "Import WIF" in menu - assert not settings_get("wifs") - - -def test_wif_store_capacity(import_wif_to_store, settings_remove, press_select, cap_story, - settings_get, cap_menu, pick_menu_item, need_keypress): - settings_remove("wifs") - - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - for _ in range(40) # MAX+1 - ] - - import_wif_to_store(wif_list[:31], early_exit=True) - time.sleep(.1) - title, story = cap_story() - assert title == "Failure" - assert "Max 30 items allowed in WIF Store" in story - assert "Attempted to import 31 keys" in story - assert "remaining WIF store capacity is only 30" - press_select() - assert not settings_get("wifs") - - # import 29 keys - import_wif_to_store(wif_list[:29]) - - assert len(settings_get("wifs", [])) == 29 - - import_wif_to_store(wif_list[-2:], early_exit=True) - time.sleep(.1) - title, story = cap_story() - assert title == "Failure" - assert "Max 30 items allowed in WIF Store" in story - assert "Attempted to import 2 keys" in story - assert "remaining WIF store capacity is only 1" - press_select() - - assert len(settings_get("wifs", [])) == 29 - - import_wif_to_store(wif_list[-1:]) - assert len(settings_get("wifs", [])) == 30 - - menu = cap_menu() - assert "Import WIF" not in menu - # remove random key to make space - # pick key at current menu item position - press_select() - time.sleep(.1) - pick_menu_item("Delete") - time.sleep(.1) - title, story = cap_story() - assert "Delete WIF key?" in story - press_select() - time.sleep(.1) - menu = cap_menu() - assert "Import WIF" in menu - - -def test_wif_store_import_duplicate(settings_remove, import_wif_to_store, settings_get, cap_menu, - goto_home): - goto_home() - settings_remove("wifs") - - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - for _ in range(4) - ] - - import_wif_to_store(wif_list) - b4 = cap_menu() - assert len(settings_get("wifs")) == 4 - - import_wif_to_store(wif_list, early_exit=True) - assert len(settings_get("wifs")) == 4 - assert len(b4) == len(cap_menu()) - - import_wif_to_store(wif_list[:1], early_exit=True) - assert len(b4) == len(cap_menu()) - assert len(settings_get("wifs")) == 4 - - -@pytest.mark.parametrize("way", ["qr", "sd", "nfc"]) -def test_wif_store_export_all(way, goto_home, settings_remove, import_wif_to_store, pick_menu_item, - load_export): - goto_home() - settings_remove("wifs") - - wif_list = [ - encode_base58_checksum(bytes([239]) + os.urandom(32) + b'\x01') - for _ in range(6) # 6*52 chars so it can be shown on mk4 too - ] - - import_wif_to_store(wif_list) - time.sleep(.1) - pick_menu_item("Export All") - conts = load_export(way, "WIF Store", is_json=False, sig_check=False) - - assert wif_list == conts.split("\n") - - @pytest.mark.onetime def test_dump_menutree(sim_execfile): # saves to ../unix/work/menudump.txt diff --git a/testing/test_wif.py b/testing/test_wif.py new file mode 100644 index 00000000..db277729 --- /dev/null +++ b/testing/test_wif.py @@ -0,0 +1,368 @@ +# (c) Copyright 2026 by Coinkite Inc. This file is covered by license found in COPYING-CC. +# +import pytest, time, os, re, hashlib, shutil +from helpers import prandom, addr_from_display_format +from charcodes import KEY_DOWN, KEY_QR, KEY_NFC, KEY_DELETE, KEY_UP +from bip32 import BIP32Node, PrivateKey +from base58 import encode_base58_checksum + + +def make_fake_wif(prefix=239): + # generate a WIF + return encode_base58_checksum(bytes([prefix]) + prandom(32) + b'\x01') + +@pytest.mark.parametrize("num_wifs", [1, 11]) +@pytest.mark.parametrize("separator", ["\n", ',']) +@pytest.mark.parametrize("way", ["sd", "nfc", "qr", "vdisk"]) +def test_wif_store_import(num_wifs, separator, way, import_wif_to_store, skip_if_useless_way, + settings_remove, goto_home): + skip_if_useless_way(way) + settings_remove("wifs") + + wif_list = [make_fake_wif() for _ in range(num_wifs)] + + import_wif_to_store(wif_list, way=way, sep=separator) + goto_home() + + +def test_wif_store_import_manual(import_wif_to_store, settings_remove, goto_home): + settings_remove("wifs") + + wif_list = [make_fake_wif()] + + import_wif_to_store(wif_list, way="input") + goto_home() + + +def test_wif_store_import_paper_wallet(goto_home, pick_menu_item, press_select, cap_story, + need_keypress, settings_remove, microsd_path, cap_menu): + settings_remove("wifs") + goto_home() + pick_menu_item('Advanced/Tools') + try: + pick_menu_item('Paper Wallets') + except: + raise pytest.skip('Feature absent') + + press_select() + pick_menu_item('GENERATE WALLET') + + time.sleep(0.1) + title, story = cap_story() + if "Press (1) to save paper wallet file to SD Card" in story: + need_keypress("1") + time.sleep(0.2) + title, story = cap_story() + assert 'Created file' in story + + story = [i for i in story.split('\n') if i] + fname = story[-2] + assert fname.endswith('.txt') + + with open(microsd_path(fname), "r") as f: + const = f.read() + + goto_home() + pick_menu_item('Advanced/Tools') + pick_menu_item("WIF Store") + time.sleep(.1) + title, story = cap_story() + if title == "WIF Store": + press_select() + pick_menu_item("Import WIF") + need_keypress("1") # SD + try: + pick_menu_item(fname) + except: + pass + + menu = cap_menu() + assert "Import WIF" in menu + pick_menu_item(menu[1]) + pick_menu_item("Detail") + time.sleep(.1) + title, story = cap_story() + assert story.split("\n\n")[0] == const.split("\n\n")[4].strip() + + +@pytest.mark.parametrize("wif,err,way", [ + ("Ky2BtsR8qRN91PjktxaTQWMgJZUWSBJLjwip642vvoNyH1PeEpUP", "chain", "qr"), # mainnet key on testnet + ("91zb4oYGEvwEroihAbkdeoBpLSKnZYMdD1CPhfQD76fxrfNSp5J", "compressed only", "sd"), # uncompressed pk + ("cWALDjUu1tszsCBMjBjL4mhYj2wHUWYDR8Q8aSjLKzjkWaXMLRaY", None, "sd"), # curve order + ("cMahea7zqjxrtgAbB7LSGbcQUr1uX1ojuat9jZodMN87J7g8rY9t", None, "nfc"), # zero + ("cPPBMnQzGV4QAqD2HNPamprjvnmv6dQ2oysHCUVSRv2yXkVvWVtX", None, "nfc"), # wrong csum + ("cPPBMnQzGV4QAqD2HNPamprjvnmv6dQ2oysHCUVSRv2yXkVvWVtX;cN7M6sNzn4LGBxAozsmphxjuxVNaHcLre7Nm163qM3DpY3BZog1v", None, "sd"), # wrong separator +]) +def test_wif_store_import_fail(way, wif, err, import_wif_to_store, skip_if_useless_way, + settings_remove, press_select, cap_story, use_testnet, settings_get): + + err = err or "no valid WIF found" + skip_if_useless_way(way) + use_testnet() + settings_remove("wifs") + + import_wif_to_store([wif], way=way, early_exit=True) + time.sleep(.1) + title, story = cap_story() + assert "Failed to import WIF" in story + assert err in story + press_select() + assert not settings_get("wifs") + + +@pytest.mark.parametrize("netcode", ["XTN", "BTC"]) +def test_wif_store_detail(netcode, import_wif_to_store, use_mainnet, cap_menu, pick_menu_item, + cap_story, need_keypress, settings_remove, cap_screen_qr, is_q1, + press_cancel, nfc_is_enabled, press_nfc, nfc_read_text, goto_home): + goto_home() + if netcode == "BTC": + use_mainnet() + + settings_remove("wifs") + + prefix = bytes([128]) if netcode == "BTC" else bytes([239]) + privkeys = [PrivateKey.parse(prandom(32)) for _ in range(5)] + wif_list = [encode_base58_checksum(prefix + bytes(sk) + b'\x01') for sk in privkeys] + + import_wif_to_store(wif_list) + + time.sleep(.1) + menu = cap_menu() + target_mi = [] + for mi in menu: + if "⋯" in mi: + target_mi.append(mi) + + assert len(target_mi) == len(wif_list) + for mi, wif, sk in zip(target_mi, wif_list, privkeys): + mi_split = mi.split(" ")[-1].split("⋯") + assert len(mi_split) == 2 + assert mi_split[0] in wif + assert mi_split[1] in wif + pick_menu_item(mi) + + time.sleep(.1) + menu = cap_menu() + assert menu[0] == "Detail" + assert menu[1] == "Addresses" + assert menu[2] == "Sign MSG" + assert menu[3] == "Delete" + + pick_menu_item("Detail") + + title, story = cap_story() + assert title == "WIF" + + split_story = story.split("\n\n") + story_wif = split_story[0] + story_sk = split_story[1].split("\n")[-1] + story_pk = split_story[2].split("\n")[-1] + + assert f'{KEY_QR if is_q1 else "(4)"} to show QR code' in story + + need_keypress(KEY_QR if is_q1 else "4") + time.sleep(.1) + wif_qr = cap_screen_qr().decode() + press_cancel() + + if nfc_is_enabled(): + assert f"{KEY_NFC if is_q1 else '(3)'} to share via NFC" in story + + press_nfc() + time.sleep(0.3) + nfc_wif = nfc_read_text() + time.sleep(0.3) + press_cancel() + assert nfc_wif == wif + + + assert story_wif == wif == wif_qr + assert story_sk == bytes(sk).hex() + assert story_pk == sk.K.sec().hex() + + press_cancel() # exit Detail + press_cancel() # exit WIF submenu + + +@pytest.mark.parametrize("netcode", ["XTN", "BTC"]) +def test_wif_store_addresses(netcode, import_wif_to_store, use_mainnet, cap_menu, pick_menu_item, + cap_story, need_keypress, settings_remove, cap_screen_qr, is_q1, + nfc_is_enabled, press_nfc, nfc_read_text, goto_home, press_cancel): + goto_home() + if netcode == "BTC": + use_mainnet() + + settings_remove("wifs") + + prefix = bytes([128]) if netcode == "BTC" else bytes([239]) + n = BIP32Node.from_master_secret(prandom(32)) + privkey = n.node.private_key + + wif_list = [ encode_base58_checksum(prefix + bytes(privkey) + b'\x01') ] + + import_wif_to_store(wif_list) + + time.sleep(.1) + menu = cap_menu() + assert len(menu) == 2 + pick_menu_item(menu[1]) + pick_menu_item("Addresses") + + for mi, af in [("P2SH-Segwit", "p2sh-p2wpkh"), ("Segwit P2WPKH", "p2wpkh"), ("Classic P2PKH", "p2pkh")]: + pick_menu_item(mi) + time.sleep(.1) + title, story = cap_story() + if is_q1: + # Q has title as it needs hint keys + assert title == mi + + target_addr = n.address(addr_fmt=af, netcode=netcode) + addr = addr_from_display_format(story.split("\n\n")[0]) + assert addr == target_addr + + if not is_q1: + assert "Press (1) to show address QR code." in story + + need_keypress(KEY_QR if is_q1 else "1") + time.sleep(.1) + qr_addr = cap_screen_qr().decode() + if af == "p2wpkh": + qr_addr = qr_addr.lower() + press_cancel() + assert qr_addr == target_addr + + if nfc_is_enabled(): + if not is_q1: + assert "(3) to share via NFC." in story + + press_nfc() + time.sleep(0.3) + nfc_addr = nfc_read_text() + time.sleep(0.3) + press_cancel() + assert nfc_addr == target_addr + + press_cancel() + press_cancel() + press_cancel() + + +def test_wif_store_clear_all(import_wif_to_store, press_select, cap_story, settings_get, + need_keypress, cap_menu, settings_remove, is_q1, goto_home): + + goto_home() + settings_remove("wifs") + wif_list = [make_fake_wif() for _ in range(30)] # 30 is the max + import_wif_to_store(wif_list) + time.sleep(.1) + + menu = cap_menu() + assert "Import WIF" not in menu # WIF store is full + assert "Clear All" in menu + need_keypress(KEY_UP if is_q1 else "5") + time.sleep(.1) + press_select() # on Clear All + time.sleep(.1) + title, story = cap_story() + assert "Remove all saved WIF keys?" in story + assert "(4)" in story + press_select() # does not work & gets you back to menu + assert len(settings_get("wifs")) == 30 + + press_select() # on Clear All + time.sleep(.1) + need_keypress("4") + time.sleep(.1) + menu = cap_menu() + assert len(menu) == 2 + assert "(none yet)" in menu + assert "Import WIF" in menu + assert not settings_get("wifs") + + +def test_wif_store_capacity(import_wif_to_store, settings_remove, press_select, cap_story, + settings_get, cap_menu, pick_menu_item, need_keypress): + settings_remove("wifs") + + wif_list = [make_fake_wif() for _ in range(40)] # MAX+1 + + import_wif_to_store(wif_list[:31], early_exit=True) + time.sleep(.1) + title, story = cap_story() + assert title == "Failure" + assert "Max 30 items allowed in WIF Store" in story + assert "Attempted to import 31 keys" in story + assert "remaining WIF store capacity is only 30" + press_select() + assert not settings_get("wifs") + + # import 29 keys + import_wif_to_store(wif_list[:29]) + + assert len(settings_get("wifs", [])) == 29 + + import_wif_to_store(wif_list[-2:], early_exit=True) + time.sleep(.1) + title, story = cap_story() + assert title == "Failure" + assert "Max 30 items allowed in WIF Store" in story + assert "Attempted to import 2 keys" in story + assert "remaining WIF store capacity is only 1" + press_select() + + assert len(settings_get("wifs", [])) == 29 + + import_wif_to_store(wif_list[-1:]) + assert len(settings_get("wifs", [])) == 30 + + menu = cap_menu() + assert "Import WIF" not in menu + # remove random key to make space + # pick key at current menu item position + press_select() + time.sleep(.1) + pick_menu_item("Delete") + time.sleep(.1) + title, story = cap_story() + assert "Delete WIF key?" in story + press_select() + time.sleep(.1) + menu = cap_menu() + assert "Import WIF" in menu + + +def test_wif_store_import_duplicate(settings_remove, import_wif_to_store, settings_get, cap_menu, cap_story, + goto_home): + goto_home() + settings_remove("wifs") + + wif_list = [make_fake_wif() for _ in range(4)] + + import_wif_to_store(wif_list) + b4 = cap_menu() + assert len(settings_get("wifs")) == 4 + + import_wif_to_store(wif_list, early_exit=True) + assert len(settings_get("wifs")) == 4 + assert len(b4) == len(cap_menu()) + + title, story = cap_story() + assert 'duplicate WIF' in story + + +@pytest.mark.parametrize("way", ["qr", "sd", "nfc"]) +def test_wif_store_export_all(way, goto_home, settings_remove, import_wif_to_store, pick_menu_item, + load_export): + goto_home() + settings_remove("wifs") + + wif_list = [make_fake_wif() for _ in range(6)] # 6*52 chars so it can be shown on mk4 too + + import_wif_to_store(wif_list) + time.sleep(.1) + pick_menu_item("Export All") + conts = load_export(way, "WIF Store", is_json=False, sig_check=False) + + assert wif_list == conts.split("\n") + +