wallet_utils.py: add testdkg command

docs/frost-wallet-dev.md

@@ -72,6 +72,7 @@ to be run permanently)
 - `dkgrm`: rm FrostWallet DKG data by `session_id` list
 - `recdkgls`: display Recovery DKG File data
 - `recdkgrm`: rm Recovery DKG File data by `session_id` list
+- `testdkg`: run only as test of DKG process
 - `testfrost`: run only as test of FROST signing
 
 ## Description of `jmclient/frost_clients.py`

src/jmclient/client_protocol.py

@@ -405,13 +405,12 @@ class JMClientProtocol(BaseClientProtocol):
 
     """DKG specifics
     """
-    async def dkg_gen(self):
+    async def dkg_gen(self, session_id=None):
         jlog.debug('Coordinator call dkg_gen')
         client = self.factory.client
         md_type_idx = None
-        session_id = None
         session = None
-
+        pub = None
         while True:
             if md_type_idx is None:
                 md_type_idx = await client.dkg_gen()
@@ -419,8 +418,9 @@ class JMClientProtocol(BaseClientProtocol):
                     jlog.debug('finished dkg_gen execution')
                     break
 
-            if session_id is None:
+            if session_id is None or session_id == b'\x00'*32:
-                session_id, _, session = self.dkg_init(*md_type_idx)
+                session_id, _, session = self.dkg_init(
+                    *md_type_idx, session_id=session_id)
                 if session_id is None:
                     jlog.warning('could not get session_id from dkg_init}')
                     await asyncio.sleep(5)
@@ -438,13 +438,14 @@ class JMClientProtocol(BaseClientProtocol):
                 session = None
                 client.dkg_gen_list.pop(0)
                 continue
+        return pub
 
-    def dkg_init(self, mixdepth, address_type, index):
+    def dkg_init(self, mixdepth, address_type, index, session_id=None):
         jlog.debug(f'Coordinator call dkg_init '
                    f'({mixdepth}, {address_type}, {index})')
         client = self.factory.client
-        hostpubkeyhash, session_id, sig = client.dkg_init(mixdepth,
+        hostpubkeyhash, session_id, sig = client.dkg_init(
-                                                          address_type, index)
+            mixdepth, address_type, index, session_id=session_id)
         coordinator = client.dkg_coordinators.get(session_id)
         session = client.dkg_sessions.get(session_id)
         if session_id and session and coordinator:

src/jmclient/frost_clients.py

@@ -155,13 +155,14 @@ class DKGClient:
             self.current_dkg_gen = None
         return self.current_dkg_gen
 
-    def dkg_init(self, mixdepth, address_type, index):
+    def dkg_init(self, mixdepth, address_type, index, session_id=None):
         try:
             wallet = self.wallet_service.wallet
             hostseckey = wallet._hostseckey[:32]
             hostpubkey = hostpubkey_gen(hostseckey)
             hostpubkeyhash = sha256(hostpubkey).digest()
-            session_id = sha256(os.urandom(32)).digest()
+            if session_id is None:
+                session_id = sha256(os.urandom(32)).digest()
             coordinator = DKGCoordinator(mixdepth=mixdepth,
                                          address_type=address_type,
                                          index=index,
@@ -392,6 +393,10 @@ class DKGClient:
             session_id = session.session_id
             coordinator = self.dkg_coordinators.get(session_id)
             coord_hostpubkey = session.coord_hostpubkey
+            if not coordinator:
+                self.dkg_sessions.pop(session_id)
+            if session_id == b'\x00'*32:
+                return True
             if coordinator:
                 dkg_man.add_coordinator_data(
                     session_id=session_id,
@@ -408,7 +413,6 @@ class DKGClient:
                     t=self.t,
                     recovery_data=session.recovery_data,
                     ext_recovery=ext_recovery)
-                self.dkg_sessions.pop(session_id)
             return True
         except Exception as e:
             jlog.error(f'finalize: {repr(e)}')

src/jmclient/frost_ipc.py

@@ -86,18 +86,30 @@ class FrostIPCServer(IPCBase):
                 jlog.error(f'FrostIPCServer.process_msgs: {repr(e)}')
             await asyncio.sleep(0.1)
 
-    async def on_get_dkg_pubkey(self, msg_id, mixdepth, address_type, index):
+    async def on_get_dkg_pubkey(self, msg_id, mixdepth, address_type, index,
+                                session_id=None):
         try:
             wallet = self.wallet
             dkg = wallet.dkg
-            new_pubkey = dkg.find_dkg_pubkey(mixdepth, address_type, index)
+            if session_id is not None:
-            if new_pubkey is None:
                 client = wallet.client_factory.getClient()
                 frost_client = wallet.client_factory.client
                 frost_client.dkg_gen_list.append(
                     (mixdepth, address_type, index))
-                await client.dkg_gen()
+                new_pubkey = await client.dkg_gen(session_id=session_id)
+            else:
                 new_pubkey = dkg.find_dkg_pubkey(mixdepth, address_type, index)
+            if session_id is None and new_pubkey is None:
+                client = wallet.client_factory.getClient()
+                frost_client = wallet.client_factory.client
+                frost_client.dkg_gen_list.append(
+                    (mixdepth, address_type, index))
+                client.dkg_gen()
+                if session_id == b'\x00'*32:
+                    new_pubkey = pub
+                else:
+                    new_pubkey = dkg.find_dkg_pubkey(
+                        mixdepth, address_type, index)
             if new_pubkey:
                 await self.send_dkg_pubkey(msg_id, new_pubkey)
             else:
@@ -203,7 +215,8 @@ class FrostIPCClient(IPCBase):
         if fut:
             fut.set_result(data)
 
-    async def get_dkg_pubkey(self, mixdepth, address_type, index):
+    async def get_dkg_pubkey(self, mixdepth, address_type, index,
+                             session_id=None):
         jlog.debug(f'FrostIPCClient.get_dkg_pubkey for mixdepth={mixdepth}, '
                    f'address_type={address_type}, index={index}')
         try:
@@ -211,7 +224,7 @@ class FrostIPCClient(IPCBase):
             msg_dict = {
                 'msg_id': self.msg_id,
                 'cmd': 'get_dkg_pubkey',
-                'data': (mixdepth, address_type, index),
+                'data': (mixdepth, address_type, index, session_id),
             }
             self.sw.write(self.encrypt_msg(msg_dict))
             await self.sw.drain()

src/jmclient/wallet_utils.py

@@ -74,6 +74,7 @@ The method is one of the following:
 (dkgrm) rm FrostWallet dkg data by session_id list
 (recdkgls) display Recovery DKG File data
 (recdkgrm) rm Recovery DKG File data by session_id list
+(testdkg) run only as test of DKG process
 (testfrost) run only as test of FROST signing
 """
     parser = OptionParser(usage='usage: %prog [options] [wallet file] [method] [args..]',
@@ -1785,8 +1786,9 @@ async def wallet_tool_main(wallet_root_path):
                               'dkgls', 'dkgrm', 'recdkgls', 'recdkgrm']
     frost_noscan_methods = ['hostpubkey', 'servefrost', 'dkgrecover',
                             'dkgls', 'dkgrm', 'recdkgls', 'recdkgrm',
-                            'testfrost']
+                            'testdkg', 'testfrost']
-    frost_readonly_methods = ['hostpubkey', 'dkgls', 'recdkgls', 'testfrost']
+    frost_readonly_methods = ['hostpubkey', 'dkgls', 'recdkgls',
+                              'testdkg', 'testfrost']
     noscan_methods.extend(frost_noscan_methods)
     readonly_methods.extend(frost_readonly_methods)
 
@@ -1965,6 +1967,14 @@ async def wallet_tool_main(wallet_root_path):
         await ipc_server.async_init()
         await ipc_server.serve_forever()
         return
+    elif method == "testdkg":
+        if not isinstance(wallet, FrostWallet):
+            return 'Command "testdgk" used only for FROST wallets'
+        md = address_type = index = 0
+        pubkey = await wallet.ipc_client.get_dkg_pubkey(
+            md, address_type, index, session_id=b'\x00'*32)
+        if pubkey:
+            return f'pubkey: {pubkey.hex()}'
     elif method == "testfrost":
         if not isinstance(wallet, FrostWallet):
             return 'Command "testfrost" used only for FROST wallets'