switch-annotate-stacktrace.py

#!/usr/bin/env python3
from __future__ import annotations
from typing import *
from datetime import datetime
from functools import lru_cache
import os, os.path, re, sys

# This tool needs to emit the stacktrace with minimal background noise, so things fail passively and we usually exit(0).
#     ./switch-annotate-stacktrace.py $filename > output_emit.log
#     In this usage, stderr will be the only sign when something has gone wrong.
# "cli" goes to stderr (regardless of err/warn level)
def print_cli(*args, file=sys.stderr, **kwargs):
    return print(*args, file=file, **kwargs)
def print_err_cli(*args, prefix="ERROR:", **kwargs):
    return print_cli(prefix, *args, **kwargs)
def print_warn_cli(*args, prefix="WARNING:", **kwargs):
    return print_cli(prefix, *args, **kwargs)
# "emit" goes to stdout (the main tool output, TODO could also be a file)
def print_emit(*args, file=sys.stdout, **kwargs):
    return print(*args, file=file, **kwargs)
def print_warn_emit(*args, prefix="WARNING:", **kwargs):
    return print_emit(prefix, *args, **kwargs)


class ModuleLookup:
    # Each class may expose utilities for AMSLog+RyuLog annotation to use: offset_info, ...

    # from_* factories are favored over ctors here because they can fail and return a different type (None),
    # whereas failing in ctors means exceptions or creating a invalid/null state for every class.
    # It's easier to keep the config section brief and heavily nested when it's not concerned with this stuff.

    class LocalELF:
        @classmethod
        def from_filename(cls, addr2line: str, elf_filename: str) -> LocalELF:
            obj = cls()

            obj.addr2line = os.path.expanduser(addr2line)
            if not os.path.exists(obj.addr2line):
                print_err_cli(f"addr2line missing: {addr2line}")
                return None # nop

            obj.elf_filename = os.path.expanduser(elf_filename)
            if not os.path.exists(obj.elf_filename):
                print_err_cli(f"elf missing: {elf_filename}")
                return None # nop

            return obj

        def offset_info(self, offset: int) -> str:
            cmdline = f'"{self.addr2line}" -e "{self.elf_filename}" {hex(offset)}'
            output = os.popen(cmdline).read().strip()
            return output

    class SymbolFile:
        @classmethod
        def from_dt12345_totk_syms(cls, sym_filename: str) -> SymbolFile:
            obj = cls()
            obj.sym_filename = sym_filename
            obj.syms: dict[int, str] = {}
            try:
                with open(os.path.expanduser(sym_filename), "r") as symfile_fd:
                    for line in symfile_fd.readlines():
                        cols = line.split(",")
                        if cols[0] == "Address":
                            continue # header
                        offset = int(cols[0], 16)
                        name = ",".join(cols[1:]).strip('"\r\n') # HACK dumb 2nd col + unquote (use a csv parser u clown)
                        obj.syms[offset] = name
                obj.sorted_offsets = list(sorted(obj.syms.keys()))
                return obj

            except Exception as e:
                print_err_cli(f"can't open {sym_filename}: {e}")
            return None # nop

        @classmethod
        def from_watertoon_totk_syms(cls, sym_filename: str) -> SymbolFile:
            obj = cls()
            obj.sym_filename = sym_filename
            obj.syms: dict[int, str] = {}
            try:
                with open(os.path.expanduser(sym_filename), "r") as symfile_fd:
                    for line in symfile_fd.readlines():
                        offset, *name, _x, _y = map(lambda s: s.strip('"\r\n '), line.split(","))
                        if offset == "Address":
                            continue # header
                        offset = int(offset, 16)
                        obj.syms[offset] = ", ".join(name)
                obj.sorted_offsets = list(sorted(obj.syms.keys()))
                return obj

            except Exception as e:
                print_err_cli(f"can't open {sym_filename}: {e}")
            return None # nop

        def offset_info(self, offset: int) -> str:
            if name := self.syms.get(offset):
                return name # exact match

            prev_symbol_offset = max(x for x in self.sorted_offsets if x < offset)
            delta_offset = offset - prev_symbol_offset
            if delta_offset < 0x4000: # near match (XXX although 16KB is still not really near)
                name = self.syms[prev_symbol_offset]
                return f"{name} + {hex(delta_offset)}"

            return ""

    class LotuskitSyms:
        @classmethod
        def from_lotuskit_syms_folder(cls, lotuskit_syms_folder: str, totk_version: str = "TOTK_121", nso_module_pyname: str = "EXKING") -> LotuskitSyms:
            obj = cls()

            obj.lotuskit_syms_folder = lotuskit_syms_folder
            obj.totk_version = totk_version
            obj.nso_module_pyname = nso_module_pyname

            prevcwd = os.path.abspath(os.getcwd())
            syms_path = os.path.expanduser(obj.lotuskit_syms_folder)
            if not os.path.exists(syms_path):
                print_err_cli(f"lotuskit syms folder missing: {obj.lotuskit_syms_folder}")
                return None # nop

            try:
                os.chdir(syms_path)
                sys.path.append(syms_path)
                import build as LotuskitSymsBuild
                LotuskitSymsBuild.run_symbol_def_imports()
                [module.postprocess() for module in LotuskitSymsBuild.NSOModule.ALL_MODULES]
            finally:
                os.chdir(prevcwd)

            syms_pymod = getattr(LotuskitSymsBuild, nso_module_pyname, None)
            if not syms_pymod:
                print_err_cli(f"lotuskit syms do not expose {totk_version} {nso_module_pyname}: {obj.lotuskit_syms_folder}")
                return None # nop

            # build int-to-str map for this module+version
            obj.syms: dict[int, str] = {}
            for name, sym in syms_pymod.sym_map.items():
                if offset := sym.address.__getitem__(totk_version):
                    obj.syms[offset] = name
            obj.sorted_offsets = list(sorted(obj.syms.keys()))

            return obj

        def offset_info(self, offset: int) -> str:
            if name := self.syms.get(offset):
                return name # exact match

            prev_symbol_offset = max(x for x in self.sorted_offsets if x < offset)
            delta_offset = offset - prev_symbol_offset
            if delta_offset < 0x4000: # near match (XXX although 16KB is still not really near)
                name = self.syms[prev_symbol_offset]
                return f"{name} + {hex(delta_offset)}"

            return ""

    class GhidraProgram:
        Symbol = type("ghidra.program.model.symbol.Symbol")
        FmtLabel = type("list(tuple[formatted_name: str, offset: int, sym: Symbol, bestguess_score: int])")

        @classmethod
        def from_project_file(cls, gpr_filename: str, projectfs_nso_filename: str, nso_addr: int = 0x7100000000) -> GhidraProgram:
            obj = cls() # FIXME better handle gpr/program/pyghidra reuse, status, etc to allow opening multiple gprs*nsos across many GhidraProgram.from_* calls
            #obj.gpr_filename = gpr_filename # XXX member?
            #obj.projectfs_nso_filename = projectfs_nso_filename # XXX member?
            obj.nso_addr = nso_addr

            try:
                import pyghidra
                if pyghidra.started():
                    raise RuntimeError("pyghidra already started")
                pyghidra.start() # seems always stderr? nice
                if not pyghidra.started():
                    raise RuntimeError("pyghidra failed to start")

                #import ghidra
                #from ghidra.app.util.headless import HeadlessAnalyzer
                #from ghidra.program.flatapi import FlatProgramAPI
                from ghidra.base.project import GhidraProject
                #from ghidra.program.model.symbol import Symbol
                #from java.lang import String

                # TODO err handle file/folder existence upfront? ghidraspeak can be cryptic
                # TODO lol do we need to be closing stuff? its readonly

                # open gpr project file
                gpr_folder = os.path.expanduser(os.path.dirname(gpr_filename))
                gpr_basename = os.path.basename(gpr_filename)
                restore_project_tool_state = True
                obj.project = GhidraProject.openProject(gpr_folder, gpr_basename, restore_project_tool_state) # XXX should this stay a member?

                # open program using special nested folder name internal to the project
                *projectfs_folder, nsoname = projectfs_nso_filename.split("/")
                projectfs_folder = "/".join(projectfs_folder) or "/"
                program_readonly = True
                obj.program = obj.project.openProgram(projectfs_folder, nsoname, program_readonly) # XXX should this stay a member?

                # some components we may want # https://ghidra.re/ghidra_docs/api/ghidra/program/model/listing/Program.html
                # program.getDataTypeManager() # gdt stuff? https://ghidra.re/ghidra_docs/api/ghidra/program/model/data/ProgramBasedDataTypeManager.html
                # program.getMemory() # prob useless until i can connect ghidra to target https://ghidra.re/ghidra_docs/api/ghidra/program/model/mem/Memory.html
                # program.getEquateTable() # is this just like storage for labels or something different??? https://ghidra.re/ghidra_docs/api/ghidra/program/model/symbol/EquateTable.html
                # program.getListing() # https://ghidra.re/ghidra_docs/api/ghidra/program/model/listing/Listing.html
                obj.symtable = obj.program.getSymbolTable() # https://ghidra.re/ghidra_docs/api/ghidra/program/model/symbol/SymbolTable.html
                obj.addrspace = obj.program.getAddressFactory().getDefaultAddressSpace() # https://ghidra.re/ghidra_docs/api/ghidra/program/model/address/AddressSpace.html

            except Exception as e:
                print_err_cli(f"can't init ghidra {gpr_filename} -> {projectfs_nso_filename}: {e}")
                return None # nop

            return obj # ok

        def _get_ghidra_syms_at_offset(self, offset: int) -> list[Symbol]: # XXX prob not a real list, theres overloads that return like 3 types lol
            # https://ghidra.re/ghidra_docs/api/ghidra/program/model/symbol/Symbol.html
            return self.symtable.getSymbols(self.addrspace.getAddress(self.nso_addr + offset))

        def _bestguess_label_mapfunc(self, sym: Symbol) -> FmtLabel:
            # deghidrafies the label a bit by adding ns, stripping "Global" root ns, ignoring noisy labels, etc
            name = sym.getName()
            ns = sym.getParentNamespace()
            offset = sym.getAddress().getOffset() - self.nso_addr

            if ns.getName() == "Global":
                if name.startswith("LAB_"):
                    return [name, offset, sym,   0] # uninteresting, random jump point?
                if name.startswith("SUB_"):
                    return [name, offset, sym,   0] # uninteresting?
                if name.startswith("thunk_FUN_"):
                    return [name, offset, sym,  20] # idk
                if name.startswith("FUN_"):
                    return [name, offset, sym, 100] # function label is a decent default
                if name.startswith("DAT_"):
                    return [name, offset, sym,  10] # good when we're at an exact offset (not really while rewinding)?
                if name.startswith("PTR_DAT_"):
                    return [name, offset, sym,  10]
                if name.startswith("PTR_PTR_"):
                    return [name, offset, sym,  10]
                if name.startswith("PTR_"):
                    return [name, offset, sym,  10]
                # TODO any more ghidra ABC_ labels to ignore/downscore?
                #print_warn_cli(f"letting through {name}")

                return [name, offset, sym,     200] # return all remaining globals, likely to be meaningful names
            return [f"{ns}::{sym.getName()}", offset, sym, 1000] # namespaced identifier probably good

        def _get_rewind_to_bestguess_label(self, offset: int) -> Optional[FmtLabel]:
            assert(offset % 4 == 0), f"instructions must be aligned to 4: {hex(offset)}"
            offset_rewind = 0 # offset from the offset :)
            limit_offset_rewind = 0x400 # XXX idk

            while offset_rewind < limit_offset_rewind:
                labels = map(self._bestguess_label_mapfunc, self._get_ghidra_syms_at_offset(offset - offset_rewind))
                labels = sorted(list(labels), key=lambda l: l[3], reverse=True) # sort score desc
                if not labels or labels[0][3] <= 0:
                    offset_rewind += 4 # rewind one instruction
                else:
                    return labels[0] # XXX anything have multiple interesting labels?

            return None # nothing nearby worth reporting

        def offset_info(self, offset: int) -> str:
            label = self._get_rewind_to_bestguess_label(offset)
            if not label or label[3] <= 0:
                return ""

            guess_name, guess_offset, guess_sym, _score = label
            delta_offset = offset - guess_offset
            if delta_offset == 0:
                return guess_name # exact match

            return f"{guess_name} + {hex(delta_offset)}"


    # TODO many Alias+Heap pointers remain unannotated. Can we svcOutputDebugString the heapinfo and use it to annotate some of this?
    # TODO similarly emitting runtime hook/jit locations might be useful, but i don't have a log format for this sort of thing on ams yet


class RyuLog:
    #                                         app   branch?   version  datetime
    LOGFILE_PAT = re.compile("^(Ryujinx_|Ryubing_)(Canary_)?([0-9.]+_)([0-9_-]+).log$")

    def setup_lookup_impls(self):
        totk_version = self.get_game_version()
        if not totk_version or not totk_version.startswith("TOTK_"):
            print_warn_emit("No supported title detected!")
            self.MODULE_LOOKUP_IMPLS = {}
            return

        self.MODULE_LOOKUP_IMPLS = {
            "rtld": [],
            "main": [
                # simple file based lookup
                ModuleLookup.SymbolFile.from_watertoon_totk_syms("~/ghidra/sym/watertoon_totk_symbols_100_20231102.csv") if totk_version == "TOTK_100" else None,
                ModuleLookup.SymbolFile.from_dt12345_totk_syms("~/ghidra/sym/dt12345_exking_symbols_f6e4a4b.csv") if totk_version == "TOTK_121" else None,

                # call into various external tools
                ModuleLookup.LotuskitSyms.from_lotuskit_syms_folder("~/totk-lotuskit/syms", totk_version, "EXKING"),
                ModuleLookup.GhidraProgram.from_project_file("~/ghidra/projects/totk100/totk100.gpr", "main") if totk_version == "TOTK_100" else None,
                ModuleLookup.GhidraProgram.from_project_file("~/ghidra/projects/totk121/totk121.gpr", "main") if totk_version == "TOTK_121" else None,
            ],
            "subsdk0": [],
            "subsdk1": [ModuleLookup.LocalELF.from_filename("/opt/devkitpro/devkitA64/bin/aarch64-none-elf-addr2line", "~/totk-lotuskit/build/subsdk9.elf")],
            "subsdk2": [],
            "subsdk3": [],
            "subsdk4": [],
            "subsdk5": [],
            "subsdk6": [],
            "subsdk7": [],
            "subsdk8": [],
            "subsdk9": [],
            "sdk": [],
        }

    @classmethod
    def get_latest_log_filename(cls, logdir: str) -> str:
        latest_logfile: tuple[datetime, str] = (datetime.min, None)

        for logfile in os.listdir(logdir):
            match = cls.LOGFILE_PAT.match(logfile)
            if not match:
                continue

            logfile_date = datetime(*map(int, match.groups()[-1].replace("_", "-").split("-")))
            if logfile_date > latest_logfile[0]:
                latest_logfile = (logfile_date, logfile)

        if not latest_logfile[1]:
            #print_warn_cli(f"(No RyuLog file found in {logdir})")
            return None
        return os.path.join(logdir, latest_logfile[1])

    @classmethod
    def from_filename(cls, logfile: str) -> RyuLog:
        print_emit(f"Reading {os.path.basename(logfile)}:")
        with open(logfile, "r") as logfile_fd:
            logfile_lines = logfile_fd.readlines()
        obj = cls(logfile_lines)
        if game_version := obj.get_game_version():
            print_emit(f"(Detected GameVersion={game_version})")
        return obj

    def __init__(self, lines: list[str]):
        self.lines = lines

    @lru_cache
    def get_game_version(self):
        # TODO make this easier to generalize across titles, maybe use app id and ignore the name?
        version_pat = re.compile(r"^[0-9.:]+ \|I\| Loader Start: Application Loaded: The Legend of Zelda: Tears of the Kingdom v([0-9.]+) \[0100f2c0115b6000\] \[64-bit\]$")
        for i in range(len(self.lines)):
            if match := version_pat.match(self.lines[i]):
                return "TOTK_" + match.group(1).replace(".", "")
        return None

    def get_stacktrace_lines(self) -> Optional[list[str]]:
        stacktrace_beginpat = re.compile(r"^[0-9.:]+ \|I\| HLE.GuestThread.[0-9]+ Cpu PrintGuestStackTrace: Guest stack trace:$")
        for i in range(len(self.lines)):
            if stacktrace_beginpat.match(self.lines[i]):
                # assert only 1 stacktrace per logfile
                # assert logfile ends shortly after stacktrace
                return self.lines[i:]
        return None

    def print_annotated_stacktrace(self) -> None:
        stacktrace = self.get_stacktrace_lines()
        if not stacktrace:
            print_emit("(No stacktrace found)")
            return

        self.setup_lookup_impls()
        module_offset_pat = re.compile(r"([a-zA-Z0-9-.]+):(0x[0-9a-f]+)")
        redaction_userhome_pat = re.compile(r"/home/([a-zA-Z0-9._-]+)/")

        for line in stacktrace:
            line = module_offset_pat.sub(self.remap_module_offset, line)
            line = redaction_userhome_pat.sub(self.remap_redaction_userhome, line)
            print_emit(line, end="")

    def remap_module_offset(self, match: re.Match) -> str:
        modulename, hexoffset = match.groups()
        offset = int(hexoffset, 16)

        annotations = []
        for annotator in filter(bool, self.MODULE_LOOKUP_IMPLS.get(modulename, [])):
            if offset_info := annotator.offset_info(offset):
                annotations.append(offset_info)

        if not annotations:
            return match.group() # unaltered
        return f"{modulename}:{hexoffset} - " + ", ".join(annotations)

    def remap_redaction_userhome(self, match: re.Match) -> str:
        return "~/"


class AMSLog:
    # See ams::creport::CrashReport::SaveReport, SaveToFile
    # https://github.com/Atmosphere-NX/Atmosphere/blob/master/stratosphere/creport/source/creport_crash_report.cpp
    #                           timestamp          appid
    LOGFILE_PAT = re.compile("^([0-9]{11})_([0-9a-f]{16}).log")

    def setup_lookup_impls(self):
        totk_version = self.get_game_version()
        if not totk_version or not totk_version.startswith("TOTK_"):
            print_warn_emit("No supported title detected!")
            self.MODULE_LOOKUP_IMPLS = {}
            return

        self.MODULE_LOOKUP_IMPLS = {
            "nnrtld": [],
            "EX-King.nss": [
                # simple file based lookup
                ModuleLookup.SymbolFile.from_watertoon_totk_syms("~/ghidra/sym/watertoon_totk_symbols_100_20231102.csv") if totk_version == "TOTK_100" else None,
                ModuleLookup.SymbolFile.from_dt12345_totk_syms("~/ghidra/sym/dt12345_exking_symbols_f6e4a4b.csv") if totk_version == "TOTK_121" else None,

                # call into various external tools
                ModuleLookup.LotuskitSyms.from_lotuskit_syms_folder("~/totk-lotuskit/syms", totk_version, "EXKING"),
                ModuleLookup.GhidraProgram.from_project_file("~/ghidra/projects/totk100/totk100.gpr", "main") if totk_version == "TOTK_100" else None,
                ModuleLookup.GhidraProgram.from_project_file("~/ghidra/projects/totk121/totk121.gpr", "main") if totk_version == "TOTK_121" else None,
            ],
            "multimedia": [],
            "subsdk9": [ModuleLookup.LocalELF.from_filename("/opt/devkitpro/devkitA64/bin/aarch64-none-elf-addr2line", "~/totk-lotuskit/build/subsdk9.elf")],
            "nnSdk": [],
        }

    @classmethod
    def get_latest_log_filename(cls, logdir: str) -> str:
        latest_logfile: tuple[int, str] = (0, None)

        for logfile in os.listdir(logdir):
            match = cls.LOGFILE_PAT.match(logfile)
            if not match:
                continue

            logfile_date = int(match.groups()[0])
            if logfile_date > latest_logfile[0]:
                latest_logfile = (logfile_date, logfile)

        if not latest_logfile[1]:
            #print_warn_cli(f"(No AMSLog file found in {logdir})")
            return None
        return os.path.join(logdir, latest_logfile[1])

    @classmethod
    def from_filename(cls, logfile: str) -> AMSLog:
        print_emit(f"Reading {os.path.basename(logfile)}:")
        with open(logfile, "r") as logfile_fd:
            logfile_lines = logfile_fd.readlines()
        obj = cls(logfile_lines)
        if game_version := obj.get_game_version():
            print_emit(f"(Detected GameVersion={game_version})")
        return obj

    def __init__(self, lines: list[str]):
        self.lines = lines

    @lru_cache
    def get_game_version(self):
        game_pat = re.compile(r"^[ \t]*Program ID:[ \t]*([0-9a-f]{16})$")
        game_prefix = ""
        for i in range(len(self.lines)):
            if match := game_pat.match(self.lines[i]):
                if match.group(1) == "0100f2c0115b6000":
                    game_prefix = "TOTK"
                break

        version = ""
        if game_prefix == "TOTK":
            dying_pat = re.compile(r"^[ \t]*Dying Message:")
            for i in range(len(self.lines)):
                if match := dying_pat.match(self.lines[i]):
                    dieline = self.lines[i+2].strip(" \t") # get to offset 0x40 (each line is 0x20)
                    version_b36 = int(dieline[0:2], 16) + 0x100 * int(dieline[2:4], 16)
                    version_major, rem = divmod(version_b36, 36*36)
                    version_minor, version_micro = divmod(rem, 36)
                    version = f"{version_major}{version_minor}{version_micro}"
                    break

        if game_prefix and version:
            return f"{game_prefix}_{version}"
        return None

    def print_annotated_stacktrace(self) -> None:
        self.setup_lookup_impls()
        module_offset_pat = re.compile(r"\(([a-zA-Z0-9-.]+) \+ (0x[0-9a-f]+)\)")
        redaction_userhome_pat = re.compile(r"/home/([a-zA-Z0-9._-]+)/")

        for line in self.lines:
            line = module_offset_pat.sub(self.remap_module_offset, line)
            line = redaction_userhome_pat.sub(self.remap_redaction_userhome, line)
            print_emit(line, end="")

    def remap_module_offset(self, match: re.Match) -> str:
        modulename, hexoffset = match.groups()
        offset = int(hexoffset, 16)

        annotations = []
        for annotator in filter(bool, self.MODULE_LOOKUP_IMPLS.get(modulename, [])):
            if offset_info := annotator.offset_info(offset):
                annotations.append(offset_info)

        if not annotations:
            return match.group() # unaltered
        return f"({modulename} + {hexoffset}) (" + ", ".join(annotations) + ")"

    def remap_redaction_userhome(self, match: re.Match) -> str:
        return "~/"


def main():
    if len(sys.argv) > 1:
        for arg in sys.argv[1:]:
            if RyuLog.LOGFILE_PAT.match(os.path.basename(arg)):
                logfile = RyuLog.from_filename(arg)
                logfile.print_annotated_stacktrace()
            elif AMSLog.LOGFILE_PAT.match(os.path.basename(arg)):
                logfile = AMSLog.from_filename(arg)
                logfile.print_annotated_stacktrace()
            else:
                # TODO set log type based on contents or argv opt (not just filename)
                # TODO accept folders here -> get_latest_log_filename
                print_warn_cli(f"Unrecognized logfile arg {arg}")

    else: # find latest log from cwd
        logdir: str = os.path.abspath(os.getcwd())
        if logfile := AMSLog.get_latest_log_filename(logdir):
            # TODO argv opt for log type to disambiguate when both types exist
            AMSLog.from_filename(logfile).print_annotated_stacktrace()
        elif logfile := RyuLog.get_latest_log_filename(logdir):
            RyuLog.from_filename(logfile).print_annotated_stacktrace()
        elif not logdir.endswith("Logs"):
            print_cli("Usage: place this script in Logs folder and run in wsl/linux dev environment, or pass a log filename")
            return sys.exit(420)

if __name__ == "__main__":
    main()