From ea6ce5d7a10634baef275a15de89f68bc49957f3 Mon Sep 17 00:00:00 2001
From: Jean-Marie Mineau <jean-marie.mineau@inria.fr>
Date: Fri, 24 Jan 2025 14:14:40 +0100
Subject: [PATCH] read apk

---
 androscalpel/src/apk.rs       | 42 ++++++++++++++++++++++++++++++++++-
 apk_frauder/src/zip_reader.rs |  6 +++--
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/androscalpel/src/apk.rs b/androscalpel/src/apk.rs
index 3bc07c7..062eca0 100644
--- a/androscalpel/src/apk.rs
+++ b/androscalpel/src/apk.rs
@@ -3,6 +3,9 @@
 use anyhow::{anyhow, bail, Context};
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
+use std::fs::File;
+use std::io::Cursor;
+use std::path::PathBuf;
 
 use log::info;
 use pyo3::prelude::*;
@@ -14,6 +17,7 @@ use crate::Result;
 use crate::*;
 use androscalpel_serializer::Instruction as InsFormat;
 use androscalpel_serializer::*;
+use apk_frauder::ZipFileReader;
 use rayon::prelude::*;
 
 #[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)]
@@ -2981,6 +2985,7 @@ impl Apk {
 impl Apk {
     #[new]
     pub fn new() -> Self {
+        // TODO take argument and dispatch to load_apk or load_apk_bin
         Self {
             dex_files: HashMap::new(),
         }
@@ -2988,7 +2993,42 @@ impl Apk {
 
     /// Load all android files in an application.
     /// This **does not include any .dex file that android would not load.
-    pub fn load_apk() -> Self {}
+    #[staticmethod]
+    #[pyo3(signature = (apk, label_each_ins=false, cache=false))]
+    pub fn load_apk(apk: PathBuf, label_each_ins: bool, cache: bool) -> Result<Self> {
+        let file = File::open(apk)?;
+        let mut apk_z = ZipFileReader::new(file);
+        let mut apk = Self::default();
+        let dex_names = apk_z
+            .get_classes_file_info()
+            .into_iter()
+            .map(|info| info.get_name())
+            .collect::<Vec<_>>();
+        for name in dex_names {
+            let data = apk_z.read_file_as_vec(&name);
+            apk.add_dex_file(&name, &data, label_each_ins, cache)?;
+        }
+        Ok(apk)
+    }
+
+    /// Load all android files in an application.
+    /// This **does not include any .dex file that android would not load.
+    #[staticmethod]
+    #[pyo3(signature = (apk, label_each_ins=false, cache=false))]
+    pub fn load_apk_bin(apk: &[u8], label_each_ins: bool, cache: bool) -> Result<Self> {
+        let mut apk_z = ZipFileReader::new(Cursor::new(apk));
+        let mut apk = Self::default();
+        let dex_names = apk_z
+            .get_classes_file_info()
+            .into_iter()
+            .map(|info| info.get_name())
+            .collect::<Vec<_>>();
+        for name in dex_names {
+            let data = apk_z.read_file_as_vec(&name);
+            apk.add_dex_file(&name, &data, label_each_ins, cache)?;
+        }
+        Ok(apk)
+    }
 
     // TODO: change cache to true when cache invalidation is setup
     /// Add the content of a dex file to the apk.
diff --git a/apk_frauder/src/zip_reader.rs b/apk_frauder/src/zip_reader.rs
index 8a17cb2..590ef87 100644
--- a/apk_frauder/src/zip_reader.rs
+++ b/apk_frauder/src/zip_reader.rs
@@ -297,14 +297,16 @@ impl<T: Read + Seek> ZipFileReader<T> {
         data
     }
 
-    pub fn read_file_as_vec(&mut self, file: &FileInfo) -> Vec<u8> {
+    pub fn read_file_as_vec(&mut self, name: &str) -> Vec<u8> {
+        let file = self.get_file_info(name).unwrap();
         let offset = file.get_file_offset();
         let size_c = file.header.compressed_size as usize;
         let size = file.header.uncompressed_size as usize;
+        let compression_method = file.header.compression_method;
         let mut data = vec![0u8; size_c];
         self.data.seek(SeekFrom::Start(offset)).unwrap();
         self.data.read_exact(&mut data).unwrap();
-        match file.header.compression_method {
+        match compression_method {
             CompressionMethod::Stored => {}
             CompressionMethod::Deflated => {
                 let mut decomp_data = vec![0u8; size];