add method to facilitate apk extention

androscalpel/src/apk.rs

@@ -2957,11 +2957,13 @@ impl Apk {
         None
     }
 
-    /// Load all android files in an application.
-    /// This **does not include any .dex file that android would not load.
+    /// Load android .dex files from an application.
+    /// This **does not include** any .dex file that **android would not load**.
     ///
-    /// - `label_ins` is a function that take an method id and instruction and return
-    ///      true is a label "label_{addr:08X}" should be added befor the instruction.
+    /// - `label_ins` is a function that take an method id, instruction and address as argument and
+    ///    and return `Some(label_name)` if a label named label_name should be added before the
+    ///    instruction.
+    /// - `cache`: if set to true, copy and cache the binary data format.
     #[cfg(feature = "external-zip-reader")]
     pub fn load_apk<F>(apk: impl Read + Seek, label_ins: F, cache: bool) -> Result<Self>
     where
@@ -2997,11 +2999,12 @@ impl Apk {
         Ok(apk)
     }
 
-    /// Load all android files in an application.
-    /// This **does not include any .dex file that android would not load.
+    /// Load android .dex files from an application.
+    /// This **does not include** any .dex file that **android would not load**.
     ///
-    /// - `label_ins`: Function that take a method id, instruction and address and return
-    ///      a label, if a label needs to be inserted before the instruction.
+    /// - `label_ins` is a function that take an method id, instruction and address as argument and
+    ///    and return `Some(label_name)` if a label named label_name should be added before the
+    ///    instruction.
     /// - `cache`: if set to true, copy and cache the binary data format.
     #[cfg(not(feature = "external-zip-reader"))]
     pub fn load_apk<F>(apk: impl Read + Seek, label_ins: F, cache: bool) -> Result<Self>
@@ -3025,6 +3028,54 @@ impl Apk {
         Ok(apk)
     }
 
+    /// Load code from a .dex/.apk/.jar
+    ///
+    /// # Parameters
+    /// - `code`: the source for the code to add. Can be a .dex, .apk or .jar (but .oat are not
+    ///     supported)
+    /// - `label_ins`: Function that take a method id, instruction and address and return
+    ///      a label, if a label needs to be inserted before the instruction.
+    /// - `cache`: if set to true, copy and cache the binary data format.
+    ///
+    pub fn add_code<F>(
+        &mut self,
+        mut code: impl Read + Seek,
+        label_ins: F,
+        cache: bool,
+    ) -> Result<()>
+    where
+        F: FnMut(&IdMethod, &instructions::Instruction, usize) -> Option<String>
+            + Clone
+            + Send
+            + Sync,
+    {
+        let mut i = 1;
+        let mut dex_name: String = "classes.dex".into();
+        while self.dex_files.contains_key(&dex_name) {
+            i += 1;
+            dex_name = format!("classes{i}.dex");
+        }
+        if let Some(_version) = crate::utils::is_dex(&mut code)? {
+            let mut data = vec![];
+            std::io::copy(&mut code, &mut data)?;
+            self.add_dex_file(&dex_name, &data, label_ins, cache)
+        } else if crate::utils::is_zip(&mut code)? {
+            let mut tmp_apk = Apk::load_apk(code, label_ins, cache)?;
+            let mut j = 1;
+            let mut tmp_dex_name: String = "classes.dex".into();
+            while let Some(dex_file) = tmp_apk.dex_files.remove(&tmp_dex_name) {
+                self.dex_files.insert(dex_name, dex_file);
+                i += 1;
+                j += 1;
+                dex_name = format!("classes{i}.dex");
+                tmp_dex_name = format!("classes{j}.dex");
+            }
+            Ok(())
+        } else {
+            bail!("Could not recognize the type of the input file")
+        }
+    }
+
     /// Add the content of a dex file to the apk.
     ///
     /// # Parameters
@@ -3033,6 +3084,10 @@ impl Apk {
     /// - `label_ins`: Function that take a method id, instruction and address and return
     ///      a label, if a label needs to be inserted before the instruction.
     /// - `cache`: if set to true, copy and cache the binary data format.
+    ///
+    /// # Infos
+    ///
+    /// `data` is as `&[u8]` to allow concurrent access to the file.
     pub fn add_dex_file<F>(
         &mut self,
         name: &str,

androscalpel/src/lib.rs

@@ -19,6 +19,7 @@ pub mod method_handle;
 #[cfg(feature = "python")]
 pub mod py_utils;
 pub mod scalar;
+pub mod utils;
 pub mod value;
 pub mod visitor;
 
@@ -39,6 +40,7 @@ pub use instructions::*;
 pub use method::*;
 pub use method_handle::*;
 pub use scalar::*;
+pub use utils::*;
 pub use value::*;
 pub use visitor::*;
 

androscalpel/src/py_utils.rs

@@ -63,37 +63,16 @@ pub fn list_defined_classes(dex: &[u8]) -> Result<HashSet<IdType>> {
 
 /// Test if a file is as .dex file an return the dex version if it is, else return None.
 #[pyfunction]
-pub fn is_dex(file: PathBuf) -> Option<usize> {
-    let mut file = match File::open(file) {
-        Ok(file) => file,
-        Err(_) => return None,
-    };
-    HeaderItem::deserialize(&mut file)
-        .ok()
-        .and_then(|header| String::from_utf8(header.magic.version.to_vec()).ok())
-        .and_then(|version| version.parse::<usize>().ok())
+pub fn is_dex(file: PathBuf) -> Result<Option<usize>> {
+    let mut file = File::open(file)?;
+    crate::utils::is_dex(file)
 }
 
 /// Test if a file is a zip file.
 #[pyfunction]
-pub fn is_zip(file: PathBuf) -> bool {
-    let mut file = match File::open(file) {
-        Ok(file) => file,
-        Err(_) => return false,
-    };
-    let ecd_off = if let Some(off) = ZipFileReader::get_end_of_central_directory_offset(&mut file) {
-        off
-    } else {
-        return false;
-    };
-    if file.seek(SeekFrom::Start(ecd_off)).is_err() {
-        return false;
-    }
-    if let Ok(sig) = apk_frauder::Signature::deserialize(&mut file) {
-        EndCentralDirectory::SIGNATURE == sig
-    } else {
-        false
-    }
+pub fn is_zip(file: PathBuf) -> Result<bool> {
+    let mut file = File::open(file)?;
+    crate::utils::is_zip(file)
 }
 
 /// Replace the dex files a an apk an resigned the apk.
@@ -103,7 +82,7 @@ pub fn is_zip(file: PathBuf) -> bool {
 /// For now, only jks keystore are allowed.
 ///
 /// The `zipalign` and `apksigner` args allow to use a specific version of the
-/// tools instead of the one in the PATH (if it even exist)
+/// toimpl Read + Seekols instead of the one in the PATH (if it even exist)
 ///
 /// `additionnal_files` is a dict of file to add, modify or remove in the apk.
 /// The keys are the file names and the values are `None` to remove the file, or

androscalpel/src/utils.rs

@@ -0,0 +1,34 @@
+use std::io::{Read, Seek, SeekFrom};
+
+use crate::Result;
+use androscalpel_serializer::{HeaderItem, Serializable};
+use apk_frauder::{end_of_central_directory::EndCentralDirectory, ZipFileReader};
+
+/// Test if a file is as .dex file an return the dex version if it is, else return None.
+pub fn is_dex(file: &mut (impl Read + Seek)) -> Result<Option<usize>> {
+    let pos = file.stream_position()?;
+    let r = HeaderItem::deserialize(file)
+        .ok()
+        .and_then(|header| String::from_utf8(header.magic.version.to_vec()).ok())
+        .and_then(|version| version.parse::<usize>().ok());
+    file.seek(SeekFrom::Start(pos))?;
+    Ok(r)
+}
+
+/// Test if a file is a zip file.
+pub fn is_zip(mut file: impl Read + Seek) -> Result<bool> {
+    let pos = file.stream_position()?;
+    let ecd_off = if let Some(off) = ZipFileReader::get_end_of_central_directory_offset(&mut file) {
+        off
+    } else {
+        return Ok(false);
+    };
+    file.seek(SeekFrom::Start(ecd_off))?;
+    let r = if let Ok(sig) = apk_frauder::Signature::deserialize(&mut file) {
+        EndCentralDirectory::SIGNATURE == sig
+    } else {
+        false
+    };
+    file.seek(SeekFrom::Start(pos))?;
+    Ok(r)
+}

androscalpel_serializer/src/file_reader.rs

@@ -11,6 +11,8 @@ use std::sync::atomic::{AtomicBool, Ordering};
 
 #[derive(Debug)]
 pub struct DexFileReader<'a> {
+    // Ideally, this would be a Read+Seek, but Read+Seek is not thread safe, while we can
+    // internally instanciate multiple cursors on the same non mutable slice.
     data: &'a [u8],
     header: HeaderItem,
     string_ids: Vec<StringIdItem>,