#!/usr/bin/env python3
"""通用解密合并 (v3): 按 PT_LOAD 映射 vaddr<->文件偏移, 把 FULL dump 的解密内存
合入 ORIGINAL .so 的正确文件位置; 数据段内陈旧运行期指针清掉.

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

def phdrs(data):
    e_phoff = struct.unpack_from('<Q', data, 0x20)[0]
    e_phentsize = struct.unpack_from('<H', data, 0x36)[0]
    e_phnum = struct.unpack_from('<H', data, 0x38)[0]
    out = []
    for i in range(e_phnum):
        p = struct.unpack_from('<IIQQQQQQ', data, e_phoff + i * e_phentsize)
        out.append(dict(type=p[0], flags=p[1], off=p[2], vaddr=p[3], filesz=p[4], memsz=p[5]))
    return out

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]
    sh = []
    for i in range(e_shnum):
        s = struct.unpack_from('<IIQQQQIIQQ', data, e_shoff + i * e_shentsize)
        sh.append(dict(name=s[0], type=s[1], addr=s[3], offset=s[4], size=s[5]))
    shstr = sh[e_shstrndx]
    tab = data[shstr['offset']:shstr['offset'] + shstr['size']]
    def nm(idx):
        e = tab.find(b'\0', idx)
        return tab[idx:e].decode(errors='replace')
    for s in sh:
        s['name'] = nm(s['name'])
    return sh

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 = bytearray(open(original, 'rb').read())
    orig_raw = open(original, 'rb').read()
    dd = open(dump, 'rb').read()
    print(f"[*] orig={len(od):#x} dump={len(dd):#x} base=0x{base:x}")

    ph = phdrs(orig_raw)
    secs = sections(orig_raw)
    dyn_ranges, zero_ranges = [], []
    for s in secs:
        if s['name'] == '.dynamic':
            dyn_ranges.append((s['offset'], s['offset'] + s['size']))
        if s['name'] in ('.init_array', '.fini_array'):
            zero_ranges.append((s['offset'], s['offset'] + s['size']))

    # 按段复制: 数据段文件偏移 <- dump(运行时 vaddr=文件段 vaddr) 内容
    DATA_SECS = ('.rodata', '.data.rel.ro', '.got', '.got.plt', '.data')
    n_copied = 0
    for s in secs:
        if s['type'] == 8 or s['name'] not in DATA_SECS:
            continue
        vaddr, foff, sz = s['addr'], s['offset'], s['size']
        if vaddr + sz > len(dd) or foff + sz > len(od):
            print(f"[!] skip {s['name']}: out of bounds vaddr={vaddr:#x} off={foff:#x} sz={sz:#x}")
            continue
        if dd[vaddr:vaddr+sz] != od[foff:foff+sz]:
            od[foff:foff+sz] = dd[vaddr:vaddr+sz]
            n_copied += sz
    print(f"[*] merged {n_copied:#x} bytes via section mapping")

    for a, b in dyn_ranges:
        od[a:b] = orig_raw[a:b]
        print(f"[*] kept original .dynamic [{a:#x}:{b:#x})")
    for a, b in zero_ranges:
        od[a:b] = b'\x00' * (b - a)
        print(f"[*] zeroed [{a:#x}:{b:#x})")

    # 清理只作用于重定位目标槽: vaddr->文件偏移, 越界运行期指针清零/模块内回退
    R_AARCH64_RELATIVE, R_AARCH64_ABS64 = 1027, 257
    R_AARCH64_GLOB_DAT, R_AARCH64_JUMP_SLOT = 1025, 1026
    def v2foff(vaddr):
        for s_ in secs:
            if s_['type'] == 8 or s_['size'] == 0:
                continue
            if s_['addr'] <= vaddr < s_['addr'] + s_['size']:
                return s_['offset'] + (vaddr - s_['addr'])
        return None
    n_zero = n_rebase = 0
    for s_ in secs:
        if s_['name'] not in ('.rela.dyn', '.rela.plt') or s_['type'] == 8:
            continue
        for i in range(0, s_['size'], 24):
            r_off, r_info, _ = struct.unpack_from('<QQq', orig_raw, s_['offset'] + i)
            rt = r_info & 0xffffffff
            if rt not in (R_AARCH64_RELATIVE, R_AARCH64_ABS64, R_AARCH64_GLOB_DAT, R_AARCH64_JUMP_SLOT):
                continue
            foff = v2foff(r_off)
            if foff is None or foff + 8 > len(od):
                continue
            v = struct.unpack_from('<Q', od, foff)[0]
            if v < 0x4000000000:
                continue
            if base <= v < base + 0x600000:
                struct.pack_into('<Q', od, foff, v - base)
                n_rebase += 1
            else:
                struct.pack_into('<Q', od, foff, 0)
                n_zero += 1
    print(f"[*] reloc slots: rebased={n_rebase} zeroed={n_zero}")

    open(out, 'wb').write(od)
    print(f"[*] -> {out}")

if __name__ == '__main__':
    main()