android_of_theseus/patcher/src/code_loading_patcher.rs

use std::collections::HashMap;
use std::fs::File;

use androscalpel::{Apk, DexString, IdType, VisitorMut};
use anyhow::{Context, Result};
use clap::ValueEnum;

use crate::runtime_data::RuntimeData;

#[derive(ValueEnum, Debug, PartialEq, Clone, Copy, Default)]
pub enum CodePatchingStrategy {
    #[default]
    Naive,
    ModelClassLoaders,
}

pub fn insert_code(
    strategy: CodePatchingStrategy,
    apk: &mut Apk,
    data: &RuntimeData,
) -> Result<()> {
    match strategy {
        CodePatchingStrategy::Naive => insert_code_naive(apk, data),
        CodePatchingStrategy::ModelClassLoaders => insert_code_model_class_loaders(apk, data),
    }
}

/// Insert statically bytecode that was loaded from other source at runtime.
/// For now, we ignore class collision.
fn insert_code_naive(apk: &mut Apk, data: &RuntimeData) -> Result<()> {
    for dyn_data in &data.dyn_code_load {
        for file in &dyn_data.files {
            let file = File::open(file)?;
            apk.add_code(file, crate::labeling, false)?;
        }
    }
    Ok(())
}

fn insert_code_model_class_loaders(apk: &mut Apk, data: &RuntimeData) -> Result<()> {
    let mut class_defined = apk.list_classes();
    let mut class_loaders = HashMap::new();
    class_loaders.insert(
        "MAIN".to_string(),
        ClassLoader {
            id: "MAIN".to_string(),
            parent: None,
            class: IdType::from_smali("Ljava/lang/Boolean;").unwrap(),
            apk: ApkOrRef::Ref(apk),
            renamed_classes: HashMap::new(),
        },
    );
    for dyn_data in &data.dyn_code_load {
        let mut apk = Apk::new();
        let class = dyn_data.classloader_class.clone();
        for file in &dyn_data.files {
            let file = File::open(file)?;
            apk.add_code(file, crate::labeling, false)?;
        }

        assert!(!class_loaders.contains_key(&dyn_data.classloader));

        let classes = apk.list_classes();
        let mut class_loader = ClassLoader {
            id: dyn_data.classloader.clone(),
            parent: None,
            class,
            apk: ApkOrRef::Owned(apk),
            renamed_classes: HashMap::new(),
        };
        let collisions = class_defined.intersection(&classes);
        for cls in collisions {
            class_loader.rename_classdef(cls)?;
        }
        class_defined.extend(classes);

        class_loaders.insert(dyn_data.classloader.clone(), class_loader);
    }
    // TODO: get the ClassLoader::parent values...
    // TODO: model the delegation behavior and rename ref to class accordingly
    // TODO: update Runtime Data to reflect the name change

    let apk = if let ApkOrRef::Ref(apk) = class_loaders.remove("MAIN").unwrap().apk {
        apk
    } else {
        panic!("Main APK is not stored as ref?")
    };
    for (_, ClassLoader { apk: other, .. }) in class_loaders.into_iter() {
        if let ApkOrRef::Owned(other) = other {
            apk.merge(other);
        } else {
            panic!("Secondary APK is not stored as owned?")
        }
    }
    //todo!()
    Ok(())
}

/// Structure modelizing a class loader.
#[derive(Debug, PartialEq)]
struct ClassLoader<'a> {
    pub id: String,
    pub parent: Option<String>,
    pub class: IdType,
    pub apk: ApkOrRef<'a>,
    pub renamed_classes: HashMap<IdType, IdType>,
}

impl ClassLoader<'_> {
    pub fn apk(&mut self) -> &mut Apk {
        match &mut self.apk {
            ApkOrRef::Owned(ref mut apk) => apk,
            ApkOrRef::Ref(ref mut apk) => apk,
        }
    }

    pub fn rename_classdef(&mut self, cls: &IdType) -> Result<()> {
        use androscalpel::SmaliName;
        let id = self.id.clone();
        let mut i = 0;
        let name = if let Some(name) = cls.get_class_name() {
            name
        } else {
            log::warn!("Tried to rename non class type {}", cls.__str__());
            return Ok(());
        };
        let new_name = loop {
            let prefix: DexString = if i == 0 {
                format!("theseus-dedup/{}/", self.id).into()
            } else {
                format!("theseus-dedup/{}-{i}/", self.id).into()
            };
            let new_name = IdType::class_from_dex_string(&prefix.concatenate(&name));
            if self.apk().get_class(&new_name).is_none() {
                break new_name;
            }
            i += 1;
        };
        println!(
            "TODO: rename {} -> {}",
            cls.try_to_smali().unwrap(),
            new_name.try_to_smali().unwrap(),
        );

        let class = self.apk().remove_class(cls, None)?.with_context(|| {
            format!(
                "Try to rename classdef of {} in class loader {}, but classdef not found",
                cls.__str__(),
                &id
            )
        })?;
        let class = RenameTypeVisitor {
            new_names: [(cls.clone(), new_name.clone())].into(),
        }
        .visit_class(class)?;
        self.apk().add_class("classes.dex", class)?;

        self.renamed_classes.insert(cls.clone(), new_name);
        Ok(())
    }
}

struct RenameTypeVisitor {
    pub new_names: HashMap<IdType, IdType>,
}

impl VisitorMut for RenameTypeVisitor {
    fn visit_type(&mut self, id: IdType) -> Result<IdType> {
        match self.new_names.get(&id) {
            Some(newid) => Ok(newid.clone()),
            None => Ok(id.clone()),
        }
    }
}

#[derive(Debug, PartialEq)]
enum ApkOrRef<'a> {
    Owned(Apk),
    Ref(&'a mut Apk),
}