annotation WIP

TODO.md

@@ -1,7 +1,8 @@
-- method
+- method (what's left to do?)
-- anotation
+- annotations (start from annotations_off in class def item)
+- generate .dex
 - code
-- edditable format
+- edditable code format
 - sanity checks
 - tests
 - DexValues will become a problem ? (eg need to clone the vector for the array)

androscalpel/src/annotation.rs

@@ -0,0 +1,74 @@
+//! Annotations (for class, fields, methods and parameters alike).
+
+use pyo3::prelude::*;
+use std::collections::HashMap;
+
+use crate::{dex_id::IdType, value::DexValue, DexString};
+
+/// An annotation.
+#[pyclass]
+#[derive(Debug, Clone)]
+pub struct Annotation {
+    // TODO: check the relation between type and encoded_value.
+    /// The type of the annotation.
+    #[pyo3(get, set)]
+    pub type_: IdType,
+    // TODO: the get/set will probably be wonky on the python side when edditing
+    /// The annotation elements.
+    #[pyo3(get, set)]
+    pub elements: HashMap<DexString, DexValue>, // TODO: check MemberName syntax?
+    // TODO: enforce exclusivity
+    /// If the annotation visibility is set to build
+    #[pyo3(get, set)]
+    pub visibility_build: bool,
+    /// If the annotation visibility is set to runtime
+    #[pyo3(get, set)]
+    pub visibility_runtime: bool,
+    /// If the annotation visibility is set to system
+    #[pyo3(get, set)]
+    pub visibility_system: bool,
+}
+
+#[pymethods]
+impl Annotation {
+    #[new]
+    pub fn new(type_: IdType, elements: HashMap<DexString, DexValue>) -> Self {
+        Self {
+            type_,
+            elements,
+            visibility_build: false,
+            visibility_runtime: false,
+            visibility_system: false,
+        }
+    }
+
+    pub fn __str__(&self) -> String {
+        self.__repr__()
+    }
+
+    pub fn __repr__(&self) -> String {
+        let type_ = self.type_.__str__();
+        let visibility =
+            if !(self.visibility_build || self.visibility_runtime || self.visibility_system) {
+                "none".into()
+            } else {
+                let mut visibility = String::new();
+                if self.visibility_build {
+                    visibility += "build";
+                }
+                if self.visibility_runtime {
+                    visibility += "runtime";
+                }
+                if self.visibility_system {
+                    visibility += "system";
+                }
+                visibility
+            };
+        let mut elts: String = "{".into();
+        for (key, val) in self.elements.iter() {
+            elts += &format!("{}: {}, ", key.__str__(), val.__str__());
+        }
+        elts += "}";
+        format!("Annotation({type_}, visibility: {visibility}, {elts})")
+    }
+}

androscalpel/src/apk.rs

@@ -1,4 +1,7 @@
 //! Representation of an apk.
+
+use std::collections::HashMap;
+
 use log::info;
 use pyo3::prelude::*;
 
@@ -86,6 +89,20 @@ impl Apk {
                 class_item.access_flags
             );
         }
+        let annotations_directory = if class_item.annotations_off == 0 {
+            None
+        } else {
+            Some(dex.get_struct_at_offset::<AnnotationDirectoryItem>(class_item.annotations_off)?)
+        };
+        let mut annotations = vec![];
+        if let Some(annotations_directory) = annotations_directory {
+            if annotations_directory.class_annotations_off != 0 {
+                annotations = Self::get_annotations_from_annotation_set_off(
+                    annotations_directory.class_annotations_off,
+                    dex,
+                )?;
+            }
+        }
         let mut static_fields = vec![];
         let mut instance_fields = vec![];
         let mut direct_methods = vec![];
@@ -139,6 +156,7 @@ impl Apk {
             static_fields,
             direct_methods,
             virtual_methods,
+            annotations,
         })
     }
 
@@ -149,6 +167,41 @@ impl Apk {
         ))
     }
 
+    pub fn get_annotations_from_annotation_set_off(
+        annotations_set_off: u32,
+        dex: &DexFileReader,
+    ) -> Result<Vec<Annotation>> {
+        let mut annotations = vec![];
+        for annotation_off in dex
+            .get_struct_at_offset::<AnnotationSetItem>(annotations_set_off)?
+            .entries
+            .iter()
+            .map(|entry| entry.annotation_off)
+        {
+            let item = dex.get_struct_at_offset::<AnnotationItem>(annotation_off)?;
+            let (visibility_build, visibility_runtime, visibility_system) = match item.visibility {
+                AnnotationVisibility::Build => (true, false, false),
+                AnnotationVisibility::Runtime => (false, true, false),
+                AnnotationVisibility::System => (false, false, true),
+            };
+            let type_ = Self::get_id_type_from_idx(item.annotation.type_idx.0 as usize, dex)?;
+            let mut elements = HashMap::new();
+            for elt in item.annotation.elements {
+                let name: DexString = dex.get_string(elt.name_idx.0)?.into();
+                let value = Self::encoded_value_to_dex_value(&elt.value, dex)?;
+                elements.insert(name, value);
+            }
+            annotations.push(Annotation {
+                type_,
+                elements,
+                visibility_build,
+                visibility_runtime,
+                visibility_system,
+            });
+        }
+        Ok(annotations)
+    }
+
     /// Return a [`IdMethodType`] from its idx.
     pub fn get_id_method_type_from_idx(idx: usize, dex: &DexFileReader) -> Result<IdMethodType> {
         let proto = dex.get_proto_id(idx)?;

androscalpel/src/class.rs

@@ -2,7 +2,7 @@
 
 use pyo3::prelude::*;
 
-use crate::{DexString, Field, Method};
+use crate::{Annotation, DexString, Field, Method};
 
 /// Represent an apk
 #[pyclass]
@@ -61,8 +61,7 @@ pub struct Class {
     pub virtual_methods: Vec<Method>,
     // Do we need to distinguish direct and virtual (all the other) methods?
     // Maybe overlapping descriptor (same name, class and proto?)
-
+    pub annotations: Vec<Annotation>,
-    // pub annotations: Option<()> // TODO
     // TODO: mix annotation data to fields / methods / class to make it more practicle
 }
 
@@ -86,6 +85,7 @@ impl Class {
             instance_fields: vec![],
             direct_methods: vec![],
             virtual_methods: vec![],
+            annotations: vec![],
         }
     }
 

androscalpel/src/lib.rs

@@ -2,6 +2,7 @@ use pyo3::class::basic::CompareOp;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
+pub mod annotation;
 pub mod apk;
 pub mod class;
 pub mod dex_id;
@@ -11,6 +12,7 @@ pub mod method_handle;
 pub mod scalar;
 pub mod value;
 
+pub use annotation::*;
 pub use apk::*;
 pub use class::*;
 pub use dex_id::*;
@@ -154,6 +156,13 @@ impl DexString {
     }
 }
 
+impl std::hash::Hash for DexString {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.get_utf16_size().hash(state);
+        self.get_bytes().hash(state);
+    }
+}
+
 /// Androscalpel.
 #[pymodule]
 fn androscalpel(_py: Python, m: &PyModule) -> PyResult<()> {
@@ -185,6 +194,7 @@ fn androscalpel(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<InvokeDirect>()?;
     m.add_class::<InvokeInterface>()?;
 
+    m.add_class::<Annotation>()?;
     m.add_class::<DexArray>()?;
     m.add_class::<IdAnnotation>()?;
     m.add_class::<FieldVisibility>()?;