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"])
+    if len(data["stack"]) == 0:
-    # addr = data["addr"]
+        return
+    caller_method = data["stack"][0]["method"]
+    addr = data["stack"][0]["bytecode_index"]
     print("Method.Invoke:")
     print(f"    called: {method}")
-    print(f"    stack:")
+    print(f"    by:     {caller_method}")
-    print_stack(data["stack"], "        ")
+    print(f"    at:     0x{addr:08x}")
-    # print(f"    by:     {caller_method}")
+    # print(f"    stack:")
-    # print(f"    at:     0x{addr:08x}")
+    # 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"])
+    if len(data["stack"]) == 0:
-    # addr = data["addr"]
+        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(f"    by:     {caller_method}")
-    print_stack(data["stack"], "        ")
+    print(f"    at:     0x{addr:08x}")
-    # print(f"    by:     {caller_method}")
+    # print(f"    stack:")
-    # print(f"    at:     0x{addr:08x}")
+    # 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"])
+    if len(data["stack"]) == 0:
-    # addr = data["addr"]
+        return
+    caller_method = data["stack"][0]["method"]
+    addr = data["stack"][0]["bytecode_index"]
     print("Constructor.newInstance:")
     print(f"    called: {constructor}")
-    print(f"    stack:")
+    print(f"    by:     {caller_method}")
-    print_stack(data["stack"], "        ")
+    print(f"    at:     0x{addr:08x}")
-    # print(f"    by:     {caller_method}")
+    # print(f"    stack:")
-    # print(f"    at:     0x{addr:08x}")
+    # 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)

patcher/Cargo.lock

@@ -815,6 +815,8 @@ dependencies = [
  "clap",
  "env_logger",
  "reqwest",
+ "serde",
+ "serde_json",
 ]
 
 [[package]]

patcher/Cargo.toml

@@ -12,3 +12,5 @@ anyhow = "1.0.95"
 clap = { version = "4.5.27", features = ["derive"] }
 env_logger = "0.11.6"
 reqwest = { version = "0.12.12", default-features = false, features = ["blocking", "rustls-tls"] }
+serde = "1.0.217"
+serde_json = "1.0.138"

patcher/src/get_apk.rs

@@ -1,97 +0,0 @@
-use androscalpel::Apk;
-use clap::Args;
-use std::fs::{read_to_string, File};
-use std::io::Cursor;
-use std::path::PathBuf;
-use std::time::Duration;
-
-use crate::labeling;
-
-#[derive(Clone, Args, Debug)]
-pub struct ApkLocation {
-    #[arg(short, long, conflicts_with = "sha256")]
-    pub path: Option<PathBuf>,
-    #[arg(long, conflicts_with = "path", requires = "androzoo_key")]
-    pub sha256: Option<String>,
-    #[command(flatten)]
-    pub androzoo_key: Option<AndrozooKey>,
-}
-
-impl ApkLocation {
-    pub fn get_id(&self) -> String {
-        match self {
-            ApkLocation {
-                path: Some(path), ..
-            } => path.as_path().file_name().unwrap().to_str().unwrap().into(),
-            ApkLocation {
-                sha256: Some(sha256),
-                ..
-            } => sha256.clone(),
-            _ => panic!("Invalid ApkLocation"),
-        }
-    }
-}
-
-#[derive(Clone, Args, Debug)]
-pub struct AndrozooKey {
-    #[arg(long, group = "androzoo_key", conflicts_with = "api_key")]
-    api_key_path: Option<PathBuf>,
-    #[arg(long, group = "androzoo_key", conflicts_with = "api_key_path")]
-    api_key: Option<String>,
-}
-
-impl AndrozooKey {
-    fn get_key(&self) -> String {
-        match self {
-            AndrozooKey {
-                api_key_path: Some(path),
-                ..
-            } => read_to_string(path)
-                .expect("Failed to read key from file")
-                .trim()
-                .to_string(),
-            AndrozooKey {
-                api_key: Some(key), ..
-            } => key.trim().to_string(),
-            _ => panic!("No key here"),
-        }
-    }
-}
-
-pub fn get_apk(location: &ApkLocation) -> Apk {
-    match location {
-        ApkLocation {
-            androzoo_key: Some(key),
-            sha256: Some(sha256),
-            ..
-        } => {
-            let key = key.get_key();
-            let url = reqwest::Url::parse_with_params(
-                "https://androzoo.uni.lu/api/download",
-                &[("apikey", key), ("sha256", sha256.clone())],
-            )
-            .expect("Failed to parse url");
-            let res = reqwest::blocking::Client::builder()
-                .timeout(Duration::from_secs(300))
-                .build()
-                .expect("Failed to build client")
-                .get(url)
-                .send()
-                .expect("Failed to download apk");
-            match res.status() {
-                reqwest::StatusCode::OK => (),
-                s => panic!("Failed to download apk: {:?}", s),
-            }
-            Apk::load_apk(
-                &mut Cursor::new(res.bytes().expect("Failed to get APK bytes")),
-                labeling,
-                false,
-            )
-            .unwrap()
-        }
-        ApkLocation {
-            path: Some(path), ..
-        } => Apk::load_apk(File::open(path).unwrap(), labeling, false).unwrap(),
-        _ => panic!("Don't know what to do with:\n{:#?}", location),
-    }
-}

patcher/src/lib.rs

@@ -4,14 +4,14 @@ use anyhow::{bail, Context, Result};
 use std::collections::{HashMap, HashSet};
 use std::sync::LazyLock;
 
-pub mod get_apk;
+use serde::{Deserialize, Serialize};
 
 // TODO:
 // Check what
 // https://cs.android.com/android/platform/superproject/main/+/main:art/runtime/reflection.cc;drc=83db0626fad8c6e0508754fffcbbd58e539d14a5;l=698
 // does.
 
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct ReflectionData {
     pub invoke_data: Vec<ReflectionInvokeData>,
     pub class_new_inst_data: Vec<ReflectionClassNewInstData>,
@@ -95,7 +95,7 @@ impl ReflectionData {
 
 /// Structure storing the runtime information of a reflection call using
 /// `java.lang.reflect.Method.invoke()`.
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct ReflectionInvokeData {
     /// The method called by `java.lang.reflect.Method.invoke()`
     pub method: IdMethod,
@@ -109,7 +109,7 @@ pub struct ReflectionInvokeData {
 
 /// Structure storing the runtime information of a reflection instanciation using
 /// `java.lang.Class.newInstance()`.
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct ReflectionClassNewInstData {
     /// The constructor called by `java.lang.Class.newInstance()`
     pub constructor: IdMethod,
@@ -121,7 +121,7 @@ pub struct ReflectionClassNewInstData {
 
 /// Structure storing the runtime information of a reflection instanciation using
 /// `java.lang.reflect.Constructor.newInstance()`.
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct ReflectionCnstrNewInstData {
     /// The constructor calleb by `java.lang.reflect.Constructor.newInstance()`
     pub constructor: IdMethod,
@@ -191,7 +191,7 @@ static CNSTR_GET_DEC_CLS: LazyLock<IdMethod> = LazyLock::new(|| {
 });
 
 /// Function passed to [`androscalpel::Apk::load_apk`] to label the instructions of interest.
-fn labeling(_mth: &IdMethod, ins: &Instruction, addr: usize) -> Option<String> {
+pub fn labeling(_mth: &IdMethod, ins: &Instruction, addr: usize) -> Option<String> {
     match ins {
         Instruction::InvokeVirtual { method, .. }
             if method == &*MTH_INVOKE

patcher/src/main.rs

@@ -1,13 +1,14 @@
 use std::collections::HashMap;
-use std::io::Cursor;
+use std::fs::File;
+use std::io::{Cursor, Read};
 use std::path::PathBuf;
 
-use androscalpel::IdMethod;
+use androscalpel::Apk;
 
-use patcher::get_apk::{get_apk, ApkLocation};
 use patcher::{
-    transform_method, ReflectionClassNewInstData, ReflectionCnstrNewInstData, ReflectionData,
+    labeling,
-    ReflectionInvokeData,
+    transform_method,
+    ReflectionData, // ReflectionInvokeData, ReflectionClassNewInstData, ReflectionCnstrNewInstData,
 };
 
 use clap::Parser;
@@ -15,8 +16,6 @@ use clap::Parser;
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None, arg_required_else_help = true)]
 struct Cli {
-    #[clap(flatten)]
-    apk: ApkLocation,
     #[arg(short, long)]
     out: PathBuf,
     #[arg(short, long)]
@@ -25,13 +24,24 @@ struct Cli {
     zipalign: Option<PathBuf>,
     #[arg(short, long)]
     apksigner: Option<PathBuf>,
+    #[arg(short, long)]
+    path: PathBuf,
+    #[arg(short, long)]
+    reflection_data: PathBuf,
 }
 
 fn main() {
     env_logger::init();
     let cli = Cli::parse();
-    let mut apk = get_apk(&cli.apk);
+    let mut apk = Apk::load_apk(File::open(&cli.path).unwrap(), labeling, false).unwrap();
     //println!("{:#?}", apk.list_classes());
+    let mut json = String::new();
+    File::open(&cli.reflection_data)
+        .unwrap()
+        .read_to_string(&mut json)
+        .unwrap();
+    let reflection_data: ReflectionData = serde_json::from_str(&json).unwrap();
+    /*
     let reflection_data = ReflectionData {
         invoke_data: vec![
             ReflectionInvokeData {
@@ -102,12 +112,15 @@ fn main() {
             addr: 0x22,
         }],
     };
+    println!("{}", serde_json::to_string(&reflection_data).unwrap());
+    */
     for method in reflection_data.get_method_referenced().iter() {
-        let class = apk.get_class_mut(&method.class_).unwrap();
+        if let Some(class) = apk.get_class_mut(&method.class_) {
-        //println!("{:#?}", class.direct_methods.keys());
+            //println!("{:#?}", class.direct_methods.keys());
-        //println!("{:#?}", class.virtual_methods.keys());
+            //println!("{:#?}", class.virtual_methods.keys());
-        let method = class.virtual_methods.get_mut(method).unwrap();
+            let method = class.virtual_methods.get_mut(method).unwrap();
-        transform_method(method, &reflection_data).unwrap();
+            transform_method(method, &reflection_data).unwrap();
+        }
     }
     let mut dex_files = vec![];
     let mut files = apk.gen_raw_dex().unwrap();
@@ -127,7 +140,7 @@ fn main() {
     }
     // TODO: aapt would be a lot more stable
     apk_frauder::replace_dex(
-        cli.apk.path.unwrap(),
+        cli.path,
         cli.out,
         &mut dex_files,
         cli.keystore,