Update Quark class to support this feature

Author: haeter525

quark/core/axmlreader/__init__.py

@@ -1,8 +1,8 @@
 import enum
 import functools
 import os.path
-import pkg_resources
 
+import pkg_resources
 import rzpipe
 
 # Resource Types Definition
@@ -71,7 +71,19 @@ class AxmlReader(object):
     A Class that parses the Android XML file
     """
 
-    def __init__(self, file_path, structure_path=None):
+    def __init__(self, file_path, structure_path=None, rizin_path=None):
+        """
+        Create an AxmlReader object to parse the given Android XML file like
+        AndroidManifest.xml.
+
+        :param file_path: a file in the Android XML format
+        :param structure_path: a plain text file defining the structures of
+        the Android XML format. Defaults to None
+        :param rizin_path: a PathLike object to specify a Rizin executable to
+        use. Defaults to None
+        :raises AxmlException: if the given file is an invalid Android XML
+        file
+        """
         if structure_path is None:
             structure_path = pkg_resources.resource_filename(
                 "quark.core.axmlreader", "axml_definition"
@@ -83,7 +95,7 @@ def __init__(self, file_path, structure_path=None):
                 f" of Rizin in {structure_path}"
             )
 
-        self._rz = rzpipe.open(file_path)
+        self._rz = rzpipe.open(file_path, rizin_home=rizin_path)
         self._rz.cmd(f"pfo {structure_path}")
 
         self._file_size = int(self._rz.cmd("i~size[1]"), 16)

quark/core/parallelquark.py

@@ -2,8 +2,8 @@
 # This file is part of Quark-Engine - https://github.com/quark-engine/quark-engine
 # See the file 'LICENSE' for copying permission.
 
-from multiprocessing.pool import Pool
 from multiprocessing import cpu_count
+from multiprocessing.pool import Pool
 
 from quark.core.analysis import QuarkAnalysis
 from quark.core.quark import Quark
@@ -13,12 +13,33 @@
 
 class ParallelQuark(Quark):
     @staticmethod
-    def _worker_initializer(apk, core_library):
+    def _worker_initializer(apk, core_library, rizin_path):
+        """
+        An initializer that creates multiple Quark object for a subprocess
+        pool.
+
+        :param apk: an APK for Quark to analyze
+        :param core_library: a string indicating which analysis library Quark
+        should use
+        :param rizin_path: a PathLike object to specify a Rizin executable for
+        the Rizin-based analysis library
+        """
+
         global _quark
-        _quark = Quark(apk, core_library)
+        _quark = Quark(apk, core_library, rizin_path)
 
     @staticmethod
     def _worker_analysis(rule_obj):
+        """
+        A function for the subprocesses in the pool to analyze the target APK
+        with the specified rule.
+
+        :param rule_obj: a Quark rule that the subprocess will analyze
+        :return: a tuple of the analysis result, including the reached stage,
+        the parent functions, the detected behavior (rule), and the
+        corresponding bytecodes.
+        """
+
         _quark.quark_analysis = QuarkAnalysis()
         _quark.run(rule_obj)
 
@@ -32,7 +53,8 @@ def to_raw_method(methodobject):
 
         reached_stage = rule_obj.check_item.count(True)
         level_4_result = tuple(
-            to_raw_method(method) for method in _quark.quark_analysis.level_4_result
+            to_raw_method(method)
+            for method in _quark.quark_analysis.level_4_result
         )
         behavior_list = [
             (
@@ -51,6 +73,13 @@ def to_raw_method(methodobject):
         )
 
     def _apply_analysis_result(self, rule_obj):
+        """
+        Append the result returned by the subprocesses into the QuarkAnalysis
+        object.
+
+        :param rule_obj: a Quark rule specifying which result this method
+        should append
+        """
         async_result = self._result_map[id(rule_obj)]
         result = async_result.get()
 
@@ -112,21 +141,43 @@ def _apply_analysis_result(self, rule_obj):
                 ]
             )
 
-    def __init__(self, apk, core_library, num_of_process=1):
+    def __init__(self, apk, core_library, num_of_process=1, rizin_path=None):
+        """
+        Create a ParallelQuark object to run the Quark analysis parallelly.
+
+        :param apk: an APK for Quark to analyze
+        :param core_library: a string indicating which analysis library Quark
+        should use
+        :param num_of_process: a value indicating the maximal number of
+        available processes for the analysis. Defaults to 1
+        :param rizin_path: a PathLike object to specify a Rizin executable for
+        the Rizin-based analysis library. Default to None
+        """
         self._result_map = {}
         self._pool = Pool(
-            min(num_of_process, cpu_count() - 1), self._worker_initializer,
-            (apk, core_library)
+            min(num_of_process, cpu_count() - 1),
+            self._worker_initializer,
+            (apk, core_library, rizin_path),
         )
 
-        super().__init__(apk, core_library)
+        super().__init__(apk, core_library, rizin_path=rizin_path)
 
     def apply_rules(self, rule_obj_list):
+        """
+        Add Quark rules to this pool.
+
+        :param rule_obj_list: a list of Quark rules
+        """
         for rule_obj in rule_obj_list:
             result = self._pool.apply_async(self._worker_analysis, (rule_obj,))
             self._result_map[id(rule_obj)] = result
 
     def run(self, rule_obj):
+        """
+        Wait for all rules to be analyzed.
+
+        :param rule_obj: _description_
+        """
         self._apply_analysis_result(rule_obj)
 
     def close(self):

quark/core/quark.py

@@ -38,14 +38,18 @@
 class Quark:
     """Quark module is used to check quark's five-stage theory"""
 
-    def __init__(self, apk, core_library="androguard"):
+    def __init__(self, apk, core_library="androguard", rizin_path=None):
         """
+        Create a Quark object.
 
-        :param apk: the filename of the apk.
+        :param apk: an APK for Quark to analyze
+        :param core_library: a string indicating which analysis library Quark should use. Defaults to "androguard"
+        :param rizin_path: a PathLike object to specify a Rizin executable for the Rizin-based analysis library. Defaults to None
+        :raises ValueError: if an unknown core library is specified
         """
         core_library = core_library.lower()
         if core_library == "rizin":
-            self.apkinfo = RizinImp(apk)
+            self.apkinfo = RizinImp(apk, rizin_path=rizin_path)
         elif core_library == "androguard":
             self.apkinfo = AndroguardImp(apk)
         else:

quark/core/rzapkinfo.py

@@ -18,7 +18,8 @@
 from quark.core.interface.baseapkinfo import BaseApkinfo
 from quark.core.struct.bytecodeobject import BytecodeObject
 from quark.core.struct.methodobject import MethodObject
-from quark.utils.tools import descriptor_to_androguard_format, remove_dup_list
+from quark.utils.tools import (descriptor_to_androguard_format,
+                               get_rizin_version, remove_dup_list)
 
 RizinCache = namedtuple("rizin_cache", "address dexindex is_imported")
 
@@ -51,9 +52,20 @@ def __init__(
         self,
         apk_filepath: Union[str, PathLike],
         tmp_dir: Union[str, PathLike] = None,
+        rizin_path: PathLike = None,
     ):
         super().__init__(apk_filepath, "rizin")
 
+        if rizin_path:
+            if not get_rizin_version(rizin_path):
+                raise ValueError(
+                    f"The file in {rizin_path} is not a valid Rizin executable."
+                )
+
+            rizin_path = os.path.dirname(rizin_path)
+
+        self.rizin_path = rizin_path
+
         if self.ret_type == "DEX":
             self._tmp_dir = None
             self._dex_list = [apk_filepath]
@@ -84,7 +96,7 @@ def __init__(
 
     @functools.lru_cache
     def _get_rz(self, index):
-        rz = rzpipe.open(self._dex_list[index])
+        rz = rzpipe.open(self._dex_list[index], rizin_home=self.rizin_path)
         rz.cmd("aa")
         return rz
 
@@ -227,7 +239,7 @@ def _get_methods_classified(self, dexindex):
 
     @functools.cached_property
     def permissions(self) -> List[str]:
-        axml = AxmlReader(self._manifest)
+        axml = AxmlReader(self._manifest, rizin_path=self.rizin_path)
         permission_list = set()
 
         for tag in axml: