read reflection data from json

frida/theseus_frida/__init__.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import subprocess
 import time
+import json
 from pathlib import Path
 
 import frida  # type: ignore
@@ -12,16 +13,16 @@ STACK_CONSUMER_B64 = Path(__file__).parent / "StackConsumer.dex.b64"
 
 
 # Define handler to event generated by the scripts
-def on_message(message, data):
+def on_message(message, data, data_storage: dict):
     if message["type"] == "error":
         print(f"[error] {message['description']}")
         print(message["stack"])
     elif message["type"] == "send" and message["payload"]["type"] == "invoke":
-        handle_invoke_data(message["payload"]["data"])
+        handle_invoke_data(message["payload"]["data"], data_storage)
     elif message["type"] == "send" and message["payload"]["type"] == "class-new-inst":
-        handle_class_new_inst_data(message["payload"]["data"])
+        handle_class_new_inst_data(message["payload"]["data"], data_storage)
     elif message["type"] == "send" and message["payload"]["type"] == "cnstr-new-isnt":
-        handle_cnstr_new_inst_data(message["payload"]["data"])
+        handle_cnstr_new_inst_data(message["payload"]["data"], data_storage)
     else:
         print("[on_message] message:", message)
 
@@ -52,42 +53,84 @@ def print_stack(stack, prefix: str):
 #    return f"{cls}->{name}({args}){ret}"
 
 
-def handle_invoke_data(data):
+def handle_invoke_data(data, data_storage: dict):
     method = data["method"]
-    # caller_method = "?"  # get_method_id(data["caller_method"])
-    # addr = data["addr"]
+    if len(data["stack"]) == 0:
+        return
+    caller_method = data["stack"][0]["method"]
+    addr = data["stack"][0]["bytecode_index"]
     print("Method.Invoke:")
     print(f"    called: {method}")
-    print(f"    stack:")
-    print_stack(data["stack"], "        ")
-    # print(f"    by:     {caller_method}")
-    # print(f"    at:     0x{addr:08x}")
+    print(f"    by:     {caller_method}")
+    print(f"    at:     0x{addr:08x}")
+    # print(f"    stack:")
+    # print_stack(data["stack"], "        ")
+    if addr < 0:
+        return
+    data_storage["invoke_data"].append(
+        {
+            "method": method,
+            "caller_method": caller_method,
+            "addr": addr,
+        }
+    )
 
 
-def handle_class_new_inst_data(data):
+def handle_class_new_inst_data(data, data_storage: dict):
     constructor = data["constructor"]
-    # caller_method = "?"  # get_method_id(data["caller_method"])
-    # addr = data["addr"]
+    if len(data["stack"]) == 0:
+        return
+    if (
+        data["stack"][0]["method"]
+        != "Ljava/lang/Class;->newInstance()Ljava/lang/Object;"
+    ):
+        frame = data["stack"][0]
+    elif len(data["stack"]) > 1:
+        frame = data["stack"][1]
+    else:
+        return
+    caller_method = frame["method"]
+    addr = frame["bytecode_index"]
     print("Class.NewInstance:")
     print(f"    called: {constructor}")
-    print(f"    stack:")
-    print_stack(data["stack"], "        ")
-    # print(f"    by:     {caller_method}")
-    # print(f"    at:     0x{addr:08x}")
+    print(f"    by:     {caller_method}")
+    print(f"    at:     0x{addr:08x}")
+    # print(f"    stack:")
+    # print_stack(data["stack"], "        ")
+    if addr < 0:
+        return
+    data_storage["class_new_inst_data"].append(
+        {
+            "constructor": constructor,
+            "caller_method": caller_method,
+            "addr": addr,
+        }
+    )
 
 
-def handle_cnstr_new_inst_data(data):
+def handle_cnstr_new_inst_data(data, data_storage: dict):
     constructor = data["constructor"]
     if not constructor.startswith("Lcom/example/theseus"):
         return
-    # caller_method = "?"  # get_method_id(data["caller_method"])
-    # addr = data["addr"]
+    if len(data["stack"]) == 0:
+        return
+    caller_method = data["stack"][0]["method"]
+    addr = data["stack"][0]["bytecode_index"]
     print("Constructor.newInstance:")
     print(f"    called: {constructor}")
-    print(f"    stack:")
-    print_stack(data["stack"], "        ")
-    # print(f"    by:     {caller_method}")
-    # print(f"    at:     0x{addr:08x}")
+    print(f"    by:     {caller_method}")
+    print(f"    at:     0x{addr:08x}")
+    # print(f"    stack:")
+    # print_stack(data["stack"], "        ")
+    if addr < 0:
+        return
+    data_storage["cnstr_new_inst_data"].append(
+        {
+            "constructor": constructor,
+            "caller_method": caller_method,
+            "addr": addr,
+        }
+    )
 
 
 def main():
@@ -104,6 +147,13 @@ def main():
         help="The android device to connect to, eg: 'emulator-5554'",
         type=str,
     )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="where to dump the collected data, default is stdout",
+        type=Path,
+    )
     args = parser.parse_args()
     env = dict(os.environ)
 
@@ -132,9 +182,15 @@ def main():
     session = device.attach(pid)
     script = session.create_script(script)
 
+    data_storage = {
+        "invoke_data": [],
+        "class_new_inst_data": [],
+        "cnstr_new_inst_data": [],
+    }
+
     script.on(
         "message",
-        on_message,
+        lambda msg, data: on_message(msg, data, data_storage),
     )
 
     # Load script
@@ -144,3 +200,8 @@ def main():
 
     print("Press ENTER to finish the analysis")
     input()
+    if args.output is None:
+        print(json.dumps(data_storage, indent="  "))
+    else:
+        with args.output.open("w") as fp:
+            json.dump(data_storage, fp)