fix dkg_recover, add test_dkg_recover

src/jmclient/wallet.py

@@ -589,13 +589,12 @@ class DKGManager:
         if save_dkg:
             self.save()
 
-    async def dkg_recover(self, dkgrec_path):
-        rec_storage = DKGRecoveryStorage(
-            dkgrec_path, create=False, read_only=True)
-        rec_dkg = rec_storage.data[self.RECOVERY_STORAGE_KEY]
+    def dkg_recover(self, dkgrec_storage):
+        rec_dkg = dkgrec_storage.data[self.RECOVERY_STORAGE_KEY]
         privkey = self.wallet._hostseckey
         wallet_rec_dkg = self.recovery_storage.data[self.RECOVERY_STORAGE_KEY]
-        jlog.info(f'Found {len(rec_dkg)} records in the {dkgrec_path}')
+        jlog.info(f'Found {len(rec_dkg)} records in the'
+                  f' {dkgrec_storage.path}')
         for session_id, (ext_recovery, recovery_data) in rec_dkg.items():
             jlog.info(f'Processing session_id {session_id.hex()}')
             try:
@@ -621,7 +620,7 @@ class DKGManager:
             self._dkg_hostpubkeys[session_id] = session_params.hostpubkeys
             self._dkg_t[session_id] = session_params.t
 
-            wallet_rec_dkg[session_id] = (ext_recovery, recovery_data)
+            wallet_rec_dkg[session_id] = [ext_recovery, recovery_data]
 
         self.save()
 
@@ -2966,6 +2965,7 @@ class LegacyWallet(ImportWalletMixin, PSBTWalletMixin, BIP32Wallet):
     def _get_bip32_base_path(self):
         return self._key_ident, 0
 
+
 class BIP32PurposedWallet(BIP32Wallet):
     """ A class to encapsulate cases like
     BIP44, 49 and 84, all of which are derivatives
@@ -2988,6 +2988,7 @@ class BIP32PurposedWallet(BIP32Wallet):
 
         return path[len(self._get_bip32_base_path())] - 2**31
 
+
 class FidelityBondMixin(object):
     BIP32_TIMELOCK_ID = 2
     BIP32_BURN_ID = 3

src/jmclient/wallet_utils.py

@@ -1958,7 +1958,9 @@ async def wallet_tool_main(wallet_root_path):
         if not isinstance(wallet, FrostWallet):
             return 'Command "dkgrecover" used only for FROST wallets'
         dkgrec_path = args[2]
-        return await wallet_service.dkg.dkg_recover(dkgrec_path)
+        dkgrec_storage = DKGRecoveryStorage(
+            dkgrec_path, create=False, read_only=True)
+        return wallet_service.dkg.dkg_recover(dkgrec_storage)
     elif method == "dkgls":
         if not isinstance(wallet, FrostWallet):
             return 'Command "dkgls" used only for FROST wallets'

test/jmclient/test_frost_wallet.py

@@ -15,11 +15,12 @@ import jmbitcoin as btc
 from jmbase import get_log
 from jmclient import (
     load_test_config, jm_single, VolatileStorage, get_network, cryptoengine,
-    create_wallet, open_test_wallet_maybe, FrostWallet, DKGManager)
+    create_wallet, open_test_wallet_maybe, FrostWallet, DKGManager,
+    WalletService)
 
 from jmfrost.chilldkg_ref.chilldkg import DKGOutput, hostpubkey_gen
 from jmclient.frost_clients import (
-    serialize_ext_recovery, decrypt_ext_recovery)
+    serialize_ext_recovery, decrypt_ext_recovery, DKGClient)
 
 pytestmark = pytest.mark.usefixtures("setup_regtest_frost_bitcoind")
 
@@ -28,12 +29,12 @@ test_create_wallet_filename = "frost_testwallet_for_create_wallet_test"
 log = get_log()
 
 
-async def get_populated_wallet():
+async def get_populated_wallet(entropy=None):
     storage = VolatileStorage()
     dkg_storage = VolatileStorage()
     recovery_storage = VolatileStorage()
     FrostWallet.initialize(storage, dkg_storage, recovery_storage,
-                           get_network())
+                           get_network(), entropy=entropy)
     wallet = FrostWallet(storage, dkg_storage, recovery_storage)
     await wallet.async_init(storage)
     return wallet
@@ -325,7 +326,86 @@ class AsyncioTestCase(IsolatedAsyncioTestCase):
         assert wlt.dkg.find_dkg_pubkey(0, 0, 1) is None
 
     async def test_dkg_recover(self):
-        assert 0  # FIXME need to add test
+        entropy1 = bytes.fromhex('8e5e5677fb302874a607b63ad03ba434')
+        entropy2 = bytes.fromhex('38dfa80fbb21b32b2b2740e00a47de9d')
+        entropy3 = bytes.fromhex('3ad9c77fcd1d537b6ef396952d1221a0')
+        wlt1 = await get_populated_wallet(entropy1)
+        hostpubkey1 = hostpubkey_gen(wlt1._hostseckey[:32])
+        wlt_svc1 = WalletService(wlt1)
+        wlt2 = await get_populated_wallet(entropy2)
+        hostpubkey2 = hostpubkey_gen(wlt2._hostseckey[:32])
+        wlt_svc2 = WalletService(wlt2)
+        wlt3 = await get_populated_wallet(entropy3)
+        hostpubkey3 = hostpubkey_gen(wlt3._hostseckey[:32])
+        wlt_svc3 = WalletService(wlt3)
+        nick1, nick2, nick3, nick4 = [
+            'nick1', 'nick2', 'nick3', 'nick4'
+        ]
+
+
+        dkgc1 = DKGClient(wlt_svc1)
+        dkgc2 = DKGClient(wlt_svc2)
+        dkgc3 = DKGClient(wlt_svc3)
+        hostpubkeyhash_hex, session_id, sig_hex = dkgc1.dkg_init(0, 0, 0)
+
+        (
+            nick1,
+            hostpubkeyhash2_hex,
+            session_id2_hex,
+            sig2_hex,
+            pmsg1_2
+        ) = dkgc2.on_dkg_init(
+            nick1, hostpubkeyhash_hex, session_id, sig_hex)
+        pmsg1_2 = dkgc2.deserialize_pmsg1(pmsg1_2)
+
+        (
+            nick1,
+            hostpubkeyhash3_hex,
+            session_id3_hex,
+            sig3_hex,
+            pmsg1_3
+        ) = dkgc3.on_dkg_init(
+            nick1, hostpubkeyhash_hex, session_id, sig_hex)
+        pmsg1_3 = dkgc2.deserialize_pmsg1(pmsg1_3)
+
+        ready_list, cmsg1 = dkgc1.on_dkg_pmsg1(
+            nick2, hostpubkeyhash2_hex, session_id, sig2_hex, pmsg1_2)
+        ready_list, cmsg1 = dkgc1.on_dkg_pmsg1(
+            nick3, hostpubkeyhash3_hex, session_id, sig3_hex, pmsg1_3)
+        cmsg1 = dkgc1.deserialize_cmsg1(cmsg1)
+
+        pmsg2_2 = dkgc2.party_step2(session_id, cmsg1)
+        pmsg2_2 = dkgc2.deserialize_pmsg2(pmsg2_2)
+        pmsg2_3 = dkgc3.party_step2(session_id, cmsg1)
+        pmsg2_3 = dkgc3.deserialize_pmsg2(pmsg2_3)
+
+        ready_list, cmsg2, ext_recovery = dkgc1.on_dkg_pmsg2(
+            nick2, session_id, pmsg2_2)
+        ready_list, cmsg2, ext_recovery = dkgc1.on_dkg_pmsg2(
+            nick3, session_id, pmsg2_3)
+        cmsg2 = dkgc3.deserialize_cmsg2(cmsg2)
+
+        assert dkgc2.finalize(session_id, cmsg2, ext_recovery)
+        assert dkgc3.finalize(session_id, cmsg2, ext_recovery)
+
+        assert not dkgc1.on_dkg_finalized(nick2, session_id)
+        assert dkgc1.on_dkg_finalized(nick3, session_id)
+
+        wlt_rec = await get_populated_wallet(entropy1)
+        wlt1._storage.data[b'created'] = wlt_rec._storage.data[b'created']
+        wlt1._dkg_storage.data[b'created'] = \
+            wlt_rec._dkg_storage.data[b'created']
+        wlt1._recovery_storage.data[b'created'] = \
+            wlt_rec._recovery_storage.data[b'created']
+        assert wlt1._storage.data == wlt_rec._storage.data  # empty wallet
+        assert wlt1._dkg_storage.data != wlt_rec._dkg_storage.data
+        assert wlt1._recovery_storage.data != wlt_rec._recovery_storage.data
+
+        wlt_rec.dkg.dkg_recover(wlt1._recovery_storage)
+
+        assert wlt1._storage.data == wlt_rec._storage.data
+        assert wlt1._dkg_storage.data == wlt_rec._dkg_storage.data
+        assert wlt1._recovery_storage.data == wlt_rec._recovery_storage.data
 
     async def test_dkg_ls(self):
         wlt = await get_populated_wallet()