#!/usr/bin/env python3
"""rebase v4: 直白版 - 原地改写 dump 的重定位槽, 无嵌套函数.
用法: python3 rebase_v4.py <original.so> <dump_full.so> <runtime_base> <out.so>
"""
from __future__ import annotations
import struct, sys

def main():
    original, dump, base_s, out = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
    base = int(base_s, 16)
    od = open(original, 'rb').read()
    dd = bytearray(open(dump, 'rb').read())

    # 原文件段表
    e_shoff = struct.unpack_from('<Q', od, 0x28)[0]
    e_shentsize = struct.unpack_from('<H', od, 0x3A)[0]
    e_shnum = struct.unpack_from('<H', od, 0x3C)[0]
    e_shstrndx = struct.unpack_from('<H', od, 0x3E)[0]
    shdrs = []
    for i in range(e_shnum):
        sh = struct.unpack_from('<IIQQQQIIQQ', od, e_shoff + i * e_shentsize)
        shdrs.append([sh[0], sh[3], sh[4], sh[5]])  # name, vaddr, offset, size
    shstr = shdrs[e_shstrndx]
    shstrtab = od[shstr[2]:shstr[2] + shstr[3]]
    def nm(idx):
        e = shstrtab.find(b'\0', idx)
        return shstrtab[idx:e].decode(errors='replace')

    reb, zer, low, oob = 0, 0, 0, 0
    for s in shdrs:
        n = nm(s[0])
        if n not in ('.rela.dyn', '.rela.plt', '.init_array', '.fini_array'):
            continue
        if n in ('.rela.dyn', '.rela.plt'):
            for i in range(0, s[3], 24):
                r_off, r_info, _ = struct.unpack_from('<QQq', od, s[2] + i)
                rt = r_info & 0xffffffff
                if r_off + 8 > len(dd):
                    oob += 1
                    continue
                v = struct.unpack_from('<Q', dd, r_off)[0]
                if rt == 1027 or rt == 257:  # RELATIVE / ABS64
                    if v == 0 or v < 0x4000000000:
                        low += 1
                        continue
                    if base <= v < base + 0x600000:
                        struct.pack_into('<Q', dd, r_off, v - base)
                        reb += 1
                    else:
                        struct.pack_into('<Q', dd, r_off, 0)
                        zer += 1
                elif rt == 1025 or rt == 1026:  # GLOB_DAT / JUMP_SLOT
                    struct.pack_into('<Q', dd, r_off, 0)
                    zer += 1
        else:  # init/fini arrays
            for i in range(0, s[3], 8):
                off = s[1] + i  # vaddr 轴
                if off + 8 > len(dd):
                    oob += 1
                    continue
                v = struct.unpack_from('<Q', dd, off)[0]
                if v == 0 or v < 0x4000000000:
                    low += 1
                    continue
                if base <= v < base + 0x600000:
                    struct.pack_into('<Q', dd, off, v - base)
                    reb += 1
                else:
                    struct.pack_into('<Q', dd, off, 0)
                    zer += 1

    # 段名表修复: dump 的 shstrtab 为空, 复制原文件段头表 + shstrtab (布局一致)
    import struct as _s
    e_shoff = _s.unpack_from('<Q', od, 0x28)[0]
    e_shentsize = _s.unpack_from('<H', od, 0x3A)[0]
    e_shnum = _s.unpack_from('<H', od, 0x3C)[0]
    e_shstrndx = _s.unpack_from('<H', od, 0x3E)[0]
    # 原文件 shstrtab 段内容
    sh = _s.unpack_from('<IIQQQQIIQQ', od, e_shoff + e_shstrndx * e_shentsize)
    shstr_off, shstr_size = sh[4], sh[5]
    if e_shoff + e_shnum * e_shentsize <= len(dd) and shstr_off + shstr_size <= len(dd):
        dd[e_shoff:e_shoff + e_shnum * e_shentsize] = od[e_shoff:e_shoff + e_shnum * e_shentsize]
        dd[shstr_off:shstr_off + shstr_size] = od[shstr_off:shstr_off + shstr_size]
        # 重定位文件为全 vaddr 轴(offset==vaddr): 逐段修正 sh_offset = sh_addr
        for si in range(e_shnum):
            basee = e_shoff + si * e_shentsize
            sh = list(_s.unpack_from('<IIQQQQIIQQ', dd, basee))
            sh[4] = sh[3]  # offset := addr
            _s.pack_into('<IIQQQQIIQQ', dd, basee, *sh)
        print(f"[*] section header table + shstrtab copied & offset:=addr (shoff={e_shoff:#x})")
    else:
        print("[!] WARN: section table copy skipped (out of bounds)")

    open(out, 'wb').write(dd)
    print(f"[*] rebase={reb} zero={zer} low={low} oob={oob}")
    # 验证
    import subprocess
    v1 = struct.unpack_from('<Q', dd, 0x3b21c0)[0]
    v2 = struct.unpack_from('<Q', dd, 0x3b9b68)[0] if 0x3b9b68 + 8 <= len(dd) else -1
    print(f"[*] check init_array[0]={v1:#x} (.got@0x3b9b68)={v2:#x}")

if __name__ == '__main__':
    main()