# Stack walker v5: dbghelp StackWalk64 on every thread of a live process.
# Usage: python stackwalk.py --bridge | --probe [secs]
import ctypes, ctypes.wintypes as w, struct, subprocess, sys, os, time

sys.stdout.reconfigure(encoding='utf-8', errors='replace')

MODE = sys.argv[1] if len(sys.argv) > 1 else "--bridge"
SECS = int(sys.argv[2]) if len(sys.argv) > 2 else 15
SHM_NAME = "SonicForge_DAW_IPC_VERIFY"
SIZE = 32768
BRIDGE = r"C:\Users\locpham\SonicForgeStudio\native_bridge\build\Release\daw_vst_bridge.exe"
PROBE = r"C:\Users\locpham\SonicForgeStudio\native_bridge\build\Release\gui_probe.exe"
NEXUS = r"C:\Program Files\Common Files\VST3\Nexus.vst3"
LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                   "stack_sample.log" if MODE == "--bridge" else "stack_probe.log")

k32 = ctypes.windll.kernel32
k32.CreateFileMappingA.argtypes = [ctypes.c_void_p, ctypes.c_void_p, w.DWORD, w.DWORD, w.DWORD, ctypes.c_char_p]
k32.CreateFileMappingA.restype = ctypes.c_void_p
k32.MapViewOfFile.argtypes = [ctypes.c_void_p, w.DWORD, w.DWORD, w.DWORD, ctypes.c_size_t]
k32.MapViewOfFile.restype = ctypes.c_void_p
INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value
hMap = k32.CreateFileMappingA(INVALID_HANDLE_VALUE, None, 0x04, 0, SIZE, SHM_NAME.encode())
ptr = k32.MapViewOfFile(hMap, 0xF001F, 0, 0, 0)

def wait_log(pattern, timeout):
    end = time.time() + timeout
    while time.time() < end:
        try:
            with open(LOG, 'r', encoding='utf-8', errors='replace') as f:
                if pattern in f.read():
                    return True
        except FileNotFoundError:
            pass
        time.sleep(0.2)
    return False

if MODE == "--bridge":
    extra = dict(os.environ)
    for k in list(extra):
        if k.startswith("SF_"):
            del extra[k]
    extra.update({"SF_SHM_NAME": SHM_NAME, "SF_SAMPLE_RATE": "48000",
                  "SF_BLOCK_SIZE": "256", "SF_ONCE": "1", "SF_AUTOGUI": "0"})
    proc = subprocess.Popen([BRIDGE, "--shm", SHM_NAME], env=extra,
                            stdout=open(LOG, 'wb'), stderr=subprocess.STDOUT)
    marker = "isPlatformTypeSupported=0"
    wait_after = 3
else:
    proc = subprocess.Popen([PROBE, NEXUS, "bridge_exact11", str(SECS)],
                            stdout=open(LOG, 'wb'), stderr=subprocess.STDOUT)
    marker = "openGUI: attached=0"
    wait_after = 2

print("pid", proc.pid, "mode", MODE)
if not wait_log(marker, 120):
    print("TIMEOUT waiting for", marker)
    sys.exit(1)
time.sleep(wait_after)
pid = proc.pid

PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_VM_READ = 0x0010
hProc = k32.OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, False, pid)
if not hProc:
    print("OpenProcess failed", ctypes.get_last_error())
    sys.exit(1)

# module map
psapi = ctypes.windll.psapi
class MODULEINFO(ctypes.Structure):
    _fields_ = [("lpBaseOfDll", ctypes.c_void_p), ("SizeOfImage", w.DWORD),
                ("pad", w.DWORD), ("EntryPoint", ctypes.c_void_p)]
psapi.EnumProcessModulesEx.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), w.DWORD, ctypes.POINTER(w.DWORD), w.DWORD]
psapi.GetModuleBaseNameA.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p, w.DWORD]
psapi.GetModuleInformation.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(MODULEINFO), w.DWORD]
mods = []
buf = (ctypes.c_void_p * 1024)()
needed = w.DWORD(0)
if psapi.EnumProcessModulesEx(hProc, buf, ctypes.sizeof(buf), ctypes.byref(needed), 3):
    n = needed.value // ctypes.sizeof(ctypes.c_void_p)
    for i in range(min(n, 1024)):
        base = buf[i]
        name = ctypes.create_string_buffer(260)
        psapi.GetModuleBaseNameA(hProc, base, name, 260)
        info = MODULEINFO()
        psapi.GetModuleInformation(hProc, base, ctypes.byref(info), ctypes.sizeof(info))
        mods.append((base, info.SizeOfImage, name.value.decode('utf-8', 'replace')))
mods.sort(key=lambda m: m[0])

def modname(addr):
    if addr == 0:
        return "null"
    for b, s, n in mods:
        if b <= addr < b + s:
            return "%s+%x" % (n, addr - b)
    return "??%x" % addr

# dbghelp stack walk
dbg = ctypes.WinDLL("dbghelp")
dbg.SymInitializeW.argtypes = [ctypes.c_void_p, ctypes.c_wchar_p, w.BOOL]
dbg.SymInitializeW.restype = w.BOOL
dbg.SymSetOptions.argtypes = [w.DWORD]
dbg.SymSetOptions.restype = w.DWORD
dbg.StackWalk64.argtypes = [w.DWORD, ctypes.c_void_p, ctypes.c_void_p,
                            ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p,
                            ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, w.DWORD]
dbg.StackWalk64.restype = w.BOOL

class STACKFRAME64(ctypes.Structure):
    class ADDR(ctypes.Structure):
        _fields_ = [("Offset", ctypes.c_uint64), ("Segment", w.DWORD),
                    ("Mode", w.DWORD)]
    _fields_ = [("AddrPC", ADDR), ("AddrReturn", ADDR), ("AddrFrame", ADDR),
                ("AddrStack", ADDR), ("AddrBStore", ADDR), ("FuncTableEntry", ctypes.c_void_p),
                ("Params", ctypes.c_uint64 * 4), ("Far", w.BOOL),
                ("Virtual", w.BOOL), ("Reserved", ctypes.c_uint64 * 3),
                ("KdHelp", ctypes.c_uint64 * 2)]

class CONTEXT64(ctypes.Structure):
    _fields_ = [
        ("P1Home", ctypes.c_uint64), ("P2Home", ctypes.c_uint64),
        ("P3Home", ctypes.c_uint64), ("P4Home", ctypes.c_uint64),
        ("P5Home", ctypes.c_uint64), ("P6Home", ctypes.c_uint64),
        ("ContextFlags", w.DWORD), ("MxCsr", w.DWORD),
        ("SegCs", w.WORD), ("SegDs", w.WORD), ("SegEs", w.WORD),
        ("SegFs", w.WORD), ("SegGs", w.WORD), ("SegSs", w.WORD),
        ("EFlags", w.DWORD),
        ("Dr0", ctypes.c_uint64), ("Dr1", ctypes.c_uint64),
        ("Dr2", ctypes.c_uint64), ("Dr3", ctypes.c_uint64),
        ("Dr6", ctypes.c_uint64), ("Dr7", ctypes.c_uint64),
        ("Rax", ctypes.c_uint64), ("Rcx", ctypes.c_uint64),
        ("Rdx", ctypes.c_uint64), ("Rbx", ctypes.c_uint64),
        ("Rsp", ctypes.c_uint64), ("Rbp", ctypes.c_uint64),
        ("Rsi", ctypes.c_uint64), ("Rdi", ctypes.c_uint64),
        ("R8", ctypes.c_uint64), ("R9", ctypes.c_uint64),
        ("R10", ctypes.c_uint64), ("R11", ctypes.c_uint64),
        ("R12", ctypes.c_uint64), ("R13", ctypes.c_uint64),
        ("R14", ctypes.c_uint64), ("R15", ctypes.c_uint64),
        ("Rip", ctypes.c_uint64),
    ]

ReadProcessMemory_cb = ctypes.CFUNCTYPE(w.BOOL, ctypes.c_void_p, ctypes.c_uint64,
                                        ctypes.c_void_p, w.DWORD, ctypes.POINTER(w.DWORD))

@ReadProcessMemory_cb
def read_mem(hProc, addr, buf, size, nread):
    return k32.ReadProcessMemory(hProc, ctypes.c_void_p(addr), buf, size, nread)

class THREADENTRY32(ctypes.Structure):
    _fields_ = [("dwSize", w.DWORD), ("cntUsage", w.DWORD), ("th32ThreadID", w.DWORD),
                ("th32OwnerProcessID", w.DWORD), ("tpBasePri", ctypes.c_long),
                ("tpDeltaPri", ctypes.c_long), ("dwFlags", w.DWORD)]

TH32CS_SNAPTHREAD = 0x4
THREAD_SUSPEND_RESUME = 0x0002
THREAD_GET_CONTEXT = 0x0008
THREAD_QUERY_INFORMATION = 0x0040
CONTEXT_CTRL_INT_SEG = 0x100001 | 0x2 | 0x20
IMAGE_FILE_MACHINE_AMD64 = 0x8664

dbg.SymInitializeW(hProc, None, True)
dbg.SymSetOptions(0x2)  # SYMOPT_DEFERRED_LOADS

snap = k32.CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
te = THREADENTRY32(); te.dwSize = ctypes.sizeof(THREADENTRY32)
threads = []
ok = k32.Thread32First(snap, ctypes.byref(te))
while ok:
    if te.th32OwnerProcessID == pid:
        threads.append(te.th32ThreadID)
    ok = k32.Thread32Next(snap, ctypes.byref(te))

print("threads(%d): %s" % (len(threads), threads))
for tid in threads:
    hT = k32.OpenThread(THREAD_SUSPEND_RESUME | THREAD_GET_CONTEXT | THREAD_QUERY_INFORMATION, False, tid)
    if not hT:
        continue
    k32.SuspendThread(hT)
    ctx = CONTEXT64(); ctx.ContextFlags = CONTEXT_CTRL_INT_SEG
    r = k32.GetThreadContext(hT, ctypes.byref(ctx))
    if not r:
        k32.ResumeThread(hT); k32.CloseHandle(hT)
        print("tid %d GetThreadContext fail" % tid)
        continue
    sf = STACKFRAME64()
    sf.AddrPC.Offset = ctx.Rip
    sf.AddrPC.Mode = 0
    sf.AddrReturn.Offset = ctx.Rsp  # will be set by walker
    sf.AddrFrame.Offset = ctx.Rbp
    sf.AddrStack.Offset = ctx.Rsp
    frames = []
    for i in range(48):
        if not dbg.StackWalk64(IMAGE_FILE_MACHINE_AMD64, hProc, hT, ctypes.byref(sf),
                               ctypes.byref(ctx), read_mem,
                               dbg.SymFunctionTableAccess64, dbg.SymGetModuleBase64,
                               None, 0):
            break
        if sf.AddrPC.Offset == 0:
            break
        frames.append(modname(sf.AddrPC.Offset))
        if sf.AddrPC.Offset == sf.AddrReturn.Offset:
            break
        if i > 0 and frames[-1] == frames[-2]:
            break
    k32.ResumeThread(hT)
    k32.CloseHandle(hT)
    print("=== tid %d (rip %s) ===" % (tid, modname(ctx.Rip)))
    for f in frames:
        print("   ", f)

print("--- log tail ---")
try:
    with open(LOG, 'rb') as f:
        f.seek(max(0, os.path.getsize(LOG) - 1200))
        print(f.read().decode('utf-8', 'replace'))
except FileNotFoundError:
    pass