#!/usr/bin/env python3
"""B9A read-only: find literal-pool xrefs to display strings in stock LK (lk.bin @0x8f600000).
usage: lk-xref.py lk.bin lk-arm.dis 'regex' ...
Prints string addr, every literal-pool word referencing it, and the function start
(nearest preceding push/stmdb) of the ldr that loads the pool word."""
import re, sys

BASE = 0x8f600000
lk = open(sys.argv[1], 'rb').read()
dis = open(sys.argv[2]).read().splitlines()

ins = {}      # addr -> (word, text)
order = []
for l in dis:
    m = re.match(r'([0-9a-f]{8}):\s+([0-9a-f]{8})\s+(.*)', l)
    if m:
        a = int(m.group(1), 16)
        ins[a] = (int(m.group(2), 16), m.group(3))
        order.append(a)

def func_start(a):
    while a >= BASE:
        t = ins.get(a)
        if t and (t[1].startswith('push') or t[1].startswith('stmdb\tsp!')):
            return a
        a -= 4
    return None

def loaders(pool):
    out = []
    for a, (w, t) in ins.items():
        m = re.search(r'ldr\w*\s+\w+, \[pc, #(-?\d+)\]', t)
        if m and a + 8 + int(m.group(1)) == pool:
            out.append(a)
    return out

for pat in sys.argv[3:]:
    for m in re.finditer(rb'[\x20-\x7e\t\n]{4,}', lk):
        s = m.group().decode()
        if not re.search(pat, s):
            continue
        sa = BASE + m.start()
        pools = [a for a, (w, _) in ins.items() if w == sa]
        print(f'STR 0x{sa:08x} {s!r}')
        for p in pools:
            for ld in loaders(p):
                fs = func_start(ld)
                print(f'   pool 0x{p:08x} <- ldr 0x{ld:08x} in func 0x{fs:08x}' if fs else f'   pool 0x{p:08x} <- 0x{ld:08x}')