From df9d258780d39b7c6c70898207d08004aa6bfaa6 Mon Sep 17 00:00:00 2001
From: Jean-Marie Mineau <jean-marie.mineau@inria.fr>
Date: Mon, 6 May 2024 18:01:27 +0200
Subject: [PATCH] add file modification

---
 androscalpel/src/py_utils.rs  | 23 +++++++++++++++++++++--
 apk_frauder/src/lib.rs        | 18 ++++++++++++++++++
 apk_frauder/src/main.rs       |  3 +++
 apk_frauder/src/zip_reader.rs |  2 ++
 4 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/androscalpel/src/py_utils.rs b/androscalpel/src/py_utils.rs
index d32cb9c..05e17ec 100644
--- a/androscalpel/src/py_utils.rs
+++ b/androscalpel/src/py_utils.rs
@@ -3,7 +3,7 @@
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fs::File;
 use std::io::{Cursor, Seek, SeekFrom};
 use std::path::PathBuf;
@@ -104,6 +104,10 @@ pub fn is_zip(file: PathBuf) -> bool {
 ///
 /// 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)
+///
+/// `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
+/// `bytes` for the content of the file.
 #[pyfunction]
 pub fn replace_dex(
     apk: PathBuf,
@@ -112,9 +116,24 @@ pub fn replace_dex(
     keystore: PathBuf,
     zipalign: Option<PathBuf>,
     apksigner: Option<PathBuf>,
+    additionnal_files: Option<HashMap<String, Option<&[u8]>>>,
 ) {
     let mut dexfiles: Vec<_> = dexfiles.iter().map(Cursor::new).collect();
-    apk_frauder::replace_dex(apk, dst, &mut dexfiles, keystore, zipalign, apksigner)
+    let additionnal_files: Option<HashMap<_, _>> = additionnal_files.map(|additionnal_files| {
+        additionnal_files
+            .into_iter()
+            .map(|(k, v)| (k, v.map(Cursor::new)))
+            .collect()
+    });
+    apk_frauder::replace_dex(
+        apk,
+        dst,
+        &mut dexfiles,
+        keystore,
+        zipalign,
+        apksigner,
+        additionnal_files,
+    )
 }
 
 /// export the function in a python module
diff --git a/apk_frauder/src/lib.rs b/apk_frauder/src/lib.rs
index f23fedb..3d35ae2 100644
--- a/apk_frauder/src/lib.rs
+++ b/apk_frauder/src/lib.rs
@@ -1,4 +1,5 @@
 use androscalpel_serializer::Serializable;
+use std::collections::HashMap;
 use std::env;
 use std::fs;
 use std::fs::File;
@@ -156,6 +157,7 @@ pub fn replace_dex(
     // 2048 -validity 10000 -alias ALIAS`
     zipalign: Option<impl AsRef<Path>>,
     apksigner: Option<impl AsRef<Path>>,
+    additionnal_files: Option<HashMap<String, Option<impl Read + Seek>>>,
 ) {
     let zipalign = if let Some(path) = &zipalign {
         path.as_ref().as_os_str()
@@ -185,6 +187,10 @@ pub fn replace_dex(
 
     apk.unlink_signature_files();
     apk.unlink_bytecode_files();
+    if let Some(additionnal_files) = &additionnal_files {
+        apk.files
+            .retain(|file| additionnal_files.get(&file.get_name()).is_none());
+    }
     for f in apk.files.clone() {
         apk_out.insert_file_from_zip(f, &mut apk);
     }
@@ -200,6 +206,18 @@ pub fn replace_dex(
             Some(file_info_ref.local_header.clone()),
         );
     }
+    if let Some(mut additionnal_files) = additionnal_files {
+        for (name, data) in &mut additionnal_files {
+            file_info_ref.set_name(name);
+            if let Some(data) = data.as_mut() {
+                apk_out.insert_file(
+                    data,
+                    file_info_ref.header.clone(),
+                    Some(file_info_ref.local_header.clone()),
+                );
+            }
+        }
+    }
     apk_out.write_central_directory();
     // TODO: we can probably do that ourself an spare ourself the trouble of finding zipalign
     Command::new(zipalign)
diff --git a/apk_frauder/src/main.rs b/apk_frauder/src/main.rs
index 0418064..77c58a0 100644
--- a/apk_frauder/src/main.rs
+++ b/apk_frauder/src/main.rs
@@ -1,5 +1,7 @@
+use std::collections::HashMap;
 use std::env;
 use std::fs::File;
+use std::io::Cursor;
 
 fn main() {
     apk_frauder::replace_dex(
@@ -15,5 +17,6 @@ fn main() {
             "{}/Android/Sdk/build-tools/34.0.0/apksigner",
             env::var("HOME").expect("$HOME not set")
         )),
+        None::<HashMap<String, Option<Cursor<&[u8]>>>>,
     );
 }
diff --git a/apk_frauder/src/zip_reader.rs b/apk_frauder/src/zip_reader.rs
index 584530b..440cb14 100644
--- a/apk_frauder/src/zip_reader.rs
+++ b/apk_frauder/src/zip_reader.rs
@@ -297,6 +297,8 @@ impl<T: Read + Seek> ZipFileReader<T> {
     }
 
     pub fn get_classes_file_info(&self) -> Vec<&FileInfo> {
+        // TODO: check for holes: if classesN.dex does not exist,
+        // classesM.dex for M>N are ignored.
         self.files
             .iter()
             .filter(|&file| match_dexfile_name(&file.get_name()))