Add core library with radare2 in axmlreader

sidra-asa · sidra-asa · commit 2b5013e319d4 · 2023-09-22T01:46:06.000+08:00
diff --git a/quark/core/axmlreader/__init__.py b/quark/core/axmlreader/__init__.py
@@ -8,6 +8,7 @@
 
 import pkg_resources
 import rzpipe
+import r2pipe
 
 # Resource Types Definition
 # Please reference to
@@ -84,7 +85,7 @@ class AxmlReader(object):
     A Class that parses the Android XML file
     """
 
-    def __init__(self, file_path, structure_path=None):
+    def __init__(self, file_path, core_library="rizin", structure_path=None):
         if structure_path is None:
             structure_path = pkg_resources.resource_filename(
                 "quark.core.axmlreader", "axml_definition"
@@ -96,10 +97,13 @@ def __init__(self, file_path, structure_path=None):
                 f" of Rizin in {structure_path}"
             )
 
-        self._rz = rzpipe.open(file_path)
-        self._rz.cmd(f"pfo {structure_path}")
+        if core_library == "rizin":
+            self._core = rzpipe.open(file_path)
+        else:
+            self._core = r2pipe.open(file_path)
+        self._core.cmd(f"pfo {structure_path}")
 
-        self._file_size = int(self._rz.cmd("i~size[1]"), 16)
+        self._file_size = int(self._core.cmd("i~size[1]"), 16)
         self._ptr = 0
 
         self._cache = {}
@@ -110,7 +114,7 @@ def __init__(self, file_path, structure_path=None):
             raise AxmlException("Filesize exceeds theoretical lower bound.")
 
         # File Header
-        header = self._rz.cmdj("pfj axml_ResChunk_header @ 0x0")
+        header = self._core.cmdj("pfj axml_ResChunk_header @ 0x0")
 
         self._data_type = header[0]["value"]
         self._axml_size = header[2]["value"]
@@ -133,7 +137,7 @@ def __init__(self, file_path, structure_path=None):
             return
 
         # String Pool
-        string_pool_header = self._rz.cmdj("pfj axml_ResStringPool_header @ 8")
+        string_pool_header = self._core.cmdj("pfj axml_ResStringPool_header @ 8")
 
         string_pool_size = string_pool_header[0]["value"][2]["value"]
 
@@ -163,18 +167,18 @@ def __init__(self, file_path, structure_path=None):
         self._stringCount = string_pool_header[1]["value"]
         stringStart = string_pool_header[4]["value"]
 
-        self._rz.cmd(f"f string_pool_header @ 0x8 ")
+        self._core.cmd(f"f string_pool_header @ 0x8 ")
         string_pool_index = header_size + self._ptr
-        self._rz.cmd(f"f string_pool_index @ { string_pool_index }")
+        self._core.cmd(f"f string_pool_index @ { string_pool_index }")
         string_pool_data = stringStart + self._ptr
-        self._rz.cmd(f"f string_pool_data @ { string_pool_data }")
+        self._core.cmd(f"f string_pool_data @ { string_pool_data }")
 
         self._ptr += string_pool_size
         if self._ptr >= self._axml_size:
             return
 
         # Resource Map (Optional)
-        header = self._rz.cmdj(f"pfj axml_ResChunk_header @ {self._ptr}")
+        header = self._core.cmdj(f"pfj axml_ResChunk_header @ {self._ptr}")
 
         header_type = header[0]["value"]
         if header_type == RES_XML_RESOURCE_MAP_TYPE:
@@ -201,7 +205,7 @@ def __iter__(self) -> Iterator[ResChunkHeader]:
         :yield: header of a resource chunk defined in the binary
         """
         while self._axml_size - self._ptr >= 16:
-            header = self._rz.cmdj(f"pfj axml_ResXMLTree_node @ {self._ptr}")
+            header = self._core.cmdj(f"pfj axml_ResXMLTree_node @ {self._ptr}")
 
             node_type = header[0]["value"][0]["value"]
             header_size = header[0]["value"][1]["value"]
@@ -224,7 +228,7 @@ def __iter__(self) -> Iterator[ResChunkHeader]:
             chunk = {"Address": self._ptr, "Type": node_type}
 
             if node_type == RES_XML_START_ELEMENT_TYPE:
-                ext = self._rz.cmdj(
+                ext = self._core.cmdj(
                     f"pfj axml_ResXMLTree_attrExt @ { ext_ptr }"
                 )
 
@@ -235,7 +239,7 @@ def __iter__(self) -> Iterator[ResChunkHeader]:
                 # node['AttrCount'] = ext[4]['value']
 
             elif node_type == RES_XML_END_ELEMENT_TYPE:
-                ext = self._rz.cmdj(
+                ext = self._core.cmdj(
                     f"pfj axml_ResXMLTree_endElementExt @ { ext_ptr }"
                 )
 
@@ -246,15 +250,15 @@ def __iter__(self) -> Iterator[ResChunkHeader]:
                 RES_XML_START_NAMESPACE_TYPE,
                 RES_XML_END_NAMESPACE_TYPE,
             ]:
-                ext = self._rz.cmdj(
+                ext = self._core.cmdj(
                     f"pfj axml_ResXMLTree_namespaceExt @ { ext_ptr }"
                 )
 
                 chunk["Prefix"] = ext[0]["value"][0]["value"]
                 chunk["Uri"] = ext[1]["value"][0]["value"]
 
             elif node_type == RES_XML_CDATA_TYPE:
-                ext = self._rz.cmdj(
+                ext = self._core.cmdj(
                     f"pfj axml_ResXMLTree_cdataExt @ { ext_ptr }"
                 )
 
@@ -281,7 +285,7 @@ def get_string(self, index):
         if index < 0 or index >= self._stringCount:
             return None
 
-        return self._rz.cmdj(
+        return self._core.cmdj(
             f"pfj Z @ string_pool_data + `pfv n4 "
             f"@ string_pool_index+ {index}*4` + 2"
         )[0]["string"]
@@ -296,14 +300,14 @@ def get_attributes(self, chunk: ResChunkHeader) -> List[ResValue]:
             return None
         extAddress = int(chunk["Address"]) + 16
 
-        attrExt = self._rz.cmdj(f"pfj axml_ResXMLTree_attrExt @ {extAddress}")
+        attrExt = self._core.cmdj(f"pfj axml_ResXMLTree_attrExt @ {extAddress}")
 
         attrAddress = extAddress + attrExt[2]["value"]
         attributeSize = attrExt[3]["value"]
         attributeCount = attrExt[4]["value"]
         attributes = []
         for _ in range(attributeCount):
-            attr = self._rz.cmdj(
+            attr = self._core.cmdj(
                 f"pfj axml_ResXMLTree_attribute @ {attrAddress}"
             )
 
@@ -402,6 +406,6 @@ def get_xml_tree(self) -> XMLElementTree:
 
     def __del__(self):
         try:
-            self._rz.quit()
+            self._core.quit()
         except BaseException:
             pass
diff --git a/quark/core/r2apkinfo.py b/quark/core/r2apkinfo.py
@@ -228,7 +228,7 @@ def permissions(self) -> List[str]:
 
         :return: a list of permissions.
         """
-        axml = AxmlReader(self._manifest)
+        axml = AxmlReader(self._manifest, core_library="radare2")
         permission_list = set()
 
         for tag in axml:
@@ -249,7 +249,7 @@ def application(self) -> XMLElement:
         :return: an application element
         """
 
-        axml = AxmlReader(self._manifest)
+        axml = AxmlReader(self._manifest, core_library="radare2")
         root = axml.get_xml_tree()
 
         return root.find("application")
@@ -261,7 +261,7 @@ def activities(self) -> List[XMLElement]:
 
         :return: a list of all activities
         """
-        axml = AxmlReader(self._manifest)
+        axml = AxmlReader(self._manifest, core_library="radare2")
         root = axml.get_xml_tree()
 
         return root.findall("application/activity")
@@ -273,7 +273,7 @@ def receivers(self) -> List[XMLElement]:
 
         :return: a list of all receivers
         """
-        axml = AxmlReader(self._manifest)
+        axml = AxmlReader(self._manifest, core_library="radare2")
         root = axml.get_xml_tree()
 
         return root.findall("application/receiver")
diff --git a/tests/core/test_axmlreader.py b/tests/core/test_axmlreader.py
@@ -11,6 +11,16 @@
 from quark.core.axmlreader import AxmlReader, ResValue
 
 
+
+@pytest.fixture(
+    scope="function",
+    params=(("radare2"), ("rizin")),
+)
+def core_library(request):
+    core_lib = request.param
+    yield core_lib
+
+
 def extractManifest(samplePath: PathLike) -> str:
     folder = Path(samplePath).parent
 
@@ -27,8 +37,8 @@ def MANIFEST_PATH_14d9f(SAMPLE_PATH_14d9f):
 
 class TestAxmlReader:
     @staticmethod
-    def testIter(MANIFEST_PATH_14d9f) -> None:
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testIter(core_library, MANIFEST_PATH_14d9f) -> None:
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         expectedTag = {"Address": 3728, "Type": 256, "Prefix": 9, "Uri": 10}
 
         tag = next(iter(axmlReader))
@@ -37,23 +47,23 @@ def testIter(MANIFEST_PATH_14d9f) -> None:
         helper.assertDictEqual(tag, expectedTag)
 
     @staticmethod
-    def testFileSize(MANIFEST_PATH_14d9f):
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testFileSize(core_library, MANIFEST_PATH_14d9f):
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         assert axmlReader.file_size == 7676
 
     @staticmethod
-    def testAxmlSize(MANIFEST_PATH_14d9f):
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testAxmlSize(core_library, MANIFEST_PATH_14d9f):
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         assert axmlReader.axml_size == 7676
 
     @staticmethod
-    def testGetString(MANIFEST_PATH_14d9f):
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testGetString(core_library, MANIFEST_PATH_14d9f):
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         assert axmlReader.get_string(13) == "manifest"
 
     @staticmethod
-    def testGetAttributes(MANIFEST_PATH_14d9f):
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testGetAttributes(core_library, MANIFEST_PATH_14d9f):
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         manifestTag = list(axmlReader)[1]
 
         expectedAttributes = [
@@ -68,8 +78,8 @@ def testGetAttributes(MANIFEST_PATH_14d9f):
             assert expectedAttrib == attrib
 
     @staticmethod
-    def testGetXmlTree(MANIFEST_PATH_14d9f):
-        axmlReader = AxmlReader(MANIFEST_PATH_14d9f)
+    def testGetXmlTree(core_library, MANIFEST_PATH_14d9f):
+        axmlReader = AxmlReader(MANIFEST_PATH_14d9f, core_library)
         xml = axmlReader.get_xml_tree()
         manifestLabel = xml.getroot()
         assert len(manifestLabel.findall("uses-sdk")) == 1
