docs(debt): say that titles pass zeros to the VRR peg call - #14395

#!/usr/bin/env python3
"""Count the call sites of imported functions in title executables and read argument registers.

Input: a directory ROOT with one run folder per title, PPSA<id>_r*/, each holding
app.registry.json (import table written by the relinker; every import lists the addresses
of its call sites) and in/eboot.elf (the title's executable as ELF). Per title the registry
with the most imports that have call sites is used (ties: the newest file).
Only the title's own executable is read: libraries / extra module files, function pointers and
code that is only reached through data are not seen.

What is counted per title and function (NAME = import name, NID computed from the name):
  * call sites: the registry addresses whose first bytes are FF 25 (PLT stub, jmp [rip+got]).
  * direct calls: E8 rel32 instructions in the executable segments whose target is such a stub.
  * tail jumps: E9 rel32 instructions whose target is such a stub. They are listed and counted
    separately and are NOT direct calls.
  * calls to a jump stub: E8 rel32 whose target is one of those tail-jump instructions
    (a thunk that forwards to the import). Listed and counted separately, not direct calls.
  * non-stub registry call sites: registry addresses that are not a stub start (FF 15 call
    [rip+got], 48 8B / 4C 8B mov reg,[rip+got], anything else, or an address outside the file).
    Per title and function with address and instruction kind. They are not scanned for calls.
Matches are byte patterns in the executable segments, no control-flow recovery: an E8 inside an
instruction or data can only be a false hit if its target happens to be a stub (not checked).

Argument value (optional, --arg NAME=REG, REG one of rdi rsi rdx rcx r8 r9; the 32/16/8-bit
names are accepted and mean the same register): for every direct call the 64..48 bytes before it
are decoded with objdump (Intel syntax), once for every start offset in that range.
  * a decoding is used only if the call address is decoded as an instruction start; if no start
    offset gives that, the site is reported as 'desync';
  * going backwards from the call, the last instruction that writes REG decides: the first operand
    is REG or one of its sub-registers (mov, lea, movzx, pop, adc, rol, bswap, cmov, ...),
    an implicit writer (rep/movs/lods/stos/cmps/scas/ins/outs, cdq/cqo/cwd, mul/div/idiv and the
    one-operand imul for rdx, cpuid, rdtsc(p), rdmsr, rdpmc, xgetbv, syscall, cmpxchg8b/16b),
    xchg/xadd/cmpxchg that name REG, or a 'call' (clobbers all argument registers);
    cmp, test, push, bt and nop only read;
  * 'xor R,R' (R the 32- or 64-bit register) = 0; 'mov R,imm' (32- or 64-bit register) = imm;
    every other writer, and any write to a 16/8-bit part, = 'computed' (not constant);
  * a branch whose target lies between the last writer and the call, or an unconditional
    jmp / ret between them, = 'computed' (the writer may not be on the path to the call);
  * no writer in the window = 'unreadable' (not constant: set before the window or by the caller);
  * if the start offsets give different answers the site is 'ambiguous' (the answers are listed in
    the site line). With --reach-neutral an 'unreadable' answer of a shorter window is ignored when
    a longer window reads a value (it only means that the writer lies beyond that window's reach);
    all other disagreements stay 'ambiguous'.
Both 'computed' and 'unreadable' and 'ambiguous' mean "not a constant" in the sums; they are kept
apart so that a reader can see why. Window limits: a writer more than ~48 bytes back is not seen,
values that flow through memory or a callee-saved register are not followed, conditional moves
and branches taking different values on different paths are reported as 'computed'.
Nothing is executed; only counts, addresses and classifications are printed.

usage: anhang_zaehlen.py ROOT OBJDUMP NAME[,NAME...] OUT [--arg NAME=REG ...] [--titles FILE]
  OBJDUMP is the objdump command (GNU binutils, x86-64); 'wsl objdump' is accepted, paths are then
  converted to /mnt/<drive>/... . --titles FILE restricts the titles (PPSA<digits> found in FILE).
Output (OUT and stdout summary): one line per site "title | function | address | kind | value",
the non-stub list, and one summary line per function.
"""
import argparse
import bisect
import collections
import glob
import hashlib
import json
import os
import re
import struct
import subprocess
import sys
import tempfile

import numpy as np

SUFFIX = bytes([0x51, 0x8D, 0x64, 0xA6, 0x35, 0xDE, 0xD8, 0xC1, 0xE6, 0xB0, 0x39, 0xB1, 0xC3, 0xE5, 0x52, 0x30])
B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-"


def nid(name):
    d = hashlib.sha1(name.encode() + SUFFIX).digest()
    r = bytes(d[7 - i] for i in range(8))
    out = []
    for i in range(0, 6, 3):
        t = (r[i] << 16) | (r[i + 1] << 8) | r[i + 2]
        out += [B64[(t >> 18) & 63], B64[(t >> 12) & 63], B64[(t >> 6) & 63], B64[t & 63]]
    t = (r[6] << 16) | (r[7] << 8)
    return ''.join(out + [B64[(t >> 18) & 63], B64[(t >> 12) & 63], B64[(t >> 6) & 63]])


# argument registers: all names of the register, and the two names a constant may be written to
PARTS = {
    'rdi': ('rdi', 'edi', 'di', 'dil'),
    'rsi': ('rsi', 'esi', 'si', 'sil'),
    'rdx': ('rdx', 'edx', 'dx', 'dl', 'dh'),
    'rcx': ('rcx', 'ecx', 'cx', 'cl', 'ch'),
    'r8': ('r8', 'r8d', 'r8w', 'r8b'),
    'r9': ('r9', 'r9d', 'r9w', 'r9b'),
}
FULL = {'rdi': ('rdi', 'edi'), 'rsi': ('rsi', 'esi'), 'rdx': ('rdx', 'edx'), 'rcx': ('rcx', 'ecx'),
        'r8': ('r8', 'r8d'), 'r9': ('r9', 'r9d')}
ALIAS = {n: r for r, names in PARTS.items() for n in names}
PREFIX = {'lock', 'rep', 'repz', 'repe', 'repnz', 'repne', 'bnd', 'notrack', 'data16', 'cs', 'ds', 'es', 'fs', 'gs', 'addr32'}
READ_ONLY = {'cmp', 'test', 'push', 'bt', 'nop', 'ptest', 'ucomiss', 'ucomisd', 'comiss', 'comisd'}
# registers written by instructions that do not name them
STRING_OPS = {'movs': ('rsi', 'rdi'), 'cmps': ('rsi', 'rdi'), 'lods': ('rsi',), 'stos': ('rdi',),
              'scas': ('rdi',), 'ins': ('rdi',), 'outs': ('rsi',)}
IMPLICIT = {'rdx': ('cdq', 'cqo', 'cwd', 'mul', 'div', 'idiv', 'rdtsc', 'rdtscp', 'cpuid', 'rdmsr', 'rdpmc',
                    'xgetbv', 'cmpxchg8b', 'cmpxchg16b'),
            'rcx': ('loop', 'cpuid', 'rdtscp', 'xgetbv', 'syscall', 'sysenter', 'rdmsr', 'rdpmc'),
            'rdi': (), 'rsi': (), 'r8': (), 'r9': ()}
BRANCH = re.compile(r'^(j\w+|loop\w*|jrcxz|jecxz)$')
LINE = re.compile(r'\s+([0-9a-f]+):\s+(?:[0-9a-f]{2} )+\s*(.*)')


def split_ins(text):
    toks = text.split()
    pfx = []
    while toks and toks[0] in PREFIX:
        pfx.append(toks.pop(0))
    if not toks:
        return pfx, '', ''
    return pfx, toks[0], ' '.join(toks[1:])


def writes_reg(pfx, op, rest, reg):
    """True if the instruction (Intel syntax, prefixes split off) writes register reg."""
    names = PARTS[reg]
    anyname = re.compile(r'\b(?:' + '|'.join(names) + r')\b')
    if op == 'call' or op == 'syscall':
        return True
    base = op.rstrip('bwlq') if op.rstrip('bwlq') in STRING_OPS else op
    if base in STRING_OPS:
        if reg in STRING_OPS[base]:
            return True
        if reg == 'rcx' and any(p.startswith('rep') for p in pfx):
            return True
        return False
    if any(p.startswith('rep') for p in pfx) and reg == 'rcx':
        return True
    if op in IMPLICIT[reg]:
        return True
    if reg == 'rdx' and op in ('imul',) and ',' not in rest:
        return True
    if op in ('xchg', 'xadd', 'cmpxchg') and anyname.search(rest):
        return True
    if op in READ_ONLY:
        return False
    first = rest.split(',')[0].strip()
    return bool(anyname.fullmatch(first))


def const_of(op, rest, reg):
    """Value if the instruction sets the full register to a constant, else None."""
    f64, f32 = FULL[reg]
    ops = [x.strip() for x in rest.split(',')]
    if op == 'xor' and len(ops) == 2 and ops[0] == ops[1] and ops[0] in (f64, f32):
        return 0
    if op == 'mov' and len(ops) == 2 and ops[0] in (f64, f32) and re.fullmatch(r'-?(0x[0-9a-f]+|\d+)', ops[1]):
        v = int(ops[1], 0)
        return v & 0xffffffff if ops[0] == f32 else v & 0xffffffffffffffff
    return None


def classify(ins, callpos, reg):
    """ins: sorted list of (pos, text) of one decoded window. Returns a value string or None (desync)."""
    if not any(p == callpos for p, _ in ins):
        return None
    last = None
    lastpos = -1
    branches = []
    for pos, text in ins:
        if pos >= callpos:
            break
        pfx, op, rest = split_ins(text)
        if not op:
            continue
        if BRANCH.match(op) or op in ('jmp', 'ret', 'retn'):
            m = re.fullmatch(r'(?:0x)?([0-9a-f]+)', rest.split()[0]) if rest else None
            branches.append((pos, op, int(m.group(1), 16) if m else None))
            continue
        if writes_reg(pfx, op, rest, reg):
            last, lastpos = (op, rest), pos
    if last is None:
        return 'unreadable'
    for pos, op, tgt in branches:
        if pos > lastpos and op in ('jmp', 'ret', 'retn'):
            return 'computed'
        if tgt is not None and lastpos < tgt <= callpos:
            return 'computed'
    v = const_of(last[0], last[1], reg)
    if v is None:
        return 'computed'
    return str(v) if v < 10 else hex(v)


def to_tool_path(objdump, p):
    if objdump.split()[0] == 'wsl':
        p = os.path.abspath(p).replace('\\', '/')
        return '/mnt/' + p[0].lower() + p[2:] if len(p) > 1 and p[1] == ':' else p
    return p


def load_elf(path):
    data = open(path, 'rb').read()
    if data[:4] != b'\x7fELF':
        return None, None
    phoff = struct.unpack_from('<Q', data, 0x20)[0]
    phentsize, phnum = struct.unpack_from('<HH', data, 0x36)
    segs = []
    for i in range(phnum):
        ptype, flags, off, vaddr, paddr, filesz, memsz, align = struct.unpack_from('<IIQQQQQQ', data, phoff + i * phentsize)
        if ptype == 1:
            segs.append((vaddr, off, filesz, flags))
    return data, segs


def file_offset(segs, va):
    for vaddr, off, filesz, flags in segs:
        if vaddr <= va < vaddr + filesz:
            return off + (va - vaddr)
    return None


def scan(arr, segs, targets):
    """E8 (call) and E9 (jmp) rel32 instructions in executable segments whose target is in targets."""
    res = {'call': [], 'jmp': []}
    if not targets:
        return res
    tarr = np.array(sorted(targets), dtype=np.int64)
    for vaddr, off, filesz, flags in segs:
        if not flags & 1:
            continue
        seg = arr[off:off + filesz]
        for op, name in ((0xE8, 'call'), (0xE9, 'jmp')):
            idx = np.nonzero(seg[:-4] == op)[0]
            if idx.size == 0:
                continue
            rel = (seg[idx + 1].astype(np.int64) | (seg[idx + 2].astype(np.int64) << 8) |
                   (seg[idx + 3].astype(np.int64) << 16) | (seg[idx + 4].astype(np.int64) << 24))
            rel = np.where(rel >= 0x80000000, rel - 0x100000000, rel)
            hit = np.isin(vaddr + idx + 5 + rel, tarr)
            res[name] += [int(vaddr + i) for i in idx[hit]]
    return res


def site_kind(data, segs, va):
    fo = file_offset(segs, va)
    if fo is None:
        return 'outside file'
    b = data[fo:fo + 3]
    if b[:2] == b'\xff\x25':
        return 'FF 25 stub'
    if b[:2] == b'\xff\x15':
        return 'FF 15 call [rip+got]'
    if b[:2] in (b'\x48\x8b', b'\x4c\x8b'):
        return '%02X 8B mov reg,[rip+got]' % b[0]
    return 'other ' + b.hex(' ')


STARTS = range(64, 47, -1)
PAD = b'\xcc' * 16


def arg_values(objdump, data, segs, jobs, tmpdir, reach_neutral=False):
    """jobs: list of (va, reg). Returns {(va, reg): value string}. One objdump run for all windows."""
    blob = bytearray()
    wins = []  # (job index, back, start in blob, callpos in blob)
    for ji, (va, reg) in enumerate(jobs):
        fo = file_offset(segs, va)
        for back in STARTS:
            w0 = max(0, fo - back)
            s = len(blob)
            blob += data[w0:fo + 5]
            wins.append((ji, s, s + (fo - w0)))
            blob += PAD
    out = {}
    if not jobs:
        return out
    tmp = os.path.join(tmpdir, 'windows.bin')
    open(tmp, 'wb').write(bytes(blob))
    cmd = objdump.split() + ['-D', '-b', 'binary', '-mi386:x86-64', '-M', 'intel', to_tool_path(objdump, tmp)]
    txt = subprocess.run(cmd, capture_output=True, text=True).stdout
    pos_list, text_list = [], []
    for line in txt.split('\n'):
        m = LINE.match(line)
        if m and m.group(2).strip():
            pos_list.append(int(m.group(1), 16))
            text_list.append(m.group(2).strip())
    per = collections.defaultdict(collections.Counter)
    for ji, s, c in wins:
        lo = bisect.bisect_left(pos_list, s)
        hi = bisect.bisect_right(pos_list, c)
        v = classify(list(zip(pos_list[lo:hi], text_list[lo:hi])), c, ALIAS_REG[jobs[ji][1]])
        per[ji][v] += 1
    for ji, (va, reg) in enumerate(jobs):
        vals = {k: n for k, n in per[ji].items() if k is not None}
        if reach_neutral and len(vals) > 1 and 'unreadable' in vals:
            del vals['unreadable']  # a shorter window simply does not reach the writer
        if not vals:
            out[(va, reg)] = 'desync'
        elif len(vals) == 1:
            out[(va, reg)] = next(iter(vals))
        else:
            out[(va, reg)] = 'ambiguous'
            DETAIL[(va, reg)] = ', '.join('%s x%d' % (k, n) for k, n in sorted(vals.items(), key=lambda kv: -kv[1]))
    return out


ALIAS_REG = ALIAS
DETAIL = {}  # (address, register) -> answers of the start offsets for 'ambiguous' sites


def main():
    ap = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    ap.add_argument('root')
    ap.add_argument('objdump')
    ap.add_argument('names')
    ap.add_argument('out')
    ap.add_argument('--arg', action='append', default=[], help='NAME=REG (rdi rsi rdx rcx r8 r9)')
    ap.add_argument('--titles')
    ap.add_argument('--reach-neutral', action='store_true',
                    help="ignore 'unreadable' answers of short windows when a longer window reads a value "
                         "(default: any disagreement between start offsets is 'ambiguous')")
    a = ap.parse_args()
    names = a.names.split(',')
    nids = {n: nid(n) for n in names}
    argreg = {}
    for x in a.arg:
        n, r = x.split('=')
        argreg[n] = ALIAS[r.lower()]
    only = None
    if a.titles:
        only = set(re.findall(r'PPSA\d+', open(a.titles, encoding='utf-8', errors='replace').read()))
    rows = {}
    for f in glob.glob(os.path.join(a.root, 'PPSA*_r*', 'app.registry.json')):
        folder = os.path.dirname(f)
        m = re.match(r'(PPSA\d+)_r\w*$', os.path.basename(folder))
        if not m or (only is not None and m.group(1) not in only):
            continue
        try:
            d = json.load(open(f, encoding='utf-8'))
        except Exception:
            continue
        key = (sum(1 for e in d if e.get('callSites')), os.path.getmtime(f))
        if m.group(1) not in rows or key > rows[m.group(1)][0]:
            rows[m.group(1)] = (key, d, folder)
    lines = ['names: ' + ', '.join('%s=%s' % (n, nids[n]) for n in names),
             'titles with a registry: %d' % len(rows), '',
             'title | function | address | kind | value']
    nonstub_lines = []
    summ = {n: {'import': set(), 'calling': set(), 'call': 0, 'jmp': 0, 'thunk': 0, 'val': collections.Counter(),
                'nonstub': collections.Counter(), 'noelf': set()} for n in names}
    with tempfile.TemporaryDirectory() as tmpdir:
        for t in sorted(rows):
            key, d, folder = rows[t]
            elf = os.path.join(folder, 'in', 'eboot.elf')
            data = segs = arr = None
            for n in names:
                sites = sorted({int(x, 16) if isinstance(x, str) else int(x)
                                for e in d if str(e.get('nid', '')).split('#')[0] == nids[n]
                                for x in (e.get('callSites') or [])})
                if not sites:
                    continue
                summ[n]['import'].add(t)
                if data is None and os.path.exists(elf):
                    data, segs = load_elf(elf)
                    arr = np.frombuffer(data, dtype=np.uint8) if data is not None else None
                if data is None:
                    summ[n]['noelf'].add(t)
                    continue
                stubs = []
                for va in sites:
                    k = site_kind(data, segs, va)
                    if k == 'FF 25 stub':
                        stubs.append(va)
                    else:
                        summ[n]['nonstub'][k.split()[0] + k.split()[1] if k.startswith('FF') else 'other'] += 1
                        nonstub_lines.append('%s | %s | 0x%x | %s' % (t, n, va, k))
                r = scan(arr, segs, stubs)
                via = scan(arr, segs, r['jmp'])['call'] if r['jmp'] else []
                vals = {}
                if n in argreg and r['call']:
                    vals = arg_values(a.objdump, data, segs, [(va, argreg[n]) for va in r['call']], tmpdir, a.reach_neutral)
                for va in sorted(r['call']):
                    v = vals.get((va, argreg.get(n)), '')
                    if n in argreg:
                        summ[n]['val'][v] += 1
                    if v == 'ambiguous':
                        v += ' [' + DETAIL.get((va, argreg[n]), '') + ' of %d start offsets]' % len(STARTS)
                    lines.append('%s | %s | 0x%x | call | %s' % (t, n, va, v))
                for va in sorted(r['jmp']):
                    lines.append('%s | %s | 0x%x | tail jmp (not counted as call) |' % (t, n, va))
                for va in sorted(via):
                    lines.append('%s | %s | 0x%x | call to jump stub (not counted as direct call) |' % (t, n, va))
                if r['call']:
                    summ[n]['calling'].add(t)
                summ[n]['call'] += len(r['call'])
                summ[n]['jmp'] += len(r['jmp'])
                summ[n]['thunk'] += len(via)
    lines += ['', 'non-stub registry call sites: title | function | address | instruction kind']
    lines += nonstub_lines or ['(none)']
    lines += ['', 'summary: function | titles that import | titles with direct calls | direct calls | tail jmp | '
              'calls to jump stub | non-stub registry sites | titles without readable ELF | argument values']
    for n in names:
        s = summ[n]
        lines.append('%s | %d | %d | %d | %d | %d | %s | %d | %s' % (
            n, len(s['import']), len(s['calling']), s['call'], s['jmp'], s['thunk'],
            json.dumps(dict(s['nonstub']), sort_keys=True) if s['nonstub'] else '0',
            len(s['noelf']), json.dumps(dict(s['val']), sort_keys=True)))
    open(a.out, 'w', encoding='utf-8').write('\n'.join(lines) + '\n')
    print('\n'.join(lines[-(len(names) + 1):]))


if __name__ == '__main__':
    main()