"""Read-only ARM64 inspection with optional on-device method-map annotations."""
import argparse
import struct
import re
from pathlib import Path
from capstone import Cs, CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN

p = argparse.ArgumentParser(description=__doc__)
p.add_argument('library', type=Path)
p.add_argument('address', type=lambda s: int(s, 0))
p.add_argument('size', type=lambda s: int(s, 0))
p.add_argument('--map', type=Path)
p.add_argument('--metadata', type=Path, help='Exact Quest metadata for literal annotations')
p.add_argument('--august', action='store_true', help='Use August 31 metadata header fields')
a = p.parse_args()
b = a.library.read_bytes()
ph = struct.unpack_from('<Q', b, 32)[0]
es, n = struct.unpack_from('<HH', b, 54)
mapping = {}
segments = [struct.unpack_from('<IIQQQQQQ', b, ph + i * es) for i in range(n)]
def qword(rva):
    for s in segments:
        if s[0] == 1 and s[3] <= rva and rva + 8 <= s[3] + s[5]:
            return struct.unpack_from('<Q', b, s[2] + rva - s[3])[0]
    raise ValueError('Unmapped pointer')
meta = a.metadata.read_bytes() if a.metadata else None
previous = None
if a.map:
    for line in a.map.read_text(encoding='utf-8-sig').splitlines():
        parts = line.split('\t')
        if len(parts) == 6:
            mapping[int(parts[0], 16)] = '.'.join(parts[2:5]) + '/' + parts[5]
for i in range(n):
    typ, flags, off, va, _, size, _, _ = struct.unpack_from('<IIQQQQQQ', b, ph + i * es)
    if typ == 1 and flags & 1 and va <= a.address < va + size:
        pos = off + a.address - va
        for inst in Cs(CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN).disasm(b[pos:pos+a.size], a.address):
            note = mapping.get(inst.address, '')
            if meta and previous and previous.mnemonic == 'adrp' and inst.mnemonic == 'ldr':
                m = re.fullmatch(r'(x\d+), \[(x\d+), #(0x[0-9a-f]+)\]', inst.op_str)
                reg, _, page = previous.op_str.partition(', #')
                if m and m[2] == reg:
                    try:
                        token = qword(qword(int(page, 0) + int(m[3], 0)))
                        if token >> 29 == 5:
                            index = (token & 0x1fffffff) >> 1
                            table = struct.unpack_from('<I', meta, 0x74 if a.august else 0x44)[0]
                            data = struct.unpack_from('<I', meta, 0xa0 if a.august else 0x28)[0]
                            length, offset = struct.unpack_from('<II', meta, table + index * 8)
                            if length < 512:
                                note += ' literal=' + repr(meta[data+offset:data+offset+length].decode('utf-8'))
                    except (ValueError, struct.error, UnicodeError):
                        pass
            if inst.mnemonic in ('b', 'bl') and inst.op_str.startswith('#0x'):
                note += ' -> ' + mapping.get(int(inst.op_str[1:], 16), '')
            print(f'{inst.address:x}: {inst.mnemonic:8} {inst.op_str:36} {note}')
            previous = inst
        break
else:
    raise ValueError('RVA is not in a file-backed executable segment')