segment_names.py

import io
import os
import hashlib
import zipfile
from collections import defaultdict, namedtuple

import ida_idaapi
import ida_kernwin
import ida_dbg
import ida_segment
import ida_bytes
import ida_nalt

try:
    from elftools.elf.elffile import ELFFile
    from elftools.elf.sections import NoteSection
    from elftools.common.exceptions import ELFError
    _IMPORT_ERROR = None
except ImportError as _exc:  # pragma: no cover - environment dependent
    ELFFile = NoteSection = ELFError = None
    _IMPORT_ERROR = _exc


# --- Tunables ---------------------------------------------------------------
PF_X = 0x1              # ELF program-header executable flag
PF_W = 0x2             # ELF program-header writable flag
HASH_WINDOW = 0x1000    # bytes hashed to fingerprint a code segment (one page)
MIN_SO_SIZE = 0x1000    # ignore .so entries smaller than this
PAGE = 0x1000           # page size / segment-location tolerance
LIB_PREFIX = "lib/arm"  # matches both armeabi-v7a/ and arm64-v8a/

# What we keep about each PT_LOAD of a library.
Load = namedtuple("Load", "vaddr is_exec is_write")
# A memory region we successfully renamed, remembered for the rebase step.
Match = namedtuple("Match", "code_seg name basename load_base")


def _log(fmt, *args):
    ida_kernwin.msg("APKRenamer: " + (fmt % args if args else fmt) + "\n")


def _addr(ea):
    """Format an address, widening the field for 64-bit values."""
    return f"{ea:#018x}" if ea > 0xffffffff else f"{ea:#010x}"


def _page_floor(value):
    return value & ~(PAGE - 1)


class _SegmentChooser(ida_kernwin.Choose):
    """Modal picker used when more than one library matches the open file."""

    def __init__(self, title, rows):
        super().__init__(
            title,
            [["Segment", 40], ["Load base", 24]],
            flags=ida_kernwin.Choose.CH_MODAL,
        )
        self.rows = rows  # list of [name, "0x..."]

    def OnGetSize(self):
        return len(self.rows)

    def OnGetLine(self, n):
        return self.rows[n]


class ApkSegmentRenamerPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_KEEP
    comment = "Rename debugger segments by matching APK .so content; offer to rebase"
    help = __doc__
    wanted_name = "APK segment renamer + rebasing"
    wanted_hotkey = "Ctrl-Alt-A"

    # -- plugin lifecycle ---------------------------------------------------
    def init(self):
        return ida_idaapi.PLUGIN_OK

    def term(self):
        pass

    def run(self, arg):
        if _IMPORT_ERROR is not None:
            ida_kernwin.warning(
                "pyelftools is required.\nInstall it with:  pip install pyelftools")
            return
        if not ida_dbg.is_debugger_on():
            ida_kernwin.warning(
                "Run this during a debugging session (start debugging and pause "
                "the process first).")
            return

        apk_path = ida_kernwin.ask_file(False, "*.apk", "Select the corresponding APK file")
        if not apk_path:
            return

        current_file = self._current_file_name()
        _log("open file: %s", current_file or "<unknown>")
        _log("scanning %s ...", os.path.basename(apk_path))

        try:
            libs, scanned = self._fingerprint_apk(apk_path)
        except zipfile.BadZipFile:
            ida_kernwin.warning(f"'{apk_path}' is not a valid APK/ZIP file.")
            return
        except Exception as exc:                       # noqa: BLE001
            ida_kernwin.warning(f"Failed to read APK: {exc}")
            return

        if not libs:
            ida_kernwin.info(
                "No usable .so libraries with a code segment were found.\n"
                "For split APKs the library may live in a split_config.*.apk.")
            return
        _log("fingerprinted %d/%d candidate libraries", sum(len(v) for v in libs.values()), scanned)

        matches = self._match_and_rename(libs)
        _log("done: renamed segments for %d librar%s; refresh the Segments window",
             len(matches), "y" if len(matches) == 1 else "ies")

        if current_file and matches:
            self._offer_rebase(current_file, matches)

    # -- helpers ------------------------------------------------------------
    @staticmethod
    def _current_file_name():
        path = ida_nalt.get_input_file_path()
        return os.path.basename(path) if path else None

    @staticmethod
    def _pick_code_segment(elf):
        """First executable PT_LOAD, preferring a non-writable one (real .text)."""
        loads = [s.header for s in elf.iter_segments() if s.header.p_type == "PT_LOAD"]
        execs = [h for h in loads if h.p_flags & PF_X]
        if not execs:
            return None
        return next((h for h in execs if not h.p_flags & PF_W), execs[0])

    @staticmethod
    def _load_list(elf):
        return [Load(s.header.p_vaddr, bool(s.header.p_flags & PF_X), bool(s.header.p_flags & PF_W))
                for s in elf.iter_segments() if s.header.p_type == "PT_LOAD"]

    @staticmethod
    def _build_id(elf):
        """Return (hex_id, descriptor_vaddr, descriptor_len) or (None, 0, 0)."""
        for section in elf.iter_sections():
            if not isinstance(section, NoteSection) or not section["sh_addr"]:
                continue  # unmapped notes can't be read back from memory
            for note in section.iter_notes():
                if note["n_type"] != "NT_GNU_BUILD_ID":
                    continue
                name_pad = (note["n_namesz"] + 3) & ~3           # name is 4-byte aligned
                desc_off = note["n_offset"] + 12 + name_pad      # file offset of descriptor
                desc_vaddr = section["sh_addr"] + (desc_off - section["sh_offset"])
                return note["n_desc"], desc_vaddr, note["n_descsz"]
        return None, 0, 0

    @staticmethod
    def _segment_near(ea, tol=PAGE):
        """Return the debugger segment whose start is within `tol` of `ea`."""
        seg = ida_segment.getseg(ea)
        if seg and abs(seg.start_ea - ea) <= tol:
            return seg
        best = None
        for i in range(ida_segment.get_segm_qty()):
            s = ida_segment.getnseg(i)
            if s and abs(s.start_ea - ea) <= tol:
                if best is None or abs(s.start_ea - ea) < abs(best.start_ea - ea):
                    best = s
        return best

    # -- APK parsing --------------------------------------------------------
    def _fingerprint_apk(self, apk_path):
        """Return ({code_hash: [lib, ...]}, candidate_count) for the APK's .so files."""
        by_hash = defaultdict(list)
        scanned = 0
        with zipfile.ZipFile(apk_path, "r") as apk:
            for item in apk.infolist():
                if (item.is_dir()
                        or item.file_size < MIN_SO_SIZE
                        or not item.filename.startswith(LIB_PREFIX)
                        or not item.filename.lower().endswith(".so")):
                    continue
                scanned += 1
                try:
                    self._fingerprint_so(apk, item, by_hash)
                except ELFError:
                    _log("skip %s: not a valid ELF", item.filename)
                except (RuntimeError, NotImplementedError) as exc:
                    _log("skip %s: unreadable zip entry (%s)", item.filename, exc)
                except Exception as exc:               # noqa: BLE001 - keep scanning
                    _log("error on %s: %s", item.filename, exc)
        return by_hash, scanned

    def _fingerprint_so(self, apk, item, by_hash):
        raw = apk.read(item.filename)
        elf = ELFFile(io.BytesIO(raw))

        code = self._pick_code_segment(elf)
        if code is None:
            _log("skip %s: no executable segment", item.filename)
            return

        # Hash the same page the loader maps: floor the file offset to a page so
        # the window lines up with the page-aligned copy seen in memory.
        start = _page_floor(code.p_offset)
        if start + HASH_WINDOW > len(raw):
            _log("skip %s: code segment truncated in archive", item.filename)
            return
        code_hash = hashlib.sha256(raw[start:start + HASH_WINDOW]).hexdigest()

        build_id, bid_vaddr, bid_len = self._build_id(elf)
        by_hash[code_hash].append({
            "basename": os.path.basename(item.filename),
            "code_vaddr": _page_floor(code.p_vaddr),
            "loads": self._load_list(elf),
            "build_id": build_id,
            "bid_vaddr": bid_vaddr,
            "bid_len": bid_len,
        })

    # -- memory matching ----------------------------------------------------
    def _match_and_rename(self, by_hash):
        matched = []
        used_bases = set()
        for i in range(ida_segment.get_segm_qty()):
            seg = ida_segment.getnseg(i)
            if not seg or (seg.end_ea - seg.start_ea) < HASH_WINDOW or seg.start_ea in used_bases:
                continue
            window = ida_bytes.get_bytes(seg.start_ea, HASH_WINDOW)
            if not window:
                continue
            h = hashlib.sha256(window).hexdigest()
            lib = self._select_candidate(by_hash.get(h), seg.start_ea)
            if lib is None:
                continue
            used_bases.add(seg.start_ea)
            load_base = seg.start_ea - lib["code_vaddr"]
            confirmed = self._confirm_build_id(lib, load_base)
            _log("matched %s at load base %s%s", lib["basename"], _addr(load_base),
                 " (build-id confirmed)" if confirmed else "")
            code_seg = self._rename_library(lib, load_base)
            if code_seg is not None:
                matched.append(Match(code_seg, ida_segment.get_segm_name(code_seg),
                                     lib["basename"], load_base))
        return matched

    @staticmethod
    def _select_candidate(candidates, seg_start):
        """Pick one library among first-page-hash collisions (unique on hash today)."""
        if not candidates:
            return None
        if len(candidates) > 1:
            _log("warning: %d libraries share the first-page hash at %s; using %s",
                 len(candidates), _addr(seg_start), candidates[0]["basename"])
        return candidates[0]

    def _confirm_build_id(self, lib, load_base):
        """Best-effort: read the build-id note from memory and compare. Never raises."""
        if not lib["build_id"] or not lib["bid_len"]:
            return False
        try:
            data = ida_bytes.get_bytes(load_base + lib["bid_vaddr"], lib["bid_len"])
            return bool(data) and data.hex() == lib["build_id"]
        except Exception:                              # noqa: BLE001
            return False

    def _rename_library(self, lib, load_base):
        """Rename every debugger segment of `lib`. Return its code segment (or None)."""
        counters = defaultdict(int)
        code_seg = None
        for load in lib["loads"]:
            stem = f"{lib['basename']}_{'code' if load.is_exec else 'data'}"
            n = counters[stem]
            counters[stem] += 1
            new_name = stem if n == 0 else f"{stem}_{n}"

            target = self._segment_near(load_base + _page_floor(load.vaddr))
            if not target:
                continue
            old_name = ida_segment.get_segm_name(target)
            if old_name != new_name and ida_segment.set_segm_name(target, new_name):
                _log("  %s  %s -> %s", _addr(target.start_ea), old_name, new_name)
            if load.is_exec and code_seg is None:
                code_seg = target
        return code_seg

    # -- rebasing -----------------------------------------------------------
    def _offer_rebase(self, target_filename, matches):
        # Exact-basename match (not substring) so 'base.so' won't hit 'database.so'.
        hits = [m for m in matches if m.basename.lower() == target_filename.lower()]
        if not hits:
            return

        if len(hits) == 1:
            m = hits[0]
            delta = m.load_base - ida_nalt.get_imagebase()
            if delta == 0:
                ida_kernwin.info(f"'{target_filename}' is already based at {_addr(m.load_base)}.")
                return
            prompt = (f"Rebase '{target_filename}' to its live load base "
                      f"{_addr(m.load_base)} (delta {delta:#x})?")
            if ida_kernwin.ask_yn(ida_kernwin.ASKBTN_YES, prompt) == ida_kernwin.ASKBTN_YES:
                self._rebase(m.load_base, target_filename)
            return

        rows = [[m.name, _addr(m.load_base)] for m in hits]
        idx = _SegmentChooser(
            f"Multiple matches for '{target_filename}' - pick one to rebase", rows).Show(True)
        if idx < 0:
            _log("rebasing cancelled")
            return
        self._rebase(hits[idx].load_base, target_filename)

    def _rebase(self, new_base, filename):
        delta = new_base - ida_nalt.get_imagebase()
        if delta == 0:
            _log("'%s' already based at %s", filename, _addr(new_base))
            return
        try:
            _log("rebasing '%s' by delta %#x -> %s", filename, delta, _addr(new_base))
            # rebase_program and the MSF_* flags live in ida_segment, not ida_loader.
            code = ida_segment.rebase_program(delta, ida_segment.MSF_NOFIX)
            if code != ida_segment.MOVE_SEGM_OK:
                ida_kernwin.warning(f"Rebase failed (rebase_program returned {code}).")
            else:
                _log("rebased to %s", _addr(new_base))
        except Exception as exc:                       # noqa: BLE001
            ida_kernwin.warning(f"Error during rebasing: {exc}")


def PLUGIN_ENTRY():
    return ApkSegmentRenamerPlugin()