make apk and dexfile visitable

androscalpel/src/apk.rs

@@ -31,6 +31,40 @@ pub struct DexFile {
     pub(crate) bin_cache: Option<Vec<u8>>, // TODO: invalidate the cache !!!
 }
 
+impl<V: Visitor> Visitable<V> for DexFile {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        for (id, class) in &self.classes {
+            v.visit_type(id)?;
+            v.visit_class(class)?;
+        }
+        for string in &self.not_referenced_strings {
+            v.visit_string(string)?;
+        }
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for DexFile {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            classes: {
+                let mut classes = HashMap::new();
+                for (id, class) in self.classes.into_iter() {
+                    classes.insert(v.visit_type(id)?, v.visit_class(class)?);
+                }
+                classes
+            },
+            not_referenced_strings: {
+                let mut classes = HashSet::new();
+                for string in self.not_referenced_strings.into_iter() {
+                    classes.insert(v.visit_string(string)?);
+                }
+                classes
+            },
+            bin_cache: None,
+        })
+    }
+}
+
 /// Represent an apk.
 #[cfg_attr(feature = "python", pyclass(eq))]
 #[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)]
@@ -38,6 +72,28 @@ pub struct Apk {
     pub dex_files: HashMap<String, DexFile>, // TODO: use accessort for chache invalidation
 }
 
+impl<V: Visitor> Visitable<V> for Apk {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        for file in self.dex_files.values() {
+            v.visit_dex_file(file)?;
+        }
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for Apk {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            dex_files: {
+                let mut dex_files = HashMap::new();
+                for (name, file) in self.dex_files.into_iter() {
+                    dex_files.insert(name, v.visit_dex_file(file)?);
+                }
+                dex_files
+            },
+        })
+    }
+}
+
 impl Apk {
     /// Extract a class from a dex file reader.
     /// `class_item_idx` if the index of the `class_def_item` of the class, **not** the

androscalpel/src/lib.rs

@@ -1,4 +1,4 @@
-#![allow(clippy::unnecessary_fallible_conversions)]
+//#![allow(clippy::unnecessary_fallible_conversions)]
 // DexString has Into<String> but it's only for
 // python, TryInto should be prefered
 use anyhow::Result;

androscalpel/src/visitor.rs

@@ -1,8 +1,8 @@
 //! The visitor trait and common implementations.
 
 use crate::{
-    ins::Instruction, scalar::*, CallSite, Class, Code, DexAnnotation, DexAnnotationItem,
+    ins::Instruction, scalar::*, Apk, CallSite, Class, Code, DexAnnotation, DexAnnotationItem,
-    DexString, DexValue, Field, FieldVisibility, HiddenApiData, HiddenApiDomain,
+    DexFile, DexString, DexValue, Field, FieldVisibility, HiddenApiData, HiddenApiDomain,
     HiddenApiPermission, IdEnum, IdField, IdMethod, IdMethodType, IdType, Method, MethodHandle,
     MethodVisibility, Result,
 };
@@ -102,6 +102,12 @@ pub trait Visitor: Sized {
     fn visit_code(&mut self, code: &Code) -> Result<()> {
         code.default_visit(self)
     }
+    fn visit_dex_file(&mut self, dex_file: &DexFile) -> Result<()> {
+        dex_file.default_visit(self)
+    }
+    fn visit_apk(&mut self, apk: &Apk) -> Result<()> {
+        apk.default_visit(self)
+    }
 }
 
 pub trait VisitorMut: Sized {
@@ -207,6 +213,12 @@ pub trait VisitorMut: Sized {
     fn visit_code(&mut self, code: Code) -> Result<Code> {
         code.default_visit_mut(self)
     }
+    fn visit_dex_file(&mut self, dex_file: DexFile) -> Result<DexFile> {
+        dex_file.default_visit_mut(self)
+    }
+    fn visit_apk(&mut self, apk: Apk) -> Result<Apk> {
+        apk.default_visit_mut(self)
+    }
 }
 
 /// Trait for structures that can be visited.