#!/usr/bin/env python3
"""重定位 dump (v3): 用 ORIGINAL .so 的段表定位 .rela.dyn/.init_array (布局一致),
   对 DUMP 内同偏移的槽做"运行期地址 -> 链接期偏移"回退 / 导入槽清零.

用法: python3 rebase_dump.py <original.so> <dump.so> <runtime_base> <out.so>
"""
from __future__ import annotations
import struct, sys

R_AARCH64_ABS64 = 257
R_AARCH64_GLOB_DAT = 1025
R_AARCH64_JUMP_SLOT = 1026
R_AARCH64_RELATIVE = 1027

def sections(data):
    e_shoff = struct.unpack_from('<Q', data, 0x28)[0]
    e_shentsize = struct.unpack_from('<H', data, 0x3A)[0]
    e_shnum = struct.unpack_from('<H', data, 0x3C)[0]
    e_shstrndx = struct.unpack_from('<H', data, 0x3E)[0]
    shdrs = []
    for i in range(e_shnum):
        sh = struct.unpack_from('<IIQQQQIIQQ', data, e_shoff + i * e_shentsize)
        shdrs.append(dict(name=sh[0], type=sh[1], addr=sh[3], offset=sh[4], size=sh[5], entsize=sh[9]))
    shstr = shdrs[e_shstrndx]
    shstrtab = data[shstr['offset']:shstr['offset'] + shstr['size']]
    def sname(idx):
        end = shstrtab.find(b'\0', idx)
        return shstrtab[idx:end].decode(errors='replace')
    for s in shdrs:
        s['name'] = sname(s['name'])
    return shdrs

def main():
    original, dump, base_s, out = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
    base = int(base_s, 16)
    odata = open(original, 'rb').read()
    ddata = bytearray(open(dump, 'rb').read())
    count = {'rebase': 0, 'zero': 0}

    # dump 文件偏移 == vaddr (运行期镜像与 vaddr 同轴, 已实证)
    def rebase_ptr(vaddr):
        off = vaddr
        if off + 8 > len(ddata):
            return
        v = struct.unpack_from('<Q', ddata, off)[0]
        if off == 0x3b23f0: print('DBG rebase_ptr sees %#x' % v)
        if v == 0 or v < 0x700000000000:
            return
        if base <= v < base + 0x600000:
            struct.pack_into('<Q', ddata, off, v - base)
            count['rebase'] += 1
        else:
            struct.pack_into('<Q', ddata, off, 0)
            count['zero'] += 1

    for s in sections(odata):
        if s['name'] in ('.rela.dyn', '.rela.plt'):
            for i in range(0, s['size'], 24):
                r_offset, r_info, _ = struct.unpack_from('<QQq', odata, s['offset'] + i)
                rtype = r_info & 0xffffffff
                if rtype in (R_AARCH64_RELATIVE, R_AARCH64_ABS64):
                    rebase_ptr(r_offset)
                elif rtype in (R_AARCH64_GLOB_DAT, R_AARCH64_JUMP_SLOT):
                    if r_offset + 8 <= len(ddata):
                        struct.pack_into('<Q', ddata, r_offset, 0)
                        count['zero'] += 1
        elif s['name'] in ('.init_array', '.fini_array'):
            for i in range(0, s['size'], 8):
                rebase_ptr(s['addr'] + i)

    open(out, 'wb').write(ddata)
    print(f"[*] rebased {count['rebase']} slots, zeroed {count['zero']} -> {out}")

if __name__ == '__main__':
    main()