From a0cb49fd77cffae3431275878e1a36b0b23ce141 Mon Sep 17 00:00:00 2001
From: Jean-Marie Mineau <jean-marie.mineau@inria.fr>
Date: Wed, 2 Apr 2025 11:17:56 +0200
Subject: [PATCH] finally

---
 frida/theseus_frida/__init__.py | 79 +++++++++++++++++++++++++--------
 frida/theseus_frida/hook.js     |  3 +-
 2 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/frida/theseus_frida/__init__.py b/frida/theseus_frida/__init__.py
index 5272d20..8c17877 100644
--- a/frida/theseus_frida/__init__.py
+++ b/frida/theseus_frida/__init__.py
@@ -32,6 +32,9 @@ def spinner(symbs: str = "-\\|/"):
             yield s
 
 
+CLASSLOADER_DONE = False
+
+
 # Define handler to event generated by the scripts
 def on_message(message, data, data_storage: dict, file_storage: Path):
     if message["type"] == "error":
@@ -47,6 +50,9 @@ def on_message(message, data, data_storage: dict, file_storage: Path):
         handle_load_dex(message["payload"]["data"], data_storage, file_storage)
     elif message["type"] == "send" and message["payload"]["type"] == "classloader":
         handle_classloader_data(message["payload"]["data"], data_storage)
+    elif message["type"] == "send" and message["payload"]["type"] == "classloader-done":
+        global CLASSLOADER_DONE
+        CLASSLOADER_DONE = True
     else:
         print("[-] message:", message)
 
@@ -362,24 +368,65 @@ def collect_runtime(apk: Path, device_name: str, file_storage: Path, output: Tex
     # Resume the execution of the APK
     device.resume(pid)
 
-    script.post({"type": "dump-class-loaders"})
-    print("==> Press ENTER to finish the analysis <==")
+    print("==> Press ENTER to end the analysis <==")
     input()
 
+    # Dump all known classloaders
+    global CLASSLOADER_DONE
+    CLASSLOADER_DONE = False
+    script.post({"type": "dump-class-loaders"})
+    t = spinner()
+    while not CLASSLOADER_DONE:
+        print(
+            f"[{t.__next__()}] Waiting for the list of classloaders to be sent",
+            end="\r",
+        )
+        time.sleep(0.3)
+    print(f"[*] Classloader list received" + " " * 20)
+
     # Try to find the Main class loader
     main_class_loader: str | None = None
-    cls = {d["id"]: d for d in data_storage["classloaders"]}
-    for load_data in data_storage["dyn_code_load"]:
-        if load_data["classloader"] in cls:
-            del cls[load_data["classloader"]]
-    for id_ in list(cls.keys()):
-        if (
-            'dalvik.system.PathClassLoader[DexPathList[[directory "."],'
-            in cls[id_]["str"]
-        ):
-            del cls[id_]
-        elif cls[id_]["cname"] == "java.lang.BootClassLoader":
-            del cls[id_]
+    # cls = {d["id"]: d for d in data_storage["classloaders"]}
+    # for load_data in data_storage["dyn_code_load"]:
+    #     if load_data["classloader"] in cls:
+    #         del cls[load_data["classloader"]]
+    # for id_ in list(cls.keys()):
+    #     if (
+    #         'dalvik.system.PathClassLoader[DexPathList[[directory "."],'
+    #         in cls[id_]["str"]
+    #     ):
+    #         del cls[id_]
+    #     elif cls[id_]["cname"] == "java.lang.BootClassLoader":
+    #         del cls[id_]
+    cls = {}
+    for cl in data_storage["classloaders"]:
+        # This is verry doubious
+        if cl["cname"] == "dalvik.system.PathClassLoader":
+            zip_files = list(
+                map(
+                    lambda s: s.removeprefix('zip file "').removesuffix('"'),
+                    filter(
+                        lambda s: s.startswith('zip file "'),
+                        (
+                            w
+                            for b in cl["str"].split("]")
+                            for a in b.split("[")
+                            for w in a.split(",")
+                        ),
+                    ),
+                )
+            )
+            if len(zip_files) == 1:
+                zip_path = Path(zip_files[0])
+                if (
+                    len(zip_path.parts) == 6
+                    and zip_path.parts[0] == "/"
+                    and zip_path.parts[1] == "data"
+                    and zip_path.parts[2] == "app"
+                    and zip_path.parts[4].startswith(app + "-")
+                    and zip_path.parts[5] == "base.apk"
+                ):
+                    cls[cl["id"]] = cl
     if len(cls) == 0:
         print("[!] No classloader found for the main APK")
     elif len(cls) > 1:
@@ -401,10 +448,6 @@ def collect_runtime(apk: Path, device_name: str, file_storage: Path, output: Tex
         main_class_loader = list(cls.keys())[0]
     data_storage["apk_cl_id"] = main_class_loader
 
-    # Dump all known classloaders
-    script.post({"type": "dump-class-loaders"})
-    time.sleep(1)  # TODO: wait for ack from frida
-
     json.dump(data_storage, output, indent="  ")
 
 
diff --git a/frida/theseus_frida/hook.js b/frida/theseus_frida/hook.js
index 5b21e50..cdf3e1a 100644
--- a/frida/theseus_frida/hook.js
+++ b/frida/theseus_frida/hook.js
@@ -10,10 +10,11 @@ function dump_classloaders() {
         "cname": cl.$className
       }}); 
     }
+    send({"type": "classloader-done"})
   });
 }
 
-recv('dump-class-loaders', function onMessage(msg) {dump_classloaders()});
+// recv('dump-class-loaders', function onMessage(msg) {dump_classloaders()});
 
 Java.perform(() => {