import re
import idc
import idaapi
import idautils
import ida_bytes
import ida_auto
SEG = ".text"
WIN = 20
APPLY = True # 先 False 看打印，核对 target 与 ADRL 符号一致再改 True
NOWRITE = {"STR", "STUR", "CMP", "CMN", "TST", "B", "BL", "BR", "BLR",
"RET", "NOP", "CBZ", "CBNZ", "TBZ", "TBNZ", "SVC", "BRK"}
CTRL = {"B", "BL", "BR", "BLR", "RET", "CBZ", "CBNZ", "TBZ", "TBNZ"}
# ---------- 基础工具 ----------
def mnem(ea):
return (idc.print_insn_mnem(ea) or "").upper()
def rn(t):
m = re.match(r"^[WX](\d+)$", (t or "").strip())
return int(m.group(1)) if m else None
def oreg(ea, i):
return rn(idc.print_operand(ea, i))
def otxt(ea, i):
return (idc.print_operand(ea, i) or "").strip()
def oimm(ea, i):
try:
return idc.get_operand_value(ea, i)
except Exception:
return None
def iscode(ea):
return ea != idaapi.BADADDR and idc.is_code(ida_bytes.get_full_flags(ea))
def xrefto(ea):
"""只取真正的地址引用，过滤 fl_F 顺序流（上次 B 到自己就是这里漏过滤）。"""
nxt = idc.next_head(ea, ea + 8)
for x in idautils.XrefsFrom(ea, 0):
if x.type == idaapi.fl_F or x.to == nxt:
continue
if iscode(x.to):
return x.to
return None
def window(br_ea, lo):
"""win[0]=BR，下标越大地址越小，最多 WIN 条。"""
win = [br_ea]
cur = br_ea
for _ in range(WIN):
p = idc.prev_head(cur, lo)
if p == idaapi.BADADDR or p < lo or not iscode(p):
break
win.append(p)
cur = p
return win
def find_def(win, r, start):
if r is None:
return None
for i in range(start, len(win)):
ea = win[i]
m = mnem(ea)
if m in CTRL:
return None
if m in NOWRITE:
continue
if oreg(ea, 0) == r:
return ea
return None
def adrl_parts(win, def_ea):
"""def_ea 是 ADRL 伪指令或 ADRP+ADD 后半时，返回 (target, [ea...])。"""
m = mnem(def_ea)
if m in ("ADRL", "ADRP"):
t = xrefto(def_ea)
return (t, [def_ea]) if t else (None, None)
if m == "ADD" and oreg(def_ea, 0) == oreg(def_ea, 1):
i = win.index(def_ea)
if i + 1 < len(win):
p = win[i + 1]
if mnem(p) == "ADRP" and oreg(p, 0) == oreg(def_ea, 0):
t = xrefto(def_ea) or xrefto(p)
if t:
return t, [def_ea, p]
return None, None
def trace_const(win, r, start):
"""追常量：MOV #imm，或一层 STR/LDR 同槽往返，支持 XZR。返回 (value, [eas])。"""
d = find_def(win, r, start)
if d is None:
return None
if mnem(d) == "MOV":
v = oimm(d, 1)
if v is not None:
return v, [d]
if mnem(d) in ("LDR", "LDUR"):
mem = otxt(d, 1)
di = win.index(d)
for j in range(di + 1, len(win)):
ea = win[j]
if mnem(ea) in ("STR", "STUR") and otxt(ea, 1) == mem:
if otxt(ea, 0) in ("XZR", "WZR"):
return 0, [d, ea]
sd = find_def(win, oreg(ea, 0), j + 1)
if sd and mnem(sd) == "MOV":
v = oimm(sd, 1)
if v is not None:
return v, [d, ea, sd]
return None
# ---------- 类型 2: MUL-BR（1 * addr） ----------
def match_mul_br(win, d):
mdi = win.index(d)
a, b = oreg(d, 1), oreg(d, 2)
for s in (a, b):
sd = find_def(win, s, mdi + 1)
if not sd:
continue
t, _ = adrl_parts(win, sd)
if t is None:
continue
other = b if s == a else a
tc = trace_const(win, other, mdi + 1)
if tc and tc[0] == 1:
return {"type": "MUL-BR", "target": t, "extra": "1*addr"}
return {"type": "MUL-BR?", "target": t, "extra": "乘数未证明为1"}
return None
# ---------- 类型 1: MADD-BR（(a&0)*b + (a^0)） ----------
def match_madd_br(win, d):
if mnem(d) != "MADD":
return None
try:
if idc.get_operand_type(d, 3) != idaapi.o_reg or otxt(d, 3) in ("XZR", "WZR"):
return None
except Exception:
return None
mdi = win.index(d)
m1, m2, add = oreg(d, 1), oreg(d, 2), oreg(d, 3)
de = find_def(win, add, mdi + 1) # 加数必须是 EOR(addr, zero)
if not de or mnem(de) != "EOR":
return None
e1, e2 = oreg(de, 1), oreg(de, 2)
ide = win.index(de)
for mc in (m1, m2): # 某个乘数是 AND(addr, zero)
da = find_def(win, mc, mdi + 1)
if not da or mnem(da) != "AND":
continue
idi = win.index(da)
a1, a2 = oreg(da, 1), oreg(da, 2)
for ar, zr in ((a1, a2), (a2, a1)):
if ar is None or zr is None or zr not in (e1, e2):
continue
er = e1 if e2 == zr else e2
if er != ar:
continue
sd = find_def(win, ar, max(idi, ide) + 1)
if not sd:
continue
t, _ = adrl_parts(win, sd)
if t is None:
continue
zc = trace_const(win, zr, max(idi, ide) + 1)
if zc and zc[0] == 0:
return {"type": "MADD-BR", "target": t, "extra": "(a&0)*b+(a^0)"}
return None
# ---------- ADD-BR 族公用：拆 base(ADRL) + off ----------
def split_base_off(win, d):
mdi = win.index(d)
for i in (1, 2):
r = oreg(d, i)
if r is None:
continue
sd = find_def(win, r, mdi + 1)
if not sd:
continue
t, _ = adrl_parts(win, sd)
if t is not None:
return t, oreg(d, 2 if i == 1 else 1), mdi
return None
# ---------- 类型 3: AND-ADD-BR（n(n+1)&1 == 0） ----------
def match_and_add(win, d):
sp = split_base_off(win, d)
if not sp:
return None
t, off, mdi = sp
da = find_def(win, off, mdi + 1)
if not da or mnem(da) != "AND" or oimm(da, 2) != 1:
return None
dm = find_def(win, oreg(da, 1), win.index(da) + 1)
if not dm or mnem(dm) != "MUL":
return None
m1, m2 = oreg(dm, 1), oreg(dm, 2)
mi = win.index(dm)
n_r, d1 = None, None
for s in (m1, m2): # 找 ADD Xs+1, Xn, #1
o = m2 if s == m1 else m1
dd = find_def(win, s, mi + 1)
if dd and mnem(dd) == "ADD" and oreg(dd, 1) == o and oimm(dd, 2) == 1:
n_r, d1 = o, dd
break
if n_r is None:
return None
extra = "n(n+1)&1"
tc = trace_const(win, n_r, win.index(d1) + 1)
if tc:
n = tc[0]
extra = "n=%d, n(n+1)&1=%d" % (n, (n * (n + 1)) & 1)
return {"type": "AND-ADD-BR", "target": t, "extra": extra}
# ---------- 类型 4: UMULH-ADD-BR（n(n+1)(n+2)%3 == 0） ----------
def has_magic3(win):
for ea in win:
if "AAAAAAAAAAAAAAAB" in otxt(ea, 1).upper():
return True
if mnem(ea) in ("MOV", "MOVZ") and oimm(ea, 1) == 0xAAAB:
return True
return False
def match_umulh_add(win, d):
sp = split_base_off(win, d)
if not sp:
return None
t, off, mdi = sp
ds = find_def(win, off, mdi + 1) # off = SUBS(prod, q*3)
if not ds or mnem(ds) != "SUBS":
return None
prod, q3 = oreg(ds, 1), oreg(ds, 2)
si = win.index(ds)
dq = find_def(win, q3, si + 1) # q*3 = ADD q, q, q, LSL#1
if not dq or mnem(dq) != "ADD":
return None
op2 = otxt(dq, 2).replace(" ", "").upper()
if "LSL#1" not in op2 or rn(op2.split(",")[0]) != oreg(dq, 1):
return None
dl = find_def(win, oreg(dq, 1), win.index(dq) + 1) # q = LSR(u, #1)
if not dl or mnem(dl) != "LSR" or oimm(dl, 2) != 1:
return None
du = find_def(win, oreg(dl, 1), win.index(dl) + 1) # u = UMULH(prod, magic)
if not du or mnem(du) != "UMULH":
return None
if not has_magic3(win):
return None
extra = "n(n+1)(n+2)%3"
dp = find_def(win, prod, si + 1) # 可选：追 n 验证
if dp and mnem(dp) == "MUL":
for s in (oreg(dp, 1), oreg(dp, 2)):
d2 = find_def(win, s, win.index(dp) + 1)
if d2 and mnem(d2) == "ADD" and oimm(d2, 2) == 2:
tc = trace_const(win, oreg(d2, 1), win.index(d2) + 1)
if tc:
n = tc[0]
extra = "n=%d, n(n+1)(n+2)%%3=%d" % (n, (n * (n + 1) * (n + 2)) % 3)
return {"type": "UMULH-ADD-BR", "target": t, "extra": extra}
# ---------- 总分类 ----------
def classify(br_ea, lo):
win = window(br_ea, lo)
br_r = oreg(br_ea, 0)
if br_r is None:
return None
d = find_def(win, br_r, 1)
if d is None:
return None
m = mnem(d)
if m in ("MUL", "MADD"):
r = match_madd_br(win, d)
if r:
return r
try:
if idc.get_operand_type(d, 3) == idaapi.o_reg and otxt(d, 3) not in ("XZR", "WZR"):
return None # 真乘加，不是花
except Exception:
pass
return match_mul_br(win, d)
if m == "ADD":
return match_and_add(win, d) or match_umulh_add(win, d)
return None
# ---------- patch：只把 BR 改成 B target ----------
def patch_one(ea, info):
t = info["target"]
delta = t - ea
if delta & 3:
print("[SKIP ] BR=0x%X -> 0x%X 未 4 对齐" % (ea, t))
return False
off = delta >> 2
if not (-(1 << 25) <= off < (1 << 25)):
print("[SKIP ] BR=0x%X -> 0x%X 超出 ±128MB" % (ea, t))
return False
word = 0x14000000 | (off & 0x3FFFFFF)
bts = " ".join("%02X" % c for c in word.to_bytes(4, "little"))
print("[PATCH] %-13s BR=0x%X -> B 0x%X | off=%d word=0x%08X bytes=%s | %s%s"
% (info["type"], ea, t, off, word, bts, info["extra"],
"" if APPLY else " (dry-run)"))
if not APPLY:
return True
ida_bytes.patch_dword(ea, word)
idc.create_insn(ea)
try:
idaapi.add_cref(ea, t, idaapi.fl_JN)
except Exception:
pass
return True
def scan():
ida_auto.auto_wait()
res = {}
for seg in idautils.Segments():
if idc.get_segm_name(seg) != SEG:
continue
for ea in idautils.Heads(seg, idc.get_segm_end(seg)):
if mnem(ea) != "BR" or oreg(ea, 0) is None:
continue
fn = idaapi.get_func(ea)
info = classify(ea, fn.start_ea if fn else seg)
if info:
res[ea] = info
print("=" * 78)
from collections import Counter
cnt = Counter(v["type"] for v in res.values())
print("[SCAN] 共 %d 个：%s" % (len(res), dict(cnt)))
ok = 0
for ea in sorted(res):
if patch_one(ea, res[ea]):
ok += 1
if APPLY:
ida_auto.auto_wait()
idc.refresh_idaview_anyway()
print("[DONE] %d/%d APPLY=%s" % (ok, len(res), APPLY))
print("=" * 78)
# ---------- 修复以前被错误 patch 的 B（可选，正常流程用不到） ----------
def repair():
ida_auto.auto_wait()
fixed = 0
for seg in idautils.Segments():
if idc.get_segm_name(seg) != SEG:
continue
for ea in idautils.Heads(seg, idc.get_segm_end(seg)):
if mnem(ea) != "B":
continue
fn = idaapi.get_func(ea)
win = window(ea, fn.start_ea if fn else seg)
if len(win) < 2 or mnem(win[1]) not in ("MUL", "MADD", "ADD"):
continue
t = None
for w in win[1:]:
tt, _ = adrl_parts(win, w)
if tt:
t = tt
break
if t is None:
continue
cur = idc.get_operand_value(ea, 0)
if cur != t:
delta = t - ea
if delta & 3 or not (-(1 << 25) <= (delta >> 2) < (1 << 25)):
continue
word = 0x14000000 | ((delta >> 2) & 0x3FFFFFF)
print("[FIX] B@0x%X: 0x%X -> 0x%X word=0x%08X" % (ea, cur, t, word))
ida_bytes.patch_dword(ea, word)
idc.create_insn(ea)
fixed += 1
print("[REPAIR] 修复 %d 处" % fixed)
scan()
# repair() # 需要修复历史误 patch 时取消注释
|
