droidasc 0.1.0__py3-none-any.whl

droidasc/__init__.py

@@ -0,0 +1,3 @@
+from droidasc.cli import main
+
+__all__ = ["main"]

droidasc/__main__.py

@@ -0,0 +1,3 @@
+from droidasc.cli import main
+
+main()

droidasc/asc_client/apk_handler.py

@@ -0,0 +1,463 @@
+import atexit
+import mmap
+import os
+import struct
+import sys
+import threading
+import time
+import types
+import zlib
+from concurrent.futures import ThreadPoolExecutor, FIRST_COMPLETED, wait
+
+from droidasc.asc_client.dex_container import iter_logical_dex_buffers
+
+
+_U16 = struct.Struct("<H")
+_U32 = struct.Struct("<I")
+_U16_FROM = _U16.unpack_from
+_U32_FROM = _U32.unpack_from
+_CD_SIG = b"PK\x01\x02"
+_LH_SIG = b"PK\x03\x04"
+_EOCD_SIG = b"PK\x05\x06"
+_DEX_SUFFIX = b".dex"
+_SLASH = ord("/")
+_DEFLATE_CHUNK = 1 << 19
+
+_WORKER_APK_PATH = None
+_WORKER_APK_FP = None
+_WORKER_APK_MM = None
+
+
+def _process_pool_context():
+    """Prefer fork so workers do not have to re-import the whole module graph.
+
+    spawn re-executes the parent's __main__ and re-imports every module a worker
+    touches; measured on the 352MB sample that is 85ms of pool cost versus 19ms for
+    fork. fork is only used when it is safe: never on Windows, and never from a
+    multi-threaded parent (fork() there can deadlock on inherited locks).
+    """
+    if os.name == "nt":
+        return None
+    try:
+        import multiprocessing
+        # decompiler.py stubs optional imports for startup; a stub is not a real
+        # module and would be accepted silently by __getattr__, so verify before
+        # handing it to ProcessPoolExecutor (a bogus context hangs instead of raising)
+        if not isinstance(multiprocessing, types.ModuleType) or not hasattr(multiprocessing, "get_context"):
+            return None
+        if threading.active_count() != 1:
+            return None
+        return multiprocessing.get_context("fork")
+    except (ImportError, ValueError, OSError):
+        return None
+
+
+def _skip_uleb128(buf, off : int) -> int:
+    while buf[off] & 0x80:
+        off += 1
+    return off + 1
+
+
+def _read_string_data_bytes(buf, str_off : int) -> bytes:
+    ptr = _skip_uleb128(buf, str_off)
+    end = buf.find(b"\x00", ptr)
+    if end < 0:
+        raise ValueError("unterminated string_data_item")
+    return buf[ptr:end]
+
+
+def _find_type_idx(buf, target_bytes : bytes) -> int:
+    if len(buf) < 0x70 or buf[:3] != b"dex":
+        return -1
+
+    string_ids_size = _U32_FROM(buf, 0x38)[0]
+    string_ids_off = _U32_FROM(buf, 0x3C)[0]
+    type_ids_size = _U32_FROM(buf, 0x40)[0]
+    type_ids_off = _U32_FROM(buf, 0x44)[0]
+
+    if string_ids_size == 0 or type_ids_size == 0:
+        return -1
+    if string_ids_off + string_ids_size * 4 > len(buf):
+        raise ValueError("bad string_ids range")
+    if type_ids_off + type_ids_size * 4 > len(buf):
+        raise ValueError("bad type_ids range")
+
+    left = 0
+    right = type_ids_size - 1
+    while left <= right:
+        mid = (left + right) >> 1
+        str_idx = _U32_FROM(buf, type_ids_off + (mid << 2))[0]
+        if str_idx >= string_ids_size:
+            raise ValueError("bad type_id->string_idx")
+        str_off = _U32_FROM(buf, string_ids_off + (str_idx << 2))[0]
+        if str_off >= len(buf):
+            raise ValueError("bad string_data_off")
+
+        cls_bytes = _read_string_data_bytes(buf, str_off)
+        if cls_bytes == target_bytes:
+            return mid
+        if cls_bytes < target_bytes:
+            left = mid + 1
+        else:
+            right = mid - 1
+
+    return -1
+
+
+def _class_defs_contains_type_idx(buf, type_idx : int) -> bool:
+    class_defs_size = _U32_FROM(buf, 0x60)[0]
+    class_defs_off = _U32_FROM(buf, 0x64)[0]
+
+    if class_defs_size == 0:
+        return False
+
+    class_defs_end = class_defs_off + (class_defs_size << 5)
+    if class_defs_end > len(buf):
+        raise ValueError("bad class_defs range")
+
+    needle = type_idx.to_bytes(4, "little")
+    pos = buf.find(needle, class_defs_off, class_defs_end)
+    while pos != -1:
+        if ((pos - class_defs_off) & 0x1F) == 0:
+            return True
+        pos = buf.find(needle, pos + 1, class_defs_end)
+    return False
+
+
+def _dex_defines_class(buf, target_bytes : bytes) -> bool:
+    type_idx = _find_type_idx(buf, target_bytes)
+    if type_idx < 0:
+        return False
+    return _class_defs_contains_type_idx(buf, type_idx)
+
+
+def _find_eocd(mm : mmap.mmap) -> int:
+    search_start = max(0, len(mm) - 65536 - 22)
+    return mm.rfind(_EOCD_SIG, search_start)
+
+
+def _parse_cd_dex_entries(mm : mmap.mmap):
+    eocd_idx = _find_eocd(mm)
+    if eocd_idx < 0:
+        raise ValueError("EOCD not found")
+
+    cd_size = _U32_FROM(mm, eocd_idx + 12)[0]
+    cd_off = _U32_FROM(mm, eocd_idx + 16)[0]
+    cd_end = cd_off + cd_size
+    entries = []
+    seen_names = set()
+    pos = cd_off
+
+    while True:
+        pos = mm.find(b"classes", pos, cd_end)
+        if pos < 0:
+            break
+        header_off = pos - 46
+        pos += 7
+        if header_off < cd_off or mm[header_off:header_off + 4] != _CD_SIG:
+            continue
+
+        name_len = _U16_FROM(mm, header_off + 28)[0]
+        name_start = pos - 7
+        name_end = name_start + name_len
+        if name_end > cd_end:
+            continue
+
+        name_bytes = mm[name_start:name_end]
+        if (
+            not name_bytes.endswith(_DEX_SUFFIX)
+            or _SLASH in name_bytes
+            or name_bytes in seen_names
+        ):
+            continue
+
+        seen_names.add(name_bytes)
+        entries.append((
+            name_bytes.decode("utf-8", errors="ignore"),
+            _U32_FROM(mm, header_off + 24)[0],
+            _U32_FROM(mm, header_off + 20)[0],
+            _U32_FROM(mm, header_off + 42)[0],
+            _U16_FROM(mm, header_off + 10)[0],
+        ))
+
+    if entries:
+        return entries
+
+    ptr = cd_off
+    while ptr + 46 <= cd_end:
+        if mm[ptr:ptr + 4] != _CD_SIG:
+            break
+
+        name_len = _U16_FROM(mm, ptr + 28)[0]
+        extra_len = _U16_FROM(mm, ptr + 30)[0]
+        comment_len = _U16_FROM(mm, ptr + 32)[0]
+        name_start = ptr + 46
+        name_end = name_start + name_len
+        name_bytes = mm[name_start:name_end]
+        if name_bytes.endswith(_DEX_SUFFIX) and _SLASH not in name_bytes:
+            entries.append((
+                name_bytes.decode("utf-8", errors="ignore"),
+                _U32_FROM(mm, ptr + 24)[0],
+                _U32_FROM(mm, ptr + 20)[0],
+                _U32_FROM(mm, ptr + 42)[0],
+                _U16_FROM(mm, ptr + 10)[0],
+            ))
+
+        ptr = name_end + extra_len + comment_len
+
+    return entries
+
+
+def _close_worker_apk():
+    global _WORKER_APK_PATH, _WORKER_APK_FP, _WORKER_APK_MM
+
+    mm = _WORKER_APK_MM
+    fp = _WORKER_APK_FP
+    _WORKER_APK_MM = None
+    _WORKER_APK_FP = None
+    _WORKER_APK_PATH = None
+
+    if mm is not None:
+        mm.close()
+    if fp is not None:
+        fp.close()
+
+
+atexit.register(_close_worker_apk)
+
+
+def _get_worker_apk_mm(apk_path : str):
+    global _WORKER_APK_PATH, _WORKER_APK_FP, _WORKER_APK_MM
+
+    mm = _WORKER_APK_MM
+    if mm is not None and _WORKER_APK_PATH == apk_path:
+        return mm
+
+    _close_worker_apk()
+    fp = open(apk_path, "rb")
+    mm = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ)
+    _WORKER_APK_PATH = apk_path
+    _WORKER_APK_FP = fp
+    _WORKER_APK_MM = mm
+    return mm
+
+
+def _inflate_deflate_chunks(comp_view, stop_event):
+    if stop_event is None:
+        # nothing to cancel: hand zlib the whole stream in one call instead of
+        # re-entering it every _DEFLATE_CHUNK bytes (measured ~4% faster per dex)
+        return zlib.decompress(comp_view, -15)
+    decomp = zlib.decompressobj(-15)
+    out = bytearray()
+    pos = 0
+    total = len(comp_view)
+    while pos < total:
+        if stop_event.is_set():
+            return None
+        end = min(pos + _DEFLATE_CHUNK, total)
+        out.extend(decomp.decompress(comp_view[pos:end]))
+        pos = end
+    out.extend(decomp.flush())
+    if stop_event.is_set():
+        return None
+    return bytes(out)
+
+
+def _inflate_dex(mm : mmap.mmap, entry, stop_event = None):
+    name, uncomp_size, comp_size, local_header_off, comp_method = entry
+    if mm[local_header_off:local_header_off + 4] != _LH_SIG:
+        raise ValueError("bad local header signature")
+
+    name_len = _U16_FROM(mm, local_header_off + 26)[0]
+    extra_len = _U16_FROM(mm, local_header_off + 28)[0]
+    data_off = local_header_off + 30 + name_len + extra_len
+    comp_view = memoryview(mm)[data_off:data_off + comp_size]
+
+    if comp_method == 0:
+        data = bytes(comp_view)
+    elif comp_method == 8:
+        data = _inflate_deflate_chunks(comp_view, stop_event)
+        if data is None:
+            return None
+    else:
+        raise ValueError(f"unsupported compression method: {comp_method}")
+
+    if uncomp_size and len(data) != uncomp_size:
+        raise ValueError(f"size mismatch: expect {uncomp_size}, got {len(data)}")
+    return data
+
+
+def _findrefs_worker(apk_path : str, entry, find_type : str, find : dict, aggregate : bool = True):
+    from droidasc.asc_client.asc_handler import AscHandler
+
+    mm = _get_worker_apk_mm(apk_path)
+    t0 = time.perf_counter()
+    data = _inflate_dex(mm, entry)
+    t1 = time.perf_counter()
+    lines = []
+    handler = AscHandler(False)
+    for dex_name, dex_buf in iter_logical_dex_buffers(entry[0], data):
+        lines.extend(handler.findrefs(dex_name, dex_buf, find_type, find, aggregate=aggregate))
+    t2 = time.perf_counter()
+    return (
+        entry[0],
+        lines,
+        (t1 - t0) * 1000000,
+        (t2 - t1) * 1000000,
+        os.getpid(),
+    )
+
+
+def _inflate_and_hit(mm : mmap.mmap, entry, target_bytes : bytes, stop_event : threading.Event, log):
+    tid = threading.get_ident() & 0xFFFF
+    name = entry[0]
+    if stop_event.is_set():
+        return False, None, None
+
+    t0 = time.perf_counter()
+    data = _inflate_dex(mm, entry, stop_event)
+    if data is None:
+        return False, None, None
+    t1 = time.perf_counter()
+    if stop_event.is_set():
+        return False, None, None
+
+    hit = False
+    hit_name = None
+    hit_data = None
+    for dex_name, dex_buf in iter_logical_dex_buffers(name, data):
+        if stop_event.is_set():
+            return False, None, None
+        if _dex_defines_class(dex_buf, target_bytes):
+            hit = True
+            hit_name = dex_name
+            hit_data = dex_buf
+            break
+    t2 = time.perf_counter()
+    if hit:
+        stop_event.set()
+    log(
+        f"[APK] [T{tid:04x}] '{name}' inflate={(t1 - t0) * 1000000:.2f} us "
+        f"lookup={(t2 - t1) * 1000000:.2f} us hit={hit}"
+    )
+    return hit, hit_name, hit_data
+
+
+class ApkHandler:
+    def __init__(self, apk_path : str, debug : bool = False, max_workers : int = 8):
+        self.apk_path = apk_path
+        self.debug = debug
+        self.max_workers = max_workers
+
+    def _log(self, msg : str):
+        if self.debug:
+            print(msg)
+
+    def _open_apk(self):
+        fp = open(self.apk_path, "rb")
+        mm = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ)
+        return fp, mm
+
+    def get_class_dex(self, dalvik_class : str):
+        target_bytes = dalvik_class.encode("utf-8")
+        t_start = time.perf_counter()
+        fp, mm = self._open_apk()
+        try:
+            entries = _parse_cd_dex_entries(mm)
+            entries.sort(key=lambda x: x[2])
+            if self.debug:
+                t_ready = time.perf_counter()
+                self._log(f"[APK] scan setup={(t_ready - t_start) * 1000000:.2f} us entries={len(entries)}")
+            if not entries:
+                return None
+
+            stop_event = threading.Event()
+            hit_name = None
+            hit_data = None
+            cap = min(len(entries), max(6, min(self.max_workers, 12)))
+
+            with ThreadPoolExecutor(max_workers=self.max_workers) as ex:
+                inflight = {}
+                idx = 0
+                while idx < len(entries) or inflight:
+                    while idx < len(entries) and len(inflight) < cap and not stop_event.is_set():
+                        entry = entries[idx]
+                        idx += 1
+                        fut = ex.submit(_inflate_and_hit, mm, entry, target_bytes, stop_event, self._log)
+                        inflight[fut] = entry
+
+                    if not inflight:
+                        break
+
+                    done, _pending = wait(list(inflight.keys()), return_when=FIRST_COMPLETED)
+                    for fut in done:
+                        entry = inflight.pop(fut)
+                        ok, dex_name, data = fut.result()
+                        if ok:
+                            hit_name = dex_name
+                            hit_data = data
+                            stop_event.set()
+                            break
+
+                    if hit_name is not None:
+                        for fut in inflight:
+                            fut.cancel()
+                        break
+
+            if self.debug:
+                t_end = time.perf_counter()
+                self._log(f"[APK] getclass total={(t_end - t_start) * 1000000:.2f} us")
+            if hit_name is None:
+                return None
+            return hit_name, hit_data
+        finally:
+            mm.close()
+            fp.close()
+
+    def for_each_findrefs(self, find_type : str, find : dict):
+        # imported before the timer so this once-per-process import is not charged to a
+        # single search; it used to happen at apk_handler import time, and keeping the
+        # getclass path (thread pool only) from paying for it is worth ~8ms of startup
+        from concurrent.futures import ProcessPoolExecutor
+
+        t_start = time.perf_counter()
+        fp, mm = self._open_apk()
+        try:
+            entries = _parse_cd_dex_entries(mm)
+            entries.sort(key=lambda x: x[2])
+        finally:
+            mm.close()
+            fp.close()
+
+        if not entries:
+            return
+
+        # a forked child inherits whatever is still sitting in the parent's stdio
+        # buffers; flush first so it cannot be emitted a second time on exit
+        sys.stdout.flush()
+        sys.stderr.flush()
+        with ProcessPoolExecutor(max_workers=self.max_workers,
+                                 mp_context=_process_pool_context()) as ex:
+            futures = {}
+            for entry in entries:
+                fut = ex.submit(_findrefs_worker, self.apk_path, entry, find_type, find)
+                futures[fut] = entry[0]
+
+            while futures:
+                done, _pending = wait(list(futures.keys()), return_when=FIRST_COMPLETED)
+                for fut in done:
+                    futures.pop(fut)
+                    dex_name, lines, inflate_us, process_us, pid = fut.result()
+                    if self.debug:
+                        self._log(
+                            f"[APK] [P{pid}] '{dex_name}' inflate={inflate_us:.2f} us "
+                            f"process={process_us:.2f} us"
+                        )
+                    yield dex_name, lines
+
+        if self.debug:
+            t_end = time.perf_counter()
+            self._log(
+                f"[APK] for_each_findrefs total={(t_end - t_start) * 1000000:.2f} us "
+                f"count={len(entries)} workers={self.max_workers}"
+            )

droidasc/asc_client/asc_handler.py

@@ -0,0 +1,129 @@
+import copy
+import importlib.util
+import os
+import sys
+import types
+from collections import defaultdict
+
+
+_DexManager = None
+_FindRefManager = None
+_decompile_dex_bytes = None
+_DEX = None
+
+
+def _install_pure_python_mutf8_shim():
+    if "mutf8.cmutf8" in sys.modules:
+        return
+
+    pkg_spec = importlib.util.find_spec("mutf8")
+    if pkg_spec is None or not pkg_spec.submodule_search_locations:
+        return
+
+    pkg_dir = pkg_spec.submodule_search_locations[0]
+    py_impl = os.path.join(pkg_dir, "mutf8.py")
+    mod_spec = importlib.util.spec_from_file_location("_asc_client_mutf8_py", py_impl)
+    if mod_spec is None or mod_spec.loader is None:
+        return
+
+    module = importlib.util.module_from_spec(mod_spec)
+    mod_spec.loader.exec_module(module)
+
+    shim = types.ModuleType("mutf8.cmutf8")
+    shim.decode_modified_utf8 = module.decode_modified_utf8
+    shim.encode_modified_utf8 = module.encode_modified_utf8
+    sys.modules["mutf8.cmutf8"] = shim
+
+
+def _lazy_import():
+    global _DexManager, _FindRefManager, _decompile_dex_bytes, _DEX
+    if _DexManager is not None:
+        return
+
+    _install_pure_python_mutf8_shim()
+    from droidasc.asc_core.core.dex.dex_manager import DexManager
+    from droidasc.asc_core.findrefs.findrefs_manager import FindRefManager
+    from droidasc.asc_core.utils.tinydex import DEX
+    from droidasc.asc_core.utils.decompiler import decompile_dex_bytes
+
+    _DexManager = DexManager
+    _FindRefManager = FindRefManager
+    _decompile_dex_bytes = decompile_dex_bytes
+    _DEX = DEX
+
+
+class AscHandler:
+    def __init__(self, debug : bool = False):
+        self.debug = debug
+
+    def getclass(self, dex_buf : bytes, dalvik_class : str) -> str:
+        _lazy_import()
+        manager = _DexManager(memoryview(dex_buf), debug=self.debug)
+        new_dex_bytes = manager.extract_and_rebuild(dalvik_class)
+        return _decompile_dex_bytes(new_dex_bytes, dalvik_class)
+
+    def _format_method(self, dex, midx : int) -> str:
+        method = dex.methods[midx]
+        return f"{method.cls.fullname}->{method.name}"
+
+    def _format_matched_name(self, dex, find_type : str, idx : int) -> str:
+        if find_type == "string":
+            return str(dex.strings[idx])
+        if find_type == "type":
+            return dex.types[idx].descriptor
+        if find_type == "method":
+            method = dex.methods[idx]
+            return f"{method.cls.fullname}->{method.name}"
+        field = dex.fields[idx]
+        return f"{field.cls.fullname}->{field.name}"
+
+    def findrefs(self, dex_name : str, dex_buf : bytes, find_type : str, find : dict, aggregate : bool = True) -> list:
+        from droidasc.asc_core.findrefs.findrefs_manager import FindRefManager
+        from droidasc.asc_core.utils.tinydex import DEX
+
+        dex = DEX.parse(memoryview(dex_buf), dex_name)
+        ref_manager = FindRefManager(dex)
+        query = copy.deepcopy(find)
+        matched_idxs = ref_manager.find_ref(query, True)
+        mids = query[find_type]
+        if not aggregate:
+            ret = []
+            for i in range(len(mids)):
+                midx = mids[i]
+                if midx is None:
+                    continue
+                idx = matched_idxs[i]
+                if isinstance(midx, list):
+                    midxs = midx
+                else:
+                    midxs = [midx]
+                matched = self._format_matched_name(dex, find_type, idx)
+                for mid in midxs:
+                    ret.append(
+                        f"{dex_name} | {self._format_method(dex, mid)} | matched=({matched})"
+                    )
+            return ret
+
+        grouped = defaultdict(set)
+        for i in range(len(mids)):
+            midx = mids[i]
+            if midx is None:
+                continue
+            idx = matched_idxs[i]
+            if isinstance(midx, list):
+                midxs = midx
+            else:
+                midxs = [midx]
+            for mid in midxs:
+                grouped[mid].add(idx)
+
+        ret = []
+        for mid in sorted(grouped):
+            matched = "; ".join(
+                self._format_matched_name(dex, find_type, idx)
+                for idx in sorted(grouped[mid])
+            )
+            ret.append(
+                f"{dex_name} | {self._format_method(dex, mid)} | matched=({matched})"
+            )
+        return ret