diff --git a/androscalpel/src/annotation.rs b/androscalpel/src/annotation.rs index 42400cf..1c4412b 100644 --- a/androscalpel/src/annotation.rs +++ b/androscalpel/src/annotation.rs @@ -111,6 +111,20 @@ impl DexAnnotationItem { } } +impl Visitable for DexAnnotationItem { + fn default_visit(&self, v: &mut V) -> Result<()> { + v.visit_annotation(&self.annotation) + } +} +impl VisitableMut for DexAnnotationItem { + fn default_visit_mut(self, v: &mut V) -> Result { + Ok(Self { + annotation: v.visit_annotation(self.annotation)?, + ..self + }) + } +} + /// An annotation. #[pyclass] #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] diff --git a/androscalpel/src/class.rs b/androscalpel/src/class.rs index c54c787..8df4661 100644 --- a/androscalpel/src/class.rs +++ b/androscalpel/src/class.rs @@ -7,7 +7,7 @@ use pyo3::prelude::*; use crate::{ DexAnnotationItem, DexString, Field, IdField, IdMethod, IdMethodType, IdType, Method, - MethodHandle, Result, + MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut, }; use androscalpel_serializer::consts::*; @@ -387,3 +387,92 @@ impl Class { self == other } } + +impl Visitable for Class { + fn default_visit(&self, v: &mut V) -> Result<()> { + v.visit_type(&self.descriptor)?; + if let Some(superclass) = &self.superclass { + v.visit_type(&superclass)?; + } + for interface in &self.interfaces { + v.visit_type(&interface)?; + } + if let Some(source_file) = &self.source_file { + v.visit_string(&source_file)?; + } + for (id, field) in &self.static_fields { + v.visit_field_id(&id)?; + v.visit_field(&field)?; + } + for (id, field) in &self.instance_fields { + v.visit_field_id(&id)?; + v.visit_field(&field)?; + } + for (id, meth) in &self.direct_methods { + v.visit_method_id(&id)?; + v.visit_method(&meth)?; + } + for (id, meth) in &self.virtual_methods { + v.visit_method_id(&id)?; + v.visit_method(&meth)?; + } + for annot in &self.annotations { + v.visit_annotation_item(&annot)?; + } + Ok(()) + } +} +impl VisitableMut for Class { + fn default_visit_mut(self, v: &mut V) -> Result { + Ok(Self { + descriptor: v.visit_type(self.descriptor)?, + superclass: self + .superclass + .map(|superclass| v.visit_type(superclass)) + .transpose()?, + interfaces: self + .interfaces + .into_iter() + .map(|interface| v.visit_type(interface)) + .collect::>()?, + source_file: self + .source_file + .map(|source_file| v.visit_string(source_file)) + .transpose()?, + static_fields: { + let mut static_fields = HashMap::new(); + for (id, field) in self.static_fields.into_iter() { + static_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?); + } + static_fields + }, + instance_fields: { + let mut instance_fields = HashMap::new(); + for (id, field) in self.instance_fields.into_iter() { + instance_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?); + } + instance_fields + }, + direct_methods: { + let mut direct_methods = HashMap::new(); + for (id, meth) in self.direct_methods.into_iter() { + direct_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?); + } + direct_methods + }, + virtual_methods: { + let mut virtual_methods = HashMap::new(); + for (id, meth) in self.virtual_methods.into_iter() { + virtual_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?); + } + virtual_methods + }, + annotations: self + .annotations + .into_iter() + .map(|annotation| v.visit_annotation_item(annotation)) + .collect::>()?, + ..self + }) + } +} diff --git a/androscalpel/src/code.rs b/androscalpel/src/code.rs index ff5db0f..806ac22 100644 --- a/androscalpel/src/code.rs +++ b/androscalpel/src/code.rs @@ -8,6 +8,7 @@ use pyo3::prelude::*; use crate::{ ins::Instruction, DexString, IdField, IdMethod, IdMethodType, IdType, MethodHandle, Result, + Visitable, VisitableMut, Visitor, VisitorMut, }; // TODO: make this easy to edit/manipulate, maybe move to Method @@ -493,3 +494,38 @@ impl Code { }) } } + +impl Visitable for Code { + fn default_visit(&self, v: &mut V) -> Result<()> { + for ins in &self.insns { + v.visit_instruction(&ins)?; + } + Ok(()) + } +} +impl VisitableMut for Code { + fn default_visit_mut(self, v: &mut V) -> Result { + let parameter_names = if let Some(parameter_names) = self.parameter_names { + let mut new_param = vec![]; + for param in parameter_names { + if let Some(param) = param { + new_param.push(Some(v.visit_string(param)?)); + } else { + new_param.push(None); + } + } + Some(new_param) + } else { + None + }; + Ok(Self { + parameter_names, + insns: self + .insns + .into_iter() + .map(|ins| v.visit_instruction(ins)) + .collect::>()?, + ..self + }) + } +} diff --git a/androscalpel/src/field.rs b/androscalpel/src/field.rs index 643541d..410e442 100644 --- a/androscalpel/src/field.rs +++ b/androscalpel/src/field.rs @@ -7,7 +7,7 @@ use pyo3::prelude::*; use crate::{ DexAnnotationItem, DexString, DexValue, HiddenApiData, IdField, IdMethod, IdMethodType, IdType, - MethodHandle, Result, + MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut, }; use androscalpel_serializer::consts::*; @@ -257,3 +257,40 @@ impl Field { self == other } } + +impl Visitable for Field { + fn default_visit(&self, v: &mut V) -> Result<()> { + v.visit_field_id(&self.descriptor)?; + v.visit_field_visibility(&self.visibility)?; + if let Some(val) = &self.value { + v.visit_value(&val)?; + } + for annot in &self.annotations { + v.visit_annotation_item(&annot)?; + } + if let Some(hiddenapi) = &self.hiddenapi { + v.visit_hidden_api_data(&hiddenapi)?; + } + + Ok(()) + } +} +impl VisitableMut for Field { + fn default_visit_mut(self, v: &mut V) -> Result { + Ok(Self { + descriptor: v.visit_field_id(self.descriptor)?, + visibility: v.visit_field_visibility(self.visibility)?, + value: self.value.map(|val| v.visit_value(val)).transpose()?, + annotations: self + .annotations + .into_iter() + .map(|annot| v.visit_annotation_item(annot)) + .collect::>()?, + hiddenapi: self + .hiddenapi + .map(|hiddenapi| v.visit_hidden_api_data(hiddenapi)) + .transpose()?, + ..self + }) + } +} diff --git a/androscalpel/src/hiddenapi.rs b/androscalpel/src/hiddenapi.rs index 662f6b1..e8ac6f3 100644 --- a/androscalpel/src/hiddenapi.rs +++ b/androscalpel/src/hiddenapi.rs @@ -1,5 +1,6 @@ //! Representation of the hidden API data +use crate::{Result, Visitable, VisitableMut, Visitor, VisitorMut}; use log::warn; use pyo3::prelude::*; use serde::{Deserialize, Serialize}; @@ -39,6 +40,22 @@ impl From for androscalpel_serializer::HiddenApiFlags { } } +impl Visitable for HiddenApiData { + fn default_visit(&self, v: &mut V) -> Result<()> { + v.visit_hidden_api_permission(&self.permission)?; + v.visit_hidden_api_domain(&self.domain)?; + Ok(()) + } +} +impl VisitableMut for HiddenApiData { + fn default_visit_mut(self, v: &mut V) -> Result { + Ok(Self { + permission: v.visit_hidden_api_permission(self.permission)?, + domain: v.visit_hidden_api_domain(self.domain)?, + }) + } +} + #[pyclass] #[derive(Debug, PartialEq, Eq, Hash, Clone, Deserialize, Serialize)] pub enum HiddenApiPermission { diff --git a/androscalpel/src/method.rs b/androscalpel/src/method.rs index 1ce0614..448718d 100644 --- a/androscalpel/src/method.rs +++ b/androscalpel/src/method.rs @@ -7,7 +7,7 @@ use pyo3::prelude::*; use crate::{ Code, DexAnnotationItem, DexString, HiddenApiData, IdField, IdMethod, IdMethodType, IdType, - MethodHandle, Result, + MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut, }; use androscalpel_serializer::consts::*; @@ -298,3 +298,54 @@ impl Method { self == other } } + +impl Visitable for Method { + fn default_visit(&self, v: &mut V) -> Result<()> { + v.visit_method_id(&self.descriptor)?; + v.visit_method_visibility(&self.visibility)?; + for annot in &self.annotations { + v.visit_annotation_item(&annot)?; + } + for parameter_annotations in &self.parameters_annotations { + for annot in parameter_annotations { + v.visit_annotation_item(&annot)?; + } + } + if let Some(hiddenapi) = &self.hiddenapi { + v.visit_hidden_api_data(&hiddenapi)?; + } + if let Some(code) = &self.code { + v.visit_code(&code)?; + } + Ok(()) + } +} +impl VisitableMut for Method { + fn default_visit_mut(self, v: &mut V) -> Result { + Ok(Self { + descriptor: v.visit_method_id(self.descriptor)?, + visibility: v.visit_method_visibility(self.visibility)?, + annotations: self + .annotations + .into_iter() + .map(|annot| v.visit_annotation_item(annot)) + .collect::>()?, + parameters_annotations: self + .parameters_annotations + .into_iter() + .map(|parameter_annotations| { + parameter_annotations + .into_iter() + .map(|annot| v.visit_annotation_item(annot)) + .collect::>() + }) + .collect::>()?, + hiddenapi: self + .hiddenapi + .map(|hiddenapi| v.visit_hidden_api_data(hiddenapi)) + .transpose()?, + code: self.code.map(|code| v.visit_code(code)).transpose()?, + ..self + }) + } +} diff --git a/androscalpel/src/visitor.rs b/androscalpel/src/visitor.rs index dbdbd40..2471b7b 100644 --- a/androscalpel/src/visitor.rs +++ b/androscalpel/src/visitor.rs @@ -1,8 +1,10 @@ //! The visitor trait and common implementations. use crate::{ - ins::Instruction, scalar::*, CallSite, DexAnnotation, DexString, DexValue, IdEnum, IdField, - IdMethod, IdMethodType, IdType, MethodHandle, Result, + ins::Instruction, scalar::*, CallSite, Class, Code, DexAnnotation, DexAnnotationItem, + DexString, DexValue, Field, FieldVisibility, HiddenApiData, HiddenApiDomain, + HiddenApiPermission, IdEnum, IdField, IdMethod, IdMethodType, IdType, Method, MethodHandle, + MethodVisibility, Result, }; use std::collections::HashSet; @@ -70,6 +72,36 @@ pub trait Visitor: Sized { fn visit_bool(&mut self, _: &DexBoolean) -> Result<()> { Ok(()) } + fn visit_class(&mut self, class: &Class) -> Result<()> { + class.default_visit(self) + } + fn visit_field(&mut self, field: &Field) -> Result<()> { + field.default_visit(self) + } + fn visit_method(&mut self, method: &Method) -> Result<()> { + method.default_visit(self) + } + fn visit_annotation_item(&mut self, annotation: &DexAnnotationItem) -> Result<()> { + annotation.default_visit(self) + } + fn visit_field_visibility(&mut self, _: &FieldVisibility) -> Result<()> { + Ok(()) + } + fn visit_method_visibility(&mut self, _: &MethodVisibility) -> Result<()> { + Ok(()) + } + fn visit_hidden_api_data(&mut self, hidden_api: &HiddenApiData) -> Result<()> { + hidden_api.default_visit(self) + } + fn visit_hidden_api_permission(&mut self, _: &HiddenApiPermission) -> Result<()> { + Ok(()) + } + fn visit_hidden_api_domain(&mut self, _: &HiddenApiDomain) -> Result<()> { + Ok(()) + } + fn visit_code(&mut self, code: &Code) -> Result<()> { + code.default_visit(self) + } } pub trait VisitorMut: Sized { @@ -136,6 +168,45 @@ pub trait VisitorMut: Sized { fn visit_bool(&mut self, val: DexBoolean) -> Result { Ok(val) } + fn visit_class(&mut self, class: Class) -> Result { + class.default_visit_mut(self) + } + fn visit_field(&mut self, field: Field) -> Result { + field.default_visit_mut(self) + } + fn visit_method(&mut self, method: Method) -> Result { + method.default_visit_mut(self) + } + fn visit_annotation_item( + &mut self, + annotation: DexAnnotationItem, + ) -> Result { + annotation.default_visit_mut(self) + } + fn visit_field_visibility(&mut self, visibility: FieldVisibility) -> Result { + Ok(visibility) + } + fn visit_method_visibility( + &mut self, + visibility: MethodVisibility, + ) -> Result { + Ok(visibility) + } + fn visit_hidden_api_data(&mut self, hidden_api: HiddenApiData) -> Result { + hidden_api.default_visit_mut(self) + } + fn visit_hidden_api_permission( + &mut self, + perm: HiddenApiPermission, + ) -> Result { + Ok(perm) + } + fn visit_hidden_api_domain(&mut self, domain: HiddenApiDomain) -> Result { + Ok(domain) + } + fn visit_code(&mut self, code: Code) -> Result { + code.default_visit_mut(self) + } } /// Trait for structures that can be visited.