From 0bd5617f38080c64a896042e052e9416a93226dd Mon Sep 17 00:00:00 2001
From: Jean-Marie Mineau <jean-marie.mineau@inria.fr>
Date: Wed, 15 Jan 2025 11:48:18 +0100
Subject: [PATCH] handle automatic file selection for classes

---
 androscalpel/src/apk.rs | 72 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 66 insertions(+), 6 deletions(-)

diff --git a/androscalpel/src/apk.rs b/androscalpel/src/apk.rs
index e8d2aa7..0d2b9a5 100644
--- a/androscalpel/src/apk.rs
+++ b/androscalpel/src/apk.rs
@@ -2931,13 +2931,72 @@ impl Apk {
         Ok(bin_dex_files)
     }
 
+    // TODO: check for android platform classes?
     /// Search for the given class. If several classes with the same name are found,
-    /// return the class used by android (classes.dex over classes2.dex over classes3.dex ...),
+    /// return the class used by android (classes.dex over classes2.dex over classes3.dex until
+    /// classesN.dex does not exists),
     /// and if not found in files used by android, look in any other files (classes0.dex,
     /// classses1.dex, classes02.dex, assets/stuff.dex, ...), and select the first one in
     /// alphabetical order of dex file name.
-    pub fn get_class_mut(&mut self, class_id: IdType) -> Option<&mut Class> {
-        todo!()
+    pub fn get_class_mut(&mut self, class_id: &IdType) -> Option<&mut Class> {
+        self.get_file_of_class(class_id, false).and_then(|file| {
+            self.dex_files
+                .get_mut(&file)
+                .expect("`get_file_of_class()` return a string not in `dex_files`")
+                .classes
+                .get_mut(class_id)
+        })
+    }
+
+    /// Search for the dex file of a given class. If several classes with the same name exist,
+    /// return the file used by android (classes.dex over classes2.dex over classes3.dex until
+    /// classesN.dex does not exists), and if `valid_file_only` is set to false, look in any
+    /// other files (classes0.dex, classses1.dex, classes02.dex, assets/stuff.dex, ...), and
+    /// select the first one with the class in alphabetical order of dex file name.
+    fn get_file_of_class(&self, class_id: &IdType, valid_file_only: bool) -> Option<String> {
+        let mut files: HashSet<&String> = self.dex_files.keys().collect();
+        let mut i = 0;
+        let mut name: String = "classes.dex".into();
+        while files.contains(&name) {
+            if self
+                .dex_files
+                .get(&name)
+                .expect(
+                    "Something went wrong: `name` is from \
+                `dex_files.keys()` but `dex_files.get(&name)` \
+                failled",
+                )
+                .classes
+                .contains_key(class_id)
+            {
+                return Some(name);
+            }
+            files.remove(&name);
+            i += 1;
+            name = format!("classes{}.dex", i + 1);
+        }
+        if valid_file_only {
+            return None;
+        }
+        let mut files = files.into_iter().collect::<Vec<&String>>();
+        files.sort();
+        for name in files.into_iter() {
+            if self
+                .dex_files
+                .get(name)
+                .expect(
+                    "Something went wrong: `name` is from \
+                    `dex_files.keys()` but `dex_files.get(&name)` \
+                     failled",
+                )
+                .classes
+                .contains_key(class_id)
+            {
+                return Some(name.to_string());
+            }
+        }
+
+        None
     }
 }
 
@@ -3012,7 +3071,7 @@ impl Apk {
                     )
                 })?
         } else {
-            self.get_class_mut(ty.clone())
+            self.get_class_mut(ty)
                 .with_context(|| format!("{} not found in apk", method_id.class_.__repr__()))?
         };
         let method = class
@@ -3047,7 +3106,6 @@ impl Apk {
     }
 
     pub fn remove_class(&mut self, class: &IdType, dex_file: Option<&str>) -> Result<()> {
-        // TODO: remove all if dex_file is not provided
         if let Some(dex_file) = dex_file {
             self.dex_files
                 .get_mut(&dex_file.to_string())
@@ -3055,7 +3113,9 @@ impl Apk {
                 .classes
                 .remove(class);
         } else {
-            todo!()
+            for DexFile { classes, .. } in self.dex_files.values_mut() {
+                classes.remove(class);
+            }
         }
         Ok(())
     }