update for pymd3 async refactor

Author: leminlimez

src/devicemanagement/device_manager.py

@@ -1,4 +1,5 @@
 import traceback
+import asyncio
 import plistlib
 import time
 from tempfile import TemporaryDirectory
@@ -111,7 +112,7 @@ def get_devices(self, settings: QSettings, show_alert=lambda x: None):
             self.pref_manager.settings = settings
         # handle errors when failing to get connected devices
         try:
-            connected_devices = usbmux.list_devices()
+            connected_devices = asyncio.run(usbmux.list_devices())
         except Exception:
             sysmsg = QCoreApplication.tr("If you are on Linux, make sure you have usbmuxd and libimobiledevice installed.")
             if os.name == 'nt':
@@ -125,7 +126,7 @@ def get_devices(self, settings: QSettings, show_alert=lambda x: None):
         for device in connected_devices:
             if self.pref_manager.apply_over_wifi or device.is_usb:
                 try:
-                    ld = create_using_usbmux(serial=device.serial)
+                    ld = asyncio.run(create_using_usbmux(serial=device.serial))
                     vals = ld.all_values
                     model = vals['ProductType']
                     hardware = vals['HardwareModel']
@@ -268,7 +269,7 @@ def current_device_books_container_uuid_callback(self, uuid: Optional[str]=None)
         self.pref_manager.settings.setValue(self.data_singleton.current_device.udid + "_books_container_uuid", uuid)
 
     def get_app_hashes(self, bundle_ids: list[str]) -> dict:
-        apps = InstallationProxyService(lockdown=self.data_singleton.current_device.ld).get_apps(application_type="Any", calculate_sizes=False)
+        apps = asyncio.run(InstallationProxyService(lockdown=self.data_singleton.current_device.ld).get_apps(application_type="Any", calculate_sizes=False))
         results = {}
         for bundle_id in bundle_ids:
             app_info = apps[bundle_id]
@@ -280,7 +281,7 @@ def send_app_hashes_afc(self, hashes: dict) -> str:
         with TemporaryDirectory() as tmpdir:
             # get the bundle id of Pocket Poster
             bundle_id = "com.leemin.Pocket-Poster"
-            apps = InstallationProxyService(lockdown=self.data_singleton.current_device.ld).get_apps(application_type="User", calculate_sizes=False)
+            apps = asyncio.run(InstallationProxyService(lockdown=self.data_singleton.current_device.ld).get_apps(application_type="User", calculate_sizes=False))
             for app in apps.values():
                 if app["CFBundleExecutable"] == "Pocket Poster":
                     bundle_id = app["CFBundleIdentifier"]
@@ -295,7 +296,7 @@ def send_app_hashes_afc(self, hashes: dict) -> str:
                 tmpf = os.path.join(tmpdir, fname)
                 with open(tmpf, "w", encoding='UTF-8') as in_file:
                     in_file.write(hashes[key])
-                afc.push(tmpf, f"/Documents/{fname}")
+                asyncio.run(afc.push(tmpf, f"/Documents/{fname}"))
 
 
     def reset_device_pairing(self):
@@ -543,7 +544,7 @@ def start_restore(self, files_to_restore: list[FileToRestore], use_bookrestore:
                 connected = False
                 while not connected and max_timeout >= time.time():
                     try:
-                        new_ld = create_using_usbmux(serial=self.get_current_device_udid(), pair_timeout=180)
+                        new_ld = asyncio.run(create_using_usbmux(serial=self.get_current_device_udid(), pair_timeout=180))
                         connected = True
                     except Exception:
                         pass

src/restore/__init__.py

@@ -1,5 +1,6 @@
 from tempfile import TemporaryDirectory
 from pathlib import Path
+import asyncio
 
 from pymobiledevice3.lockdown import create_using_usbmux
 from pymobiledevice3.services.mobilebackup2 import Mobilebackup2Service
@@ -13,7 +14,7 @@ def reboot_device(reboot: bool = False, lockdown_client: LockdownClient = None):
     if reboot and lockdown_client != None:
         print("Success! Rebooting your device...")
         with DiagnosticsService(lockdown_client) as diagnostics_service:
-            diagnostics_service.restart()
+            asyncio.run(diagnostics_service.restart())
         print("Remember to turn Find My back on!")
 
 def perform_restore(backup: backup.Backup, reboot: bool = False, lockdown_client: LockdownClient = None, progress_callback = lambda x: None):
@@ -22,9 +23,9 @@ def perform_restore(backup: backup.Backup, reboot: bool = False, lockdown_client
             backup.write_to_directory(Path(backup_dir))
 
             if lockdown_client == None:
-                lockdown_client = create_using_usbmux()
+                lockdown_client = asyncio.run(create_using_usbmux())
             with Mobilebackup2Service(lockdown_client) as mb:
-                mb.restore(backup_dir, system=True, reboot=False, copy=False, source=".", progress_callback=progress_callback, skip_apps=True)
+                asyncio.run(mb.restore(backup_dir, system=True, reboot=False, copy=False, source=".", progress_callback=progress_callback, skip_apps=True))
             # reboot the device
             reboot_device(reboot, lockdown_client)
     except PyMobileDevice3Exception as e:

src/restore/bookrestore.py

@@ -24,7 +24,7 @@
 from pymobiledevice3.services.afc import AfcService
 from pymobiledevice3.services.amfi import AmfiService
 from pymobiledevice3.remote.remote_service_discovery import RemoteServiceDiscoveryService
-from pymobiledevice3.services.dvt.dvt_secure_socket_proxy import DvtSecureSocketProxyService
+from pymobiledevice3.services.dvt.instruments.dvt_provider import DvtProvider
 from pymobiledevice3.services.dvt.instruments.process_control import ProcessControl
 from pymobiledevice3.services.os_trace import OsTraceService
 from http.server import HTTPServer, SimpleHTTPRequestHandler
@@ -178,7 +178,7 @@ async def async_connection():
                     loop = asyncio.get_running_loop()
 
                     def run_blocking_callback():
-                        with DvtSecureSocketProxyService(rsd) as dvt:
+                        with DvtProvider(rsd) as dvt:
                             apply_bookrestore_files(files, rsd, dvt, current_device_uuid_callback, progress, transfer_mode, do_full_reboot)
 
                     await loop.run_in_executor(None, run_blocking_callback)
@@ -265,7 +265,7 @@ def generate_bldbmanager(files: list[FileToRestore], out_file: str, afc: AfcServ
                 zassetpath = f'{file.restore_path}.zassetpath'
                 media_folder = file_name#f'{nugget_media_folder}/{file_name}'
                 zplistpath = f'/var/mobile/Media/{media_folder}'
-                afc.set_file_contents(media_folder, file.contents)
+                asyncio.run(afc.set_file_contents(media_folder, file.contents))
             else:
                 zdownloadid = ""
                 zassetpath = file.restore_path
@@ -283,7 +283,7 @@ def generate_bldbmanager(files: list[FileToRestore], out_file: str, afc: AfcServ
     # return the number of files in the thing
     return z_id
 
-def apply_bookrestore_files(files: list[FileToRestore], lockdown_client: LockdownClient, dvt: DvtSecureSocketProxyService,
+def apply_bookrestore_files(files: list[FileToRestore], lockdown_client: LockdownClient, dvt: DvtProvider,
                             current_device_uuid_callback = lambda x: None, progress_callback = lambda x: None,
                             transfer_mode: BookRestoreFileTransferMethod = BookRestoreFileTransferMethod.LocalHost,
                             do_full_reboot: bool = False):
@@ -369,13 +369,13 @@ def apply_bookrestore_files(files: list[FileToRestore], lockdown_client: Lockdow
         """)
         connection.commit()
 
-        procs = OsTraceService(lockdown=lockdown_client).get_pid_list().get("Payload")
+        procs = asyncio.run(OsTraceService(lockdown=lockdown_client).get_pid_list()).get("Payload")
         pid_bookassetd = next((pid for pid, p in procs.items() if p['ProcessName'] == 'bookassetd'), None)
         pid_books = next((pid for pid, p in procs.items() if p['ProcessName'] == 'Books'), None)
         if pid_bookassetd:
-            pc.signal(pid_bookassetd, 19)
+            asyncio.run(pc.signal(pid_bookassetd, 19))
         if pid_books:
-            pc.kill(pid_books)
+            asyncio.run(pc.kill(pid_books))
 
         progress_callback("Uploading files...")
 
@@ -389,7 +389,7 @@ def apply_bookrestore_files(files: list[FileToRestore], lockdown_client: Lockdow
                 _, file_name = os.path.split(file.restore_path)
                 print(f"including {file.restore_path}")
                 media_folder = file_name
-                afc.set_file_contents(media_folder, file.contents)
+                asyncio.run(afc.set_file_contents(media_folder, file.contents))
 
         def fast_upload(local_path, remote_path):
             content = b''
@@ -399,7 +399,7 @@ def fast_upload(local_path, remote_path):
                         content = f.read()
                 except OSError:
                     content = b''
-            afc.set_file_contents(remote_path, content)
+            asyncio.run(afc.set_file_contents(remote_path, content))
 
         fast_upload(temp_db_path, "Downloads/downloads.28.sqlitedb")
         fast_upload(temp_db_path + "-shm", "Downloads/downloads.28.sqlitedb-shm")
@@ -418,10 +418,10 @@ def fast_upload(local_path, remote_path):
             except Exception:
                 pass
 
-    procs = OsTraceService(lockdown=lockdown_client).get_pid_list().get("Payload")
+    procs = asyncio.run(OsTraceService(lockdown=lockdown_client).get_pid_list()).get("Payload")
     pid_itunesstored = next((pid for pid, p in procs.items() if p['ProcessName'] == 'itunesstored'), None)
     if pid_itunesstored:
-        pc.kill(pid_itunesstored)
+        asyncio.run(pc.kill(pid_itunesstored))
 
     timeout = time.time() + 120 
     progress_callback("Waiting for itunesstored to finish download..." + "\n" + "(This might take a minute)")
@@ -435,12 +435,12 @@ def fast_upload(local_path, remote_path):
     pid_bookassetd = next((pid for pid, p in procs.items() if p['ProcessName'] == 'bookassetd'), None)
     pid_books = next((pid for pid, p in procs.items() if p['ProcessName'] == 'Books'), None)
     if pid_bookassetd:
-        pc.kill(pid_bookassetd)
+        asyncio.run(pc.kill(pid_bookassetd))
     if pid_books:
-        pc.kill(pid_books)
+        asyncio.run(pc.kill(pid_books))
 
     try:
-        pc.launch("com.apple.iBooks")
+        asyncio.run(pc.launch("com.apple.iBooks"))
     except Exception as e:
         raise NuggetException("Error launching Books app", detailed_text=repr(e))
 
@@ -461,7 +461,7 @@ def fast_upload(local_path, remote_path):
             # respring anyway even if it is not detected that all files overwrote
             break
             # raise Exception("Timed out waiting for file, please try again.")
-    pc.kill(pid_bookassetd)
+    asyncio.run(pc.kill(pid_bookassetd))
     if transfer_mode == BookRestoreFileTransferMethod.LocalHost:
         close_dl_connection()
         remove_db_files(temp_dl_manager)
@@ -471,9 +471,9 @@ def fast_upload(local_path, remote_path):
         reboot_device(True, lockdown_client=lockdown_client)
     else:
         progress_callback("Respringing")
-        procs = OsTraceService(lockdown=lockdown_client).get_pid_list().get("Payload")
+        procs = asyncio.run(OsTraceService(lockdown=lockdown_client).get_pid_list()).get("Payload")
         pid = next((pid for pid, p in procs.items() if p['ProcessName'] == 'backboardd'), None)
-        pc.kill(pid)
+        asyncio.run(pc.kill(pid))
 
 def perform_bookrestore(files: list[FileToRestore], lockdown_client: LockdownClient,
                         current_device_books_uuid_callback = lambda x: None, progress_callback = lambda x: None,
@@ -482,7 +482,7 @@ def perform_bookrestore(files: list[FileToRestore], lockdown_client: LockdownCli
     if not lockdown_client.developer_mode_status:
         # enable developer mode
         progress_callback("Enabling Developer Mode...")
-        AmfiService(lockdown=lockdown_client).reveal_developer_mode_option_in_ui()
+        asyncio.run(AmfiService(lockdown=lockdown_client).reveal_developer_mode_option_in_ui())
         raise NuggetException("You must enable developer mode on your device. You can do it in the Settings app.\n\nClick \"Show Details\" for more information.",
                               detailed_text="BookRestore tweaks with the AFC method require developer mode to apply.\n\nYou can enable this at the bottom of Settings > Privacy & Security > Developer Mode on your iPhone or iPad.")
     if os.name == 'nt':

src/restore/restore.py

@@ -6,6 +6,7 @@
 import os
 import plistlib
 import ssl
+import asyncio
 
 class FileToRestore:
     def __init__(self,
@@ -167,7 +168,7 @@ def restore_files(files: list[FileToRestore], reboot: bool = False, lockdown_cli
                 bundle_id = last_domain.removeprefix("AppDomain-")
                 if not bundle_id in active_bundle_ids:
                     if apps == None:
-                        apps = InstallationProxyService(lockdown=lockdown_client).get_apps(application_type="Any", calculate_sizes=False)
+                        apps = asyncio.run(InstallationProxyService(lockdown=lockdown_client).get_apps(application_type="Any", calculate_sizes=False))
                     app_info = apps[bundle_id]
                     active_bundle_ids.append(bundle_id)
                     apps_list.append(backup.AppBundle(