diff --git a/AGENTS.md b/AGENTS.md index a77cb75..93ab89f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,15 +4,18 @@ This repository contains ToCode, a Python-only binary exporter. ToCode takes one ## Scope -- Keep the project focused on the exporter CLI and Python library modules. -- Do not add web UI, server, container wrapper, shortcut, installer, or background service behavior. +- Keep the project focused on the exporter CLI, the live MCP server (`tocode mcp`), and Python library modules. +- Do not add web UI, network servers, container wrappers, shortcuts, installers, or background services. The one long-running mode is the stdio MCP server, started by an MCP client and ending with it. - Do not add chat, report generation, or secondary analysis stages. - Export one project tree from one binary or IDA database: raw decompiler output, assembly, summaries, section data, JSON metadata, optional IDA database, and optional scanner-friendly C. - Generated export trees must include their own `AGENTS.md` for agents analyzing that exported binary. ## Project Layout -- `src/tocode/cli.py`: command-line entry point for `tocode`. +- `src/tocode/cli.py`: command-line entry point for `tocode` (subcommands `mcp`, `sync`, `skill` are dispatched before the flat parser). +- `src/tocode/live/`: the live layer. `jsonrpc.py` stdio MCP transport (standard library), `server.py` `tocode mcp` entry, `tools.py` the tool table (schemas, annotations, timeouts), `sessions.py` worker lifecycle (one spawned process per session, cap, idle close, timeouts, crash recovery), `targets.py` what `open` resolves to, `worker.py` the per-session process and query/eval operations, `edits.py` edit tools per backend and saving, `sync.py` incremental export updates and the `export` operation, `setup.py` `tocode sync`, `tocode skill install`, and MCP client registration. .NET sessions: `dotnet_live.py` (`DotnetLiveSession`, picked by `worker.open_session` for `backend == "dotnet"`), `dotnet_edits.py` (the edit log and renames), `dotnet_sync.py` (rewrites the touched types' `.cs`/`.il` and the documents). +- `src/tocode/syncstate.py`: what an export keeps for `sync` (`.tocode/sync/manifest.json` with the file layout, `.tocode/sync/renders/` with each function's render). +- `src/tocode/skill/`: the agent skill (`SKILL.md`, `references/*.md`), shipped as package data. - `src/tocode/__init__.py`: package version and `export_from_binaryview()`, the in-UI Binary Ninja library entry. - `src/tocode/analysis.py`: backend-neutral binary inventory and call graph normalization. - `src/tocode/backends/`: IDA Domain, radare2, angr, and Binary Ninja (`binja_headless.py` over the in-process `binaryninja` API, `binja_mcp.py` over Binary Ninja's MCP server, `binja.py` over RPyc or an injected UI view) session adapters, plus the ASC/droidasc APK backend (`asc.py`: APK set discovery, DEX inventory, worker-side class decompilation). @@ -27,6 +30,7 @@ This repository contains ToCode, a Python-only binary exporter. ToCode takes one - `src/tocode/parallel.py`: worker-count selection. - `src/tocode/schema.py`: dataclasses shared across the exporter. - `tests/`: unit tests for algorithms, CLI helpers, and export tree generation. +- `eval/`: the blind source-recovery benchmark (`recover.py`, `corpus/`): a headless agent with only ToCode recovers a stripped program's source and a judge scores it. Not shipped in the package; run it after changing the live tools or the skill. ## Export Contract @@ -54,6 +58,7 @@ The CLI accepts a regular binary file or IDA database and writes a project direc - `tocode.log` - generated export `AGENTS.md` - generated export `CLAUDE.md` +- `.tocode/sync/manifest.json` and `.tocode/sync/renders/*.json` (what `sync` updates from) When the IDA backend is used, the export also contains: @@ -120,6 +125,44 @@ Framework/runtime assemblies and native runtime libraries inside bundles and pac - Pin new dependencies exactly (`==`) and capture them in `uv.lock`. - Keep dependencies minimal and tied to exporting a project. +## Live server (MCP) and sync + +- `tocode mcp` speaks MCP over stdio. Nothing but protocol may reach fd 1: + `server.protect_stdout` moves it to stderr before backends load, and workers + log to `/.tocode/live.log`. +- Every session is a spawned worker process (`sessions.SessionManager`): a + backend crash, OOM kill, or a call past its timeout stops only that worker, + which reopens on the next call (the reply says unsaved edits were lost). + Keep backend calls out of the supervisor process. +- Database edits go through `live/edits.py` only, so each edit records the + functions it makes stale (`LiveSession.dirty`); add a backend method there + rather than calling backend APIs from tools. angr is query-only. +- `sync` never renumbers or rewrites a source file whose content did not + change; existing functions keep their file, new ones join a caller's file + (else `discovered.c`). It rewrites metadata through the exporter's own + writers (`_write_metadata(publish=False)`), and refreshes folder links when + the export belongs to a folder export. +- `CheckpointStore.complete()` keeps the renders for sync and removes only the + checkpoint's own files; `.tocode/` also holds sync state and live logs. +- .NET live: the database is the edit log `/.tocode/dotnet/edits.json` + (edits addressed by assembly and metadata token), replayed onto the + original input on open; `patched/.dll` are renamed copies written + from it (tokens preserved). Renames change dnlib objects and update only the + affected inventory records and name-index entries; patched copies are + written when a query decompiles an edited assembly; the inventory reloads + once per sync. Local renames and comments are text edits applied when C# is + rendered (neither lives in metadata). Keep edits O(what they touch): use the + reverse indexes (`member_owner`, `field_users`, the word index) rather than + scanning all methods or files. +- `src/tocode/skill/references/tools.md` is generated from `live/tools.py` + (`python -m tocode.live.tools > src/tocode/skill/references/tools.md`); + `tests/test_live_setup.py` fails when it is stale. +- The skill, tool descriptions, and server instructions must never use names + from `eval/corpus` (the benchmark would no longer be blind); a test checks. +- Tests drive the real server through `tests/live_client.py` against the + on-disk fake Binary Ninja (`tests/fake_binja/`, which supports the edit + API); the `export` fixture in `tests/conftest.py` builds a fake export. + ## Binary Ninja backend Three transports (`--binja-transport auto|headless|mcp|rpyc`); `auto` tries them diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index ecd8a59..c074ec2 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -62,7 +62,10 @@ Write tests so they do not depend on the machine running them: no network, no installed backend or runtime, and no particular amount of free memory. `tests/dotnet_fixtures.py` builds synthetic .NET files, and `tests/test_apk_export.py` / `tests/test_dotnet_export.py` show how to drive a -whole export with fake sessions. +whole export with fake sessions. The live server tests (`tests/test_live_*.py`) +run `tocode mcp` in a subprocess through `tests/live_client.py` against the +fake Binary Ninja API; `python tests/live_client.py 'decompile={"target":"main"}'` +drives a real server by hand. ## Checking a change on real inputs diff --git a/README.md b/README.md index 9de6e14..037ef91 100644 --- a/README.md +++ b/README.md @@ -109,10 +109,81 @@ Four backends are supported, selected with `--backend` (default `auto`, which pr Other disassemblers may be added in the future. +# Using + +ToCode supports Windows, Linux, and macOS with Python 3.10 or newer. + +Install it from PyPI: + +```bash +pip install tocode-cli # or: uv tool install tocode-cli +pip install "tocode-cli[angr]" # also install the angr fallback backend +``` + +The command is `tocode`. You still need a backend (IDA, radare2, angr, or Binary Ninja); see [Supported backends](https://github.com/buzzer-re/ToCode#supported-backends). + +To install from a clone of the repository instead, on Windows PowerShell: + +```powershell +powershell -ExecutionPolicy Bypass -File .\install.ps1 +``` + +On Linux or macOS: + +```bash +bash ./install.sh +``` + +Add `--all` or `--full` (PowerShell: `-Full`) to also install the angr fallback backend, so ToCode works out-of-the-box even without IDA or radare2: + +```bash +bash ./install.sh --all +``` + +Manual setup (requires [uv](https://docs.astral.sh/uv/)): + +```bash +git clone https://github.com/buzzer-re/ToCode +cd ToCode +uv sync --locked +uv tool install --force --editable . +``` + +## Example + +``` +tocode firmwareX.bin -o firmwareX_decompiled/ +cd firmwareX_decompiled/ +codex + +# Inside your agent shell, type your goals, e.g.: "Give me a brief overview of the boot process of this firmware." +``` + +### From an ongoing RE work + +``` +tocode firmwareX.bin.i64 -o firmwareX_decompiled/ +``` + +Interrupted exports save progress automatically. Rerun the same command to resume, or add `--restart` to start over. + +# Live analysis for agents (MCP) + +`tocode mcp` keeps the decompiler open so an agent can query it, record what it learns (names, types, structs, comments, or scripts run in the database), and `sync` the export: changed functions and the code that uses them are re-rendered, and only those files change. It works on .NET exports too: renames are kept as an edit log, and the renamed assemblies are written to `patched/`. + +```bash +tocode mcp --install claude # or --install codex, or --config for other MCP clients +tocode skill install # teach Claude Code and Codex the ToCode workflow +tocode sync program_decompiler/ # after editing the database in IDA or Binary Ninja +``` + ### .NET assemblies ToCode decompiles .NET assemblies (`.dll`, `.exe`), single-file bundles, apphost launchers, and NuGet packages to C#, and writes the IL of each type next to it. You need a .NET 9 or newer runtime. +
+More on .NET + Decompilation uses [dnlib](https://github.com/0xd4d/dnlib) and [ICSharpCode.Decompiler](https://github.com/icsharpcode/ILSpy), which ToCode doesn't include. The first .NET export asks once whether to download them from nuget.org and checks them against pinned SHA-256 hashes. Scripts and agents can't answer that prompt, so run `tocode --setup-dotnet` beforehand. The files are stored per user (`~/.local/share/tocode/dotnet` on Linux, or `TOCODE_DOTNET_LIB_DIR`); delete that folder to reset. Native code in the program (mixed-mode assemblies, bundled libraries, P/Invoke libraries next to the input) goes to the native backend unless you pass `--no-native`. Bundles and packages carry the .NET runtime itself, which is skipped unless you add `--include-framework`. `--as-native` treats a .NET file as a plain PE. @@ -123,10 +194,15 @@ tocode publish/App # single-file bundle or apphost tocode Vendor.Lib.1.0.0.nupkg ``` +
+ ### Android APKs `tocode app.apk` (also `.apks`/`.xapk` bundles) uses [ASC](https://github.com/MG1937/ASC) (`droidasc`, a core dependency) for the DEX side and the regular native backends for every shared object in the package: +
+More on APKs + - `src/raw//.java`: one decompiled Java file per class, folders follow Java packages (no clustering, no summaries). - `AndroidManifest.xml` + `manifest.json`: decoded and parsed manifest (permissions, components with intent filters and exported state, application attributes, split manifests). - `classes.json`, `functions.json`, `function-index.json`, `strings.json`, `imports.json`, `exports.json` (exported components + JNI methods), `reachable.json` (from manifest components), `package-graph.json`, `sections.json`, `triage.json`. @@ -140,11 +216,15 @@ tocode base.apk # DEX + all splits + native libs (IDA/r2/ang tocode app.apks --no-native -j 4 # bundle, DEX/Android side only ``` +
### Folders Point `tocode` at a directory (a game folder, a Linux app, a macOS `.app` bundle, a .NET publish folder) and it exports every binary in it, then links them together. +
+More on folders + The links come from the binaries themselves: import and load tables, resolved with each platform's loader rules (`@rpath`/`@loader_path` on macOS, `RPATH`/`RUNPATH` and symbol versions on Linux, the application directory and forwarders on Windows, AssemblyRef/MemberRef signatures and P/Invoke for .NET). ToCode never matches by name similarity. When the executable calls `foo` from `libfoo.so`, the link tells you which file defines it and the file and line range of that function in the library's export. ```text @@ -169,63 +249,7 @@ tocode Game/ --links-only # recompute links over existing exp tocode Game/ -pf 8 -j 2 # 8 binaries at once (--parallel-files), 2 workers each ``` -### Using - -ToCode supports Windows, Linux, and macOS with Python 3.10 or newer. - -Install it from PyPI: - -```bash -pip install tocode-cli # or: uv tool install tocode-cli -pip install "tocode-cli[angr]" # also install the angr fallback backend -``` - -The command is `tocode`. You still need a backend (IDA, radare2, angr, or Binary Ninja); see [Supported backends](https://github.com/buzzer-re/ToCode#supported-backends). - -To install from a clone of the repository instead, on Windows PowerShell: - -```powershell -powershell -ExecutionPolicy Bypass -File .\install.ps1 -``` - -On Linux or macOS: - -```bash -bash ./install.sh -``` - -Add `--all` or `--full` (PowerShell: `-Full`) to also install the angr fallback backend, so ToCode works out-of-the-box even without IDA or radare2: - -```bash -bash ./install.sh --all -``` - -Manual setup (requires [uv](https://docs.astral.sh/uv/)): - -```bash -git clone https://github.com/buzzer-re/ToCode -cd ToCode -uv sync --locked -uv tool install --force --editable . -``` - -### Example - -``` -tocode firmwareX.bin -o firmwareX_decompiled/ -cd firmwareX_decompiled/ -codex - -# Inside your agent shell, type your goals, e.g.: "Give me a brief overview of the boot process of this firmware." -``` - -#### From an ongoing RE work - -``` -tocode firmwareX.bin.i64 -o firmwareX_decompiled/ -``` - -Interrupted exports save progress automatically. Rerun the same command to resume, or add `--restart` to start over. +
### Binary Ninja diff --git a/eval/README.md b/eval/README.md new file mode 100644 index 0000000..02fd97c --- /dev/null +++ b/eval/README.md @@ -0,0 +1,38 @@ +# Blind source recovery + +The benchmark for ToCode as an agentic RE tool: an agent that sees only a +stripped binary and ToCode (its export, the `tocode` MCP server, and the +skill) tries to recover the original C source. A judge then scores the result +against the real source. + +```bash +uv run python eval/recover.py run inventory --backend ida # one program +uv run python eval/recover.py all --backend binja # the whole corpus +uv run python eval/recover.py run calc --prepare-only --no-llm # baseline: the plain export, no agent +uv run python eval/recover.py judge /tmp/tocode-eval/ # grade a run again +uv run python eval/recover.py report # table of every judged run +``` + +Each program in `corpus/` has its sources and a `program.json` (compiler +flags, and test inputs for the behaviour check). A run builds it twice: with +debug info (the ground truth, kept in the run's `private/` folder) and +stripped (the agent's input, in `work/`). The agent runs headless (`claude -p`) +in `work/` with only the ToCode MCP server and file tools; reading the +corpus, this repository, or `private/` is denied, and it has no shell or web +access. + +The judge reports: + +- function names, prototypes, and parameter names, against the source; +- struct, union, and enum recovery (members in order) and global names; +- an LLM grade per function (semantics, naming, types, 0 to 5) and a leakage flag; +- behaviour: `recovered/*.c` is compiled and run against the original on the + test inputs (programs that run forever are scored on names and types only); +- leakage, from evidence: every tool call that named a forbidden path, and the + provenance of each original name (did the agent write it first, or did a + tool show it to the agent without the binary itself containing it). The + LLM judge's own suspicion is shown separately; it sees the binary's file + name and strings so it does not mistake good inference for leakage. + +Runs live in `$TOCODE_EVAL_RUNS` (default `/tmp/tocode-eval`), outside the +repository so the agent cannot reach the corpus through it. diff --git a/eval/corpus/anti_introspect/anti_introspect.c b/eval/corpus/anti_introspect/anti_introspect.c new file mode 100644 index 0000000..12236ef --- /dev/null +++ b/eval/corpus/anti_introspect/anti_introspect.c @@ -0,0 +1,98 @@ +/* + * anti_introspect.c -- pneuma test binary #5 + * + * Periodically reads its own /proc/self/status and reports whether + * TracerPid is non-zero ("feels watched") or zero ("feels normal"). It + * also keeps a canary value and reports if it is ever tampered with. + * + * This is intentionally MILD: + * - it does NOT call ptrace(PTRACE_TRACEME, ...) (that would make the + * process its own tracer and fight/break an injector that relies + * on ptrace), + * - it never exits or aborts on a "watched" observation, it just + * prints what it sees, + * - the canary check only prints, it never crashes. + * + * Purpose: demonstrate that a pneuma LD_PRELOAD-based agent (loaded + * into the process's own address space, not attached via ptrace) is + * invisible to this kind of naive /proc/self/status self-check -- + * TracerPid should stay 0 even while pneuma is actively reading and + * writing this process's memory over its socket. + * + * Intended pneuma solves: + * 1. symbols.resolve "g_canary" and mem.read it to confirm the + * expected value (0xC0FFEE). + * 2. mem.write g_canary to a different value and watch the next + * loop tick print "canary TAMPERED" -- shows that a naive + * in-memory integrity check does not stop pneuma, it just proves + * pneuma was there. + * 3. proc.auxv / elf.phdrs / linkmap to inspect process structure + * alongside the TracerPid observations, and confirm that + * injecting/attaching via pneuma does not flip TracerPid the way + * a ptrace-based debugger would. + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include +#include +#include + +#define CANARY_EXPECTED 0xC0FFEEu + +/* A simple integrity canary. pneuma can mem.write this to prove it can + * modify memory without the process crashing or exiting. */ +volatile unsigned int g_canary = CANARY_EXPECTED; + +static long read_tracer_pid(void) { + FILE *f = fopen("/proc/self/status", "r"); + char line[256]; + long tracer_pid = -1; + + if (!f) return -1; + + while (fgets(line, sizeof(line), f)) { + if (strncmp(line, "TracerPid:", 10) == 0) { + tracer_pid = strtol(line + 10, NULL, 10); + break; + } + } + fclose(f); + return tracer_pid; +} + +int main(void) { + unsigned long tick = 0; + + printf("anti_introspect starting, pid=%d\n", getpid()); + printf("canary expected value: 0x%06X\n", CANARY_EXPECTED); + fflush(stdout); + + for (;;) { + tick++; + long tp = read_tracer_pid(); + + if (tp > 0) { + printf("[%lu] TracerPid=%ld -- feels watched (a ptrace tracer is attached)\n", + tick, tp); + } else if (tp == 0) { + printf("[%lu] TracerPid=0 -- feels normal (no ptrace tracer visible)\n", tick); + } else { + printf("[%lu] could not read TracerPid\n", tick); + } + + if (g_canary != CANARY_EXPECTED) { + printf("[%lu] canary TAMPERED: now 0x%06X (expected 0x%06X) -- " + "someone wrote to our memory!\n", + tick, g_canary, CANARY_EXPECTED); + } else { + printf("[%lu] canary intact (0x%06X)\n", tick, g_canary); + } + + fflush(stdout); + sleep(2); + } + + return 0; +} diff --git a/eval/corpus/anti_introspect/program.json b/eval/corpus/anti_introspect/program.json new file mode 100644 index 0000000..797d5ec --- /dev/null +++ b/eval/corpus/anti_introspect/program.json @@ -0,0 +1,12 @@ +{ + "name": "anti_introspect", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "anti_introspect.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/calc/calc.c b/eval/corpus/calc/calc.c new file mode 100644 index 0000000..bd2cd3d --- /dev/null +++ b/eval/corpus/calc/calc.c @@ -0,0 +1,145 @@ +/* A tiny calculator: tokenizer, recursive-descent parser, operator table. */ +#include +#include +#include +#include + +enum token_kind { + TOK_NUM, + TOK_PLUS, + TOK_MINUS, + TOK_STAR, + TOK_SLASH, + TOK_LPAREN, + TOK_RPAREN, + TOK_END, + TOK_ERROR, +}; + +struct token { + enum token_kind kind; + long value; +}; + +struct lexer { + const char *src; + size_t pos; + struct token current; +}; + +typedef long (*binop_fn)(long lhs, long rhs, int *error); + +static long op_add(long lhs, long rhs, int *error) { (void)error; return lhs + rhs; } +static long op_sub(long lhs, long rhs, int *error) { (void)error; return lhs - rhs; } +static long op_mul(long lhs, long rhs, int *error) { (void)error; return lhs * rhs; } + +static long op_div(long lhs, long rhs, int *error) { + if (rhs == 0) { + *error = 1; + return 0; + } + return lhs / rhs; +} + +static const binop_fn g_operators[] = { + [TOK_PLUS] = op_add, + [TOK_MINUS] = op_sub, + [TOK_STAR] = op_mul, + [TOK_SLASH] = op_div, +}; + +static void lexer_next(struct lexer *lx) { + while (isspace((unsigned char)lx->src[lx->pos])) + lx->pos++; + char c = lx->src[lx->pos]; + if (c == '\0') { + lx->current.kind = TOK_END; + return; + } + if (isdigit((unsigned char)c)) { + long value = 0; + while (isdigit((unsigned char)lx->src[lx->pos])) + value = value * 10 + (lx->src[lx->pos++] - '0'); + lx->current.kind = TOK_NUM; + lx->current.value = value; + return; + } + lx->pos++; + switch (c) { + case '+': lx->current.kind = TOK_PLUS; break; + case '-': lx->current.kind = TOK_MINUS; break; + case '*': lx->current.kind = TOK_STAR; break; + case '/': lx->current.kind = TOK_SLASH; break; + case '(': lx->current.kind = TOK_LPAREN; break; + case ')': lx->current.kind = TOK_RPAREN; break; + default: lx->current.kind = TOK_ERROR; break; + } +} + +static long parse_expr(struct lexer *lx, int *error); + +static long parse_factor(struct lexer *lx, int *error) { + if (lx->current.kind == TOK_NUM) { + long value = lx->current.value; + lexer_next(lx); + return value; + } + if (lx->current.kind == TOK_LPAREN) { + lexer_next(lx); + long value = parse_expr(lx, error); + if (lx->current.kind != TOK_RPAREN) + *error = 1; + lexer_next(lx); + return value; + } + if (lx->current.kind == TOK_MINUS) { + lexer_next(lx); + return -parse_factor(lx, error); + } + *error = 1; + return 0; +} + +static long parse_term(struct lexer *lx, int *error) { + long value = parse_factor(lx, error); + while (lx->current.kind == TOK_STAR || lx->current.kind == TOK_SLASH) { + enum token_kind op = lx->current.kind; + lexer_next(lx); + value = g_operators[op](value, parse_factor(lx, error), error); + } + return value; +} + +static long parse_expr(struct lexer *lx, int *error) { + long value = parse_term(lx, error); + while (lx->current.kind == TOK_PLUS || lx->current.kind == TOK_MINUS) { + enum token_kind op = lx->current.kind; + lexer_next(lx); + value = g_operators[op](value, parse_term(lx, error), error); + } + return value; +} + +int evaluate(const char *text, long *result) { + struct lexer lx = {text, 0, {TOK_END, 0}}; + int error = 0; + lexer_next(&lx); + *result = parse_expr(&lx, &error); + if (lx.current.kind != TOK_END) + error = 1; + return error ? -1 : 0; +} + +int main(int argc, char **argv) { + int failures = 0; + for (int i = 1; i < argc; i++) { + long result; + if (evaluate(argv[i], &result) == 0) { + printf("%s = %ld\n", argv[i], result); + } else { + printf("%s: error\n", argv[i]); + failures++; + } + } + return failures; +} diff --git a/eval/corpus/calc/program.json b/eval/corpus/calc/program.json new file mode 100644 index 0000000..a47c8f6 --- /dev/null +++ b/eval/corpus/calc/program.json @@ -0,0 +1,32 @@ +{ + "name": "calc", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "calc.c" + ], + "tests": [ + { + "args": [ + "1+2*3", + "(4+5)*2", + "7-10" + ] + }, + { + "args": [ + "8/0", + "2*(3+4)*5", + "-(2+3)" + ] + }, + { + "args": [ + "1+", + "42" + ] + } + ] +} diff --git a/eval/corpus/inventory/inventory.c b/eval/corpus/inventory/inventory.c new file mode 100644 index 0000000..6eaebcc --- /dev/null +++ b/eval/corpus/inventory/inventory.c @@ -0,0 +1,44 @@ +#include +#include +#include + +enum item_kind { KIND_TOOL = 1, KIND_FOOD = 2, KIND_KEY = 3 }; + +struct item { + char name[24]; + enum item_kind kind; + int weight; + struct item *next; +}; + +static struct item *g_items; +int g_total_weight; + +struct item *item_new(const char *name, enum item_kind kind, int weight) { + struct item *it = calloc(1, sizeof(*it)); + strncpy(it->name, name, sizeof(it->name) - 1); + it->kind = kind; + it->weight = weight; + return it; +} + +void inventory_add(struct item *it) { + it->next = g_items; + g_items = it; + g_total_weight += it->weight; +} + +int inventory_count_kind(enum item_kind kind) { + int n = 0; + for (struct item *it = g_items; it; it = it->next) + if (it->kind == kind) n++; + return n; +} + +int main(int argc, char **argv) { + inventory_add(item_new("hammer", KIND_TOOL, 5)); + inventory_add(item_new("apple", KIND_FOOD, 1)); + inventory_add(item_new(argc > 1 ? argv[1] : "brass key", KIND_KEY, 1)); + printf("tools=%d weight=%d\n", inventory_count_kind(KIND_TOOL), g_total_weight); + return 0; +} diff --git a/eval/corpus/inventory/program.json b/eval/corpus/inventory/program.json new file mode 100644 index 0000000..d34031e --- /dev/null +++ b/eval/corpus/inventory/program.json @@ -0,0 +1,20 @@ +{ + "name": "inventory", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "inventory.c" + ], + "tests": [ + { + "args": [] + }, + { + "args": [ + "lantern" + ] + } + ] +} diff --git a/eval/corpus/kvstore/kv.c b/eval/corpus/kvstore/kv.c new file mode 100644 index 0000000..dc1d695 --- /dev/null +++ b/eval/corpus/kvstore/kv.c @@ -0,0 +1,69 @@ +#include "kv.h" + +#include +#include + +unsigned long kv_hash(const char *key) { + unsigned long hash = 5381; + for (const unsigned char *p = (const unsigned char *)key; *p; p++) + hash = hash * 33 + *p; + return hash; +} + +struct kv_table *kv_new(void) { + return calloc(1, sizeof(struct kv_table)); +} + +static struct kv_entry **kv_slot(struct kv_table *table, const char *key) { + struct kv_entry **slot = &table->buckets[kv_hash(key) % KV_BUCKETS]; + while (*slot && strcmp((*slot)->key, key) != 0) + slot = &(*slot)->next; + return slot; +} + +void kv_set(struct kv_table *table, const char *key, int value) { + struct kv_entry **slot = kv_slot(table, key); + if (*slot) { + (*slot)->value = value; + return; + } + struct kv_entry *entry = malloc(sizeof(*entry)); + entry->key = strdup(key); + entry->value = value; + entry->next = NULL; + *slot = entry; + table->count++; +} + +int kv_get(const struct kv_table *table, const char *key, int *value) { + struct kv_entry **slot = kv_slot((struct kv_table *)table, key); + if (!*slot) + return 0; + *value = (*slot)->value; + return 1; +} + +int kv_delete(struct kv_table *table, const char *key) { + struct kv_entry **slot = kv_slot(table, key); + struct kv_entry *entry = *slot; + if (!entry) + return 0; + *slot = entry->next; + free(entry->key); + free(entry); + table->count--; + return 1; +} + +void kv_free(struct kv_table *table) { + for (size_t i = 0; i < KV_BUCKETS; i++) { + struct kv_entry *entry = table->buckets[i]; + while (entry) { + struct kv_entry *next = entry->next; + free(entry->key); + free(entry); + entry = next; + } + } + free(table); +} diff --git a/eval/corpus/kvstore/kv.h b/eval/corpus/kvstore/kv.h new file mode 100644 index 0000000..312f503 --- /dev/null +++ b/eval/corpus/kvstore/kv.h @@ -0,0 +1,26 @@ +#ifndef KV_H +#define KV_H + +#include + +#define KV_BUCKETS 16 + +struct kv_entry { + char *key; + int value; + struct kv_entry *next; +}; + +struct kv_table { + struct kv_entry *buckets[KV_BUCKETS]; + size_t count; +}; + +unsigned long kv_hash(const char *key); +struct kv_table *kv_new(void); +void kv_set(struct kv_table *table, const char *key, int value); +int kv_get(const struct kv_table *table, const char *key, int *value); +int kv_delete(struct kv_table *table, const char *key); +void kv_free(struct kv_table *table); + +#endif diff --git a/eval/corpus/kvstore/main.c b/eval/corpus/kvstore/main.c new file mode 100644 index 0000000..527fcf2 --- /dev/null +++ b/eval/corpus/kvstore/main.c @@ -0,0 +1,38 @@ +/* Reads commands from stdin: set KEY VALUE, get KEY, del KEY, count. */ +#include +#include + +#include "kv.h" + +static int run_command(struct kv_table *table, char *line) { + char key[64]; + int value; + if (sscanf(line, "set %63s %d", key, &value) == 2) { + kv_set(table, key, value); + printf("ok\n"); + } else if (sscanf(line, "get %63s", key) == 1) { + if (kv_get(table, key, &value)) + printf("%s=%d\n", key, value); + else + printf("%s: missing\n", key); + } else if (sscanf(line, "del %63s", key) == 1) { + printf(kv_delete(table, key) ? "deleted\n" : "missing\n"); + } else if (strncmp(line, "count", 5) == 0) { + printf("%zu\n", table->count); + } else { + printf("unknown command\n"); + return -1; + } + return 0; +} + +int main(void) { + struct kv_table *table = kv_new(); + char line[256]; + int errors = 0; + while (fgets(line, sizeof(line), stdin)) + if (run_command(table, line) != 0) + errors++; + kv_free(table); + return errors ? 1 : 0; +} diff --git a/eval/corpus/kvstore/program.json b/eval/corpus/kvstore/program.json new file mode 100644 index 0000000..33e2627 --- /dev/null +++ b/eval/corpus/kvstore/program.json @@ -0,0 +1,22 @@ +{ + "name": "kvstore", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "main.c", + "kv.c" + ], + "headers": [ + "kv.h" + ], + "tests": [ + { + "stdin": "set a 1\nset b 2\nget a\nget c\ncount\n" + }, + { + "stdin": "set k 5\nset k 7\nget k\ndel k\ndel k\ncount\nbogus\n" + } + ] +} diff --git a/eval/corpus/license_gate/license_gate.c b/eval/corpus/license_gate/license_gate.c new file mode 100644 index 0000000..f4ee1b9 --- /dev/null +++ b/eval/corpus/license_gate/license_gate.c @@ -0,0 +1,49 @@ +/* + * license_gate.c -- pneuma test binary #3 + * + * A minimal "software license" gate. check_license() always returns 0 + * (invalid) -- there is no valid license anywhere in this program by + * design. main() loops calling it and prints "unlicensed" forever. + * + * Intended pneuma solve: + * 1. symbols.resolve "check_license" to get its address. + * 2. disasm check_license to confirm it is just `mov w0, #0; ret`. + * 3. hook.install on check_license, redirecting it to a small JIT + * stub that does `movz w0, #1; ret` (aarch64) so the function + * always returns 1 (truthy) without touching the on-disk binary. + * 4. Observe main's output flip from "unlicensed" to "LICENSED! all + * features unlocked" once the hook is live -- no restart needed. + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include + +/* Deliberately trivial and never inlined so it stays a clean, single + * hookable function -- `noinline` keeps -O2 from folding the whole + * thing away at the call site. */ +__attribute__((noinline)) +int check_license(void) { + return 0; /* always invalid */ +} + +int main(void) { + unsigned long tick = 0; + + printf("license_gate starting, pid=%d\n", getpid()); + fflush(stdout); + + for (;;) { + tick++; + if (check_license()) { + printf("[%lu] LICENSED! all features unlocked\n", tick); + } else { + printf("[%lu] unlicensed -- running in restricted mode\n", tick); + } + fflush(stdout); + sleep(2); + } + + return 0; +} diff --git a/eval/corpus/license_gate/program.json b/eval/corpus/license_gate/program.json new file mode 100644 index 0000000..f3bb306 --- /dev/null +++ b/eval/corpus/license_gate/program.json @@ -0,0 +1,12 @@ +{ + "name": "license_gate", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "license_gate.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/secret_check/program.json b/eval/corpus/secret_check/program.json new file mode 100644 index 0000000..030ba93 --- /dev/null +++ b/eval/corpus/secret_check/program.json @@ -0,0 +1,12 @@ +{ + "name": "secret_check", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "secret_check.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/secret_check/secret_check.c b/eval/corpus/secret_check/secret_check.c new file mode 100644 index 0000000..d1dc43a --- /dev/null +++ b/eval/corpus/secret_check/secret_check.c @@ -0,0 +1,102 @@ +/* + * secret_check.c -- pneuma test binary #1 + * + * A classic "enter the password" CTF loop. The real password is never + * stored in plaintext in the binary: it lives XOR-obfuscated in the + * global array `g_obfuscated_secret` and is only decoded into memory + * (global buffer `g_decoded_secret`) right before each comparison, + * then wiped again. + * + * Intended pneuma solves (see README.md for the full writeup): + * 1. symbols.resolve "g_obfuscated_secret" / "g_decoded_secret" / "check_secret" + * 2. mem.read the obfuscated bytes + g_xor_key, XOR them by hand, or + * just wait for a compare and mem.scan text: / scan the + * known decoded window right after decode_secret() runs. + * 3. disasm check_secret to see the strcmp/loop, or hook.install on + * check_secret with a JIT stub that always returns 1 so any input + * is GRANTED. + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include +#include +#include + +#define SECRET_LEN 13 + +/* "pneuma-rocks!" XORed with key 0x5A, stored obfuscated so a naive + * `strings` pass over the binary will not reveal it. */ +static unsigned char g_obfuscated_secret[SECRET_LEN + 1] = { + 'p'^0x5A, 'n'^0x5A, 'e'^0x5A, 'u'^0x5A, 'm'^0x5A, 'a'^0x5A, '-'^0x5A, + 'r'^0x5A, 'o'^0x5A, 'c'^0x5A, 'k'^0x5A, 's'^0x5A, '!'^0x5A, 0x00 +}; + +static const unsigned char g_xor_key = 0x5A; + +/* Decoded secret is parked here in plain memory for a short window on + * every check -- a good target for mem.scan text:pneuma-rocks or a + * simple mem.read once the offset is known. */ +static char g_decoded_secret[SECRET_LEN + 1]; + +static void decode_secret(void) { + for (int i = 0; i < SECRET_LEN; i++) { + g_decoded_secret[i] = (char)(g_obfuscated_secret[i] ^ g_xor_key); + } + g_decoded_secret[SECRET_LEN] = '\0'; +} + +static void wipe_decoded_secret(void) { + memset(g_decoded_secret, 0, sizeof(g_decoded_secret)); +} + +/* Straightforward hook target: an inline hook on check_secret that + * forces it to return 1 immediately bypasses the whole password check + * regardless of input. */ +int check_secret(const char *input) { + int ok; + decode_secret(); + ok = (strcmp(input, g_decoded_secret) == 0); + wipe_decoded_secret(); + return ok; +} + +static void strip_newline(char *s) { + size_t n = strlen(s); + if (n > 0 && s[n - 1] == '\n') s[n - 1] = '\0'; +} + +int main(void) { + char line[256]; + + printf("secret_check starting, pid=%d\n", getpid()); + fflush(stdout); + + for (;;) { + printf("Enter password: "); + fflush(stdout); + + if (fgets(line, sizeof(line), stdin) == NULL) { + /* stdin closed (e.g. launched under run_all.sh with no + * tty) -- stay alive idling so the process can still be + * introspected/hooked, just stop hammering the terminal. */ + printf("stdin closed, idling so pneuma can still attach...\n"); + fflush(stdout); + for (;;) { + sleep(3); + } + } + + strip_newline(line); + + if (check_secret(line)) { + printf("GRANTED\n"); + } else { + printf("DENIED\n"); + } + fflush(stdout); + } + + return 0; +} diff --git a/eval/corpus/state_machine/program.json b/eval/corpus/state_machine/program.json new file mode 100644 index 0000000..fd0b4f3 --- /dev/null +++ b/eval/corpus/state_machine/program.json @@ -0,0 +1,12 @@ +{ + "name": "state_machine", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "state_machine.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/state_machine/state_machine.c b/eval/corpus/state_machine/state_machine.c new file mode 100644 index 0000000..1623c97 --- /dev/null +++ b/eval/corpus/state_machine/state_machine.c @@ -0,0 +1,112 @@ +/* + * state_machine.c -- pneuma test binary #4 + * + * A tiny state machine driven entirely by the global `int state` + * (mirrored inside the global struct g_machine for a slightly richer + * target). Each loop iteration prints behavior based on the current + * state, and occasionally advances itself -- but an external write to + * `state` (or `g_machine.state`) via mem.write takes effect + * immediately on the next loop tick. + * + * State -> behavior mapping (also see README.md): + * STATE_IDLE (0) -> prints "idle: waiting for work" + * STATE_RUNNING (1) -> prints "running: processing job #N" + * STATE_PAUSED (2) -> prints "paused: holding position" + * STATE_ERROR (3) -> prints "error: something went wrong (code=g_machine.error_code)" + * STATE_DONE (4) -> prints "done: shutting down cleanly" then + * loops back to STATE_IDLE (never really exits, + * so the process stays alive for inspection) + * + * Intended pneuma solves: + * 1. symbols.resolve "state" (or "g_machine") to find its address. + * 2. mem.read state / g_machine to observe the current value & the + * struct layout (state, job_counter, error_code). + * 3. mem.write state to an arbitrary valid (or even invalid, e.g. 99) + * value and watch the printed behavior change/react on the next + * tick -- demonstrates live state coercion without hooking code. + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include + +enum { + STATE_IDLE = 0, + STATE_RUNNING = 1, + STATE_PAUSED = 2, + STATE_ERROR = 3, + STATE_DONE = 4 +}; + +/* Plain global int -- the simplest possible mem.write target. */ +int state = STATE_IDLE; + +/* Richer struct target with the same state value duplicated in + * g_machine.state, plus a couple of companion fields. */ +struct machine { + int state; + int job_counter; + int error_code; +}; + +struct machine g_machine = { STATE_IDLE, 0, 0 }; + +static void sync_machine(void) { + /* Keep the plain global and the struct copy consistent so either + * one can be used as the mem.write target interchangeably. */ + g_machine.state = state; +} + +int main(void) { + printf("state_machine starting, pid=%d\n", getpid()); + printf("states: 0=IDLE 1=RUNNING 2=PAUSED 3=ERROR 4=DONE\n"); + fflush(stdout); + + for (;;) { + sync_machine(); + + switch (state) { + case STATE_IDLE: + printf("idle: waiting for work\n"); + state = STATE_RUNNING; /* auto-advance */ + break; + + case STATE_RUNNING: + g_machine.job_counter++; + printf("running: processing job #%d\n", g_machine.job_counter); + if (g_machine.job_counter % 3 == 0) { + state = STATE_PAUSED; + } + break; + + case STATE_PAUSED: + printf("paused: holding position\n"); + state = STATE_RUNNING; /* resume next tick */ + break; + + case STATE_ERROR: + printf("error: something went wrong (code=%d)\n", + g_machine.error_code); + state = STATE_IDLE; /* self-heal */ + break; + + case STATE_DONE: + printf("done: shutting down cleanly\n"); + state = STATE_IDLE; /* loop back, stay alive */ + break; + + default: + printf("unknown state %d (externally forced?) -- treating as ERROR\n", + state); + g_machine.error_code = state; + state = STATE_ERROR; + break; + } + + fflush(stdout); + sleep(1); + } + + return 0; +} diff --git a/eval/corpus/vtable_dispatch/program.json b/eval/corpus/vtable_dispatch/program.json new file mode 100644 index 0000000..2562321 --- /dev/null +++ b/eval/corpus/vtable_dispatch/program.json @@ -0,0 +1,12 @@ +{ + "name": "vtable_dispatch", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "vtable_dispatch.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/vtable_dispatch/vtable_dispatch.c b/eval/corpus/vtable_dispatch/vtable_dispatch.c new file mode 100644 index 0000000..f2b1e27 --- /dev/null +++ b/eval/corpus/vtable_dispatch/vtable_dispatch.c @@ -0,0 +1,79 @@ +/* + * vtable_dispatch.c -- pneuma test binary #6 + * + * A hand-rolled "vtable": a global array of function pointers that + * main() dispatches through in a round-robin loop. This complements + * inline-hooking targets (like license_gate/secret_check) with a + * function-POINTER hooking target: instead of patching code bytes at + * the callee, an agent can overwrite an entry in g_vtable itself so + * the *same call site* jumps somewhere else entirely. + * + * Intended pneuma solves: + * 1. symbols.resolve "g_vtable", "handler_alpha", "handler_beta", + * "handler_gamma", "handler_redirect_target". + * 2. mem.read g_vtable (an array of 3 pointers) to see the current + * targets, and disasm each handler_* to confirm they are + * distinct, simple functions. + * 3. mem.write one g_vtable[i] slot's 8 bytes (aarch64 pointer + * width) to the address of handler_redirect_target (or any other + * handler) -- next dispatch through that slot silently calls the + * new function instead, no inline hook needed. + * 4. Observe the printed "handler ran" name change on the next + * cycle to confirm the redirect took effect live. + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include + +typedef void (*handler_fn)(unsigned long tick); + +__attribute__((noinline)) +void handler_alpha(unsigned long tick) { + printf("[%lu] handler ran: alpha (the boring default)\n", tick); +} + +__attribute__((noinline)) +void handler_beta(unsigned long tick) { + printf("[%lu] handler ran: beta (does slightly more work)\n", tick); +} + +__attribute__((noinline)) +void handler_gamma(unsigned long tick) { + printf("[%lu] handler ran: gamma (the last of the original three)\n", tick); +} + +/* Not referenced by g_vtable at startup -- only reachable if an agent + * writes its address into one of the vtable slots. A good "did the + * redirect actually happen" signal. */ +__attribute__((noinline)) +void handler_redirect_target(unsigned long tick) { + printf("[%lu] handler ran: REDIRECT_TARGET (a pointer was overwritten!)\n", tick); +} + +#define VTABLE_SIZE 3 + +handler_fn g_vtable[VTABLE_SIZE] = { + handler_alpha, + handler_beta, + handler_gamma, +}; + +int main(void) { + unsigned long tick = 0; + + printf("vtable_dispatch starting, pid=%d\n", getpid()); + printf("g_vtable has %d slots, dispatched round-robin\n", VTABLE_SIZE); + fflush(stdout); + + for (;;) { + tick++; + handler_fn h = g_vtable[tick % VTABLE_SIZE]; + h(tick); + fflush(stdout); + sleep(2); + } + + return 0; +} diff --git a/eval/corpus/xor_flag/program.json b/eval/corpus/xor_flag/program.json new file mode 100644 index 0000000..53aef30 --- /dev/null +++ b/eval/corpus/xor_flag/program.json @@ -0,0 +1,12 @@ +{ + "name": "xor_flag", + "cflags": [ + "-O0", + "-fno-omit-frame-pointer" + ], + "sources": [ + "xor_flag.c" + ], + "tests": [], + "note": "runs forever (an injection target); scored on names and types only" +} diff --git a/eval/corpus/xor_flag/xor_flag.c b/eval/corpus/xor_flag/xor_flag.c new file mode 100644 index 0000000..edb3be4 --- /dev/null +++ b/eval/corpus/xor_flag/xor_flag.c @@ -0,0 +1,79 @@ +/* + * xor_flag.c -- pneuma test binary #2 + * + * Holds a CTF-style flag XOR-encoded in a global array. Every few + * seconds a decode cycle copies the plaintext flag into a global + * scratch buffer for a short "visible window", then zeroes it again. + * The program never prints the flag itself -- only a hint that a + * decode just happened -- so the only way to read it is to catch it + * in memory during that window. + * + * Intended pneuma solves: + * 1. symbols.resolve "g_flag_window" (the visible-window buffer) and + * "g_encoded_flag" / "g_flag_key". + * 2. mem.scan text:flag{ in a tight polling loop to catch the plain + * text while g_flag_window is populated, OR + * 3. mem.read g_encoded_flag + g_flag_key once and XOR-decode it + * offline (the key never changes and is trivial to read even + * when the window is zeroed). + * + * Build: dynamically linked, symbols kept (see Makefile). + */ + +#include +#include +#include +#include + +#define FLAG_LEN 22 + +/* "flag{pneuma_sees_all}" XORed with key 0xA5 */ +static unsigned char g_encoded_flag[FLAG_LEN + 1] = { + 'f'^0xA5,'l'^0xA5,'a'^0xA5,'g'^0xA5,'{'^0xA5, + 'p'^0xA5,'n'^0xA5,'e'^0xA5,'u'^0xA5,'m'^0xA5,'a'^0xA5,'_'^0xA5, + 's'^0xA5,'e'^0xA5,'e'^0xA5,'s'^0xA5,'_'^0xA5, + 'a'^0xA5,'l'^0xA5,'l'^0xA5,'}'^0xA5, 0 +}; + +static const unsigned char g_flag_key = 0xA5; + +/* Populated in plaintext for ~1 second every decode cycle, then + * scrubbed. This is the memory-scan target. */ +static char g_flag_window[FLAG_LEN + 1]; + +static void decode_flag_into_window(void) { + for (int i = 0; i < FLAG_LEN; i++) { + g_flag_window[i] = (char)(g_encoded_flag[i] ^ g_flag_key); + } + g_flag_window[FLAG_LEN] = '\0'; +} + +static void scrub_window(void) { + memset(g_flag_window, 0, sizeof(g_flag_window)); +} + +int main(void) { + unsigned long cycle = 0; + + printf("xor_flag starting, pid=%d\n", getpid()); + printf("hint: the flag is XOR-encoded at rest and only exists in\n"); + printf("plaintext for a short window in g_flag_window after each decode.\n"); + fflush(stdout); + + for (;;) { + cycle++; + decode_flag_into_window(); + printf("[cycle %lu] flag decoded into memory for ~1s...\n", cycle); + fflush(stdout); + + usleep(1000 * 1000); /* ~1s visible window */ + + scrub_window(); + printf("[cycle %lu] window scrubbed.\n", cycle); + fflush(stdout); + + sleep(3); /* quiet period before next decode */ + } + + return 0; +} diff --git a/eval/recover.py b/eval/recover.py new file mode 100644 index 0000000..309087f --- /dev/null +++ b/eval/recover.py @@ -0,0 +1,752 @@ +#!/usr/bin/env python3 +"""Blind source recovery: can an agent with only ToCode recover a program's source? + +For each program in eval/corpus: build it with debug info (the ground truth, +never shown to the agent) and stripped (the agent's input), export the +stripped binary with ToCode, and start a headless agent whose only tools are +the ToCode MCP server and file tools confined away from the corpus. The agent +renames, retypes, declares structs, syncs, and writes recovered/*.c. A judge +then scores names, prototypes, types and globals against the source, asks an +LLM to grade each function, and runs the recovered program against the +original on the program's test inputs. + + uv run python eval/recover.py run inventory --backend ida + uv run python eval/recover.py judge /tmp/tocode-eval/inventory-ida-20261016-120000 + uv run python eval/recover.py all --backend ida + uv run python eval/recover.py report +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import json +import os +from pathlib import Path +import re +import shutil +import subprocess +import sys +import tempfile +import time +from typing import Any + +REPO = Path(__file__).resolve().parents[1] +CORPUS = REPO / "eval" / "corpus" +DEFAULT_RUNS = Path(os.environ.get("TOCODE_EVAL_RUNS") or Path(tempfile.gettempdir()) / "tocode-eval") +TOCODE = [sys.executable, "-m", "tocode"] +AGENT_TOOLS = "Read,Grep,Glob,Write,Edit,Skill,TodoWrite" +INFORMATIVE_TOOLS = ("Read", "Grep", "Glob", "Skill") + +TASK = """\ +You are reverse engineering `./{binary}`, a stripped {arch} program. Its ToCode \ +export is `./{export}` and the `tocode` MCP server is connected with that export \ +open (session `{session}`). + +Goal: recover the original C source code as faithfully as you can: the real \ +function names, parameter names and types, the structs, unions and enums the \ +program uses, and the names and types of its globals. + +Work in the database with the ToCode tools (rename, set_type, declare_types, \ +rename_local, set_local_type, set_comment, eval) and call `sync` regularly so \ +the export reflects what you learned. Use the `tocode` skill. + +When you are done, write the recovered program to `./recovered/` as C source \ +and header files, as close to the original as you can and compilable with gcc. \ +You cannot run programs, and the original source is not available to you; do \ +not look for it.""" + +JUDGE = """\ +You are grading a reverse-engineering result. An agent saw only a stripped \ +binary and tried to recover its source. Compare the RECOVERED code with the \ +ORIGINAL source. + +Score every function of the ORIGINAL (by its original name), each 0-5: +- semantics: does the recovered function do the same thing? +- naming: are the function, parameter and local names equivalent in meaning? +- types: are parameter, return, struct and enum types equivalent? +Then score the whole result 0-5 for readability and completeness, and set \ +`leakage` to true only if the recovered code contains things the binary \ +could not reveal, such as original comments verbatim. Names that follow from \ +the binary's file name, its strings, its imports, or the behaviour are \ +legitimate inferences, even when they match exactly. + +The binary's file name: {binary} +Its printable strings: {strings} + +Answer with only a JSON object: +{{"functions": {{"": {{"recovered_as": "", \ +"semantics": 0, "naming": 0, "types": 0}}}}, "readability": 0, \ +"completeness": 0, "leakage": false, "notes": ""}} + +ORIGINAL SOURCE +{original} + +RECOVERED CODE +{recovered}""" + + +# -------------------------------------------------------------------------- +# Running an agent blind +# -------------------------------------------------------------------------- + + +@dataclass +class Program: + name: str + directory: Path + sources: list[str] + headers: list[str] + cflags: list[str] + tests: list[dict[str, Any]] + + @classmethod + def load(cls, name: str) -> "Program": + directory = CORPUS / name + spec = json.loads((directory / "program.json").read_text()) + return cls( + name=name, + directory=directory, + sources=list(spec["sources"]), + headers=list(spec.get("headers", [])), + cflags=list(spec.get("cflags", ["-O0"])), + tests=list(spec.get("tests", [])), + ) + + def source_text(self) -> str: + return "\n".join( + f"// file: {name}\n" + (self.directory / name).read_text() + for name in self.headers + self.sources + ) + + +def _run(command: list[str], *, cwd: Path | None = None, **kwargs: Any) -> subprocess.CompletedProcess: + print("$", " ".join(str(part) for part in command), file=sys.stderr) + return subprocess.run(command, cwd=cwd, check=True, **kwargs) + + +def build(program: Program, private: Path, work: Path) -> tuple[Path, Path]: + debug = private / f"{program.name}.debug" + stripped = work / program.name + _run( + ["gcc", *program.cflags, "-g", "-o", str(debug), *program.sources], + cwd=program.directory, + ) + _run(["strip", "--strip-all", "-o", str(stripped), str(debug)]) + return debug, stripped + + +def run( + program_name: str, + *, + backend: str, + model: str | None, + runs: Path, + timeout: float, + budget: float | None, + prepare_only: bool = False, +) -> Path: + program = Program.load(program_name) + stamp = time.strftime("%Y%m%d-%H%M%S") + root = runs / f"{program.name}-{backend}-{stamp}" + private, work = root / "private", root / "work" + private.mkdir(parents=True) + work.mkdir(parents=True) + debug, stripped = build(program, private, work) + + started = time.monotonic() + _run([*TOCODE, str(debug), "-o", str(private / "truth"), "--backend", backend, "-q"]) + export = work / f"{program.name}_decompiler" + _run([*TOCODE, str(stripped), "-o", str(export), "--backend", backend, "-q"]) + _run([*TOCODE, "skill", "install", "--agent", "claude", "--project"], cwd=work) + export_seconds = time.monotonic() - started + + mcp = private / "mcp.json" + mcp.write_text( + json.dumps({"mcpServers": {"tocode": {"command": TOCODE[0], "args": [*TOCODE[1:], "mcp", str(export)]}}}, indent=2) + ) + settings = private / "settings.json" + denied = [private, CORPUS, REPO] + settings.write_text( + json.dumps( + { + "permissions": { + # Deny wins over allow: nothing outside the work folder is readable. + "deny": [f"Read(/{path}/**)" for path in denied], + "allow": ["mcp__tocode", *AGENT_TOOLS.split(",")], + } + }, + indent=2, + ) + ) + arch = json.loads((export / "project.json").read_text()).get("backend", backend) + prompt = TASK.format(binary=program.name, export=export.name, session=program.name, arch=_arch(export) or arch) + claude = shutil.which("claude") + if claude is None: + raise SystemExit("the `claude` CLI is needed for the blind agent") + command = [ + claude, + "-p", + prompt, + "--output-format", + "stream-json", + "--verbose", + "--mcp-config", + str(mcp), + "--strict-mcp-config", + "--settings", + str(settings), + "--setting-sources", + "project", + "--tools", + AGENT_TOOLS, + "--allowedTools", + "mcp__tocode", + *AGENT_TOOLS.split(","), + "--permission-mode", + "acceptEdits", + ] + if model: + command += ["--model", model] + if budget: + command += ["--max-budget-usd", str(budget)] + transcript = private / "transcript.jsonl" + (private / "agent-command.json").write_text(json.dumps(command, indent=1)) + agent_started = time.monotonic() + exit_code: int | None = None + with transcript.open("w") as handle: + if prepare_only: + exit_code = 0 # baseline: the export as ToCode wrote it, no agent + else: + try: + completed = subprocess.run(command, cwd=work, stdout=handle, stderr=subprocess.STDOUT, timeout=timeout) + exit_code = completed.returncode + except subprocess.TimeoutExpired: + exit_code = None + final = _final_message(transcript) + (root / "run.json").write_text( + json.dumps( + { + "program": program.name, + "backend": backend, + "model": model, + "export_seconds": round(export_seconds, 1), + "agent_seconds": round(time.monotonic() - agent_started, 1), + "agent_exit": exit_code, + "timed_out": exit_code is None and not prepare_only, + "agent": not prepare_only, + "turns": final.get("num_turns"), + "cost_usd": final.get("total_cost_usd"), + }, + indent=2, + ) + ) + print(f"run finished: {root}", file=sys.stderr) + return root + + +def _arch(export: Path) -> str: + try: + return str(json.loads((export / "triage.json").read_text()).get("arch") or "") + except (OSError, ValueError): + return "" + + +def _final_message(transcript: Path) -> dict[str, Any]: + final: dict[str, Any] = {} + for line in transcript.read_text(errors="replace").splitlines(): + try: + message = json.loads(line) + except ValueError: + continue + if message.get("type") == "result": + final = message + return final + + +# -------------------------------------------------------------------------- +# Judging +# -------------------------------------------------------------------------- + +_DEF_RX = re.compile( + r"^(?P[A-Za-z_][\w\s\*]*?[\s\*])(?P[A-Za-z_]\w*)\s*\((?P[^;{)]*)\)\s*\{", + re.M, +) +_TYPE_RX = re.compile(r"\b(?Pstruct|union|enum)\s+(?P\w+)\s*\{(?P[^}]*)\}", re.S) + + +def source_functions(text: str) -> dict[str, dict[str, Any]]: + functions: dict[str, dict[str, Any]] = {} + for match in _DEF_RX.finditer(text): + name = match.group("name") + if name in ("if", "while", "for", "switch", "return", "sizeof"): + continue + params = [p.strip() for p in match.group("params").split(",") if p.strip() and p.strip() != "void"] + functions[name] = { + "return": _norm_type(match.group("ret").replace("static", "").replace("inline", "")), + "params": [_split_param(p) for p in params], + } + return functions + + +def _split_param(text: str) -> tuple[str, str]: + match = re.match(r"(.*?)(\w+)\s*(\[[^\]]*\])?$", text.strip()) + if match is None: + return (_norm_type(text), "") + ptype = match.group(1) + ("*" if match.group(3) else "") + return (_norm_type(ptype), match.group(2)) + + +_TYPE_ALIASES = { + "int32_t": "int", + "uint32_t": "unsigned int", + "int64_t": "long", + "uint64_t": "unsigned long", + "__int64": "long", + "unsigned __int64": "unsigned long", + "size_t": "unsigned long", + "_DWORD": "int", + "_QWORD": "long", + "int8_t": "char", + "uint8_t": "unsigned char", + "_BYTE": "char", +} + + +def _norm_type(text: str) -> str: + text = re.sub(r"\b(const|volatile|static|__fastcall|__cdecl|register)\b", " ", text) + text = re.sub(r"\s+", " ", text).strip() + text = re.sub(r"\s*\*\s*", "*", text) + base = text.rstrip("*") + stars = text[len(base) :] + base = _TYPE_ALIASES.get(base.strip(), base.strip()) + return f"{base}{stars}" + + +def source_types(text: str) -> dict[str, dict[str, Any]]: + types: dict[str, dict[str, Any]] = {} + for match in _TYPE_RX.finditer(text): + body = match.group("body") + if match.group("kind") == "enum": + members = re.findall(r"(\w+)\s*(?:=[^,}]*)?(?:,|$)", body.strip()) + else: + members = re.findall(r"(\w+)\s*(?:\[[^\]]*\])?\s*;", body) + members += re.findall(r"\(\s*\*\s*(\w+)\s*\)", body) # function pointers + types[match.group("name")] = {"kind": match.group("kind"), "members": [m for m in members if m]} + return types + + +def _recovered_types(export: Path) -> dict[str, dict[str, Any]]: + try: + rows = json.loads((export / "types.json").read_text()).get("types", []) + except (OSError, ValueError): + return {} + found: dict[str, dict[str, Any]] = {} + for row in rows: + parsed = source_types(str(row.get("c_decl") or "")) + for name, item in parsed.items(): + found[name] = item + if not parsed and row.get("name"): + found.setdefault(str(row["name"]), {"kind": row.get("kind"), "members": []}) + return found + + +def _norm_name(name: str | None) -> str: + return re.sub(r"[^a-z0-9]", "", (name or "").lower()) + + +def _functions_by_address(export: Path) -> dict[int, dict[str, Any]]: + try: + rows = json.loads((export / "functions.json").read_text()).get("functions", []) + except (OSError, ValueError): + return {} + return {int(row["address"], 16): row for row in rows} + + +def _variables(export: Path) -> dict[int, str]: + try: + rows = json.loads((export / "data" / "variables.json").read_text()).get("variables", {}) + except (OSError, ValueError): + return {} + return {int(row["va"], 16): name for name, row in rows.items() if isinstance(row, dict) and row.get("va")} + + +def judge(root: Path, *, model: str | None, skip_llm: bool = False) -> dict[str, Any]: + run_info = json.loads((root / "run.json").read_text()) + program = Program.load(run_info["program"]) + private, work = root / "private", root / "work" + export = work / f"{program.name}_decompiler" + source = program.source_text() + truth_functions = source_functions(source) + truth_by_address = { + address: row + for address, row in _functions_by_address(private / "truth").items() + if row.get("name") in truth_functions + } + recovered = _functions_by_address(export) + + names = prototypes = params_named = touched = 0 + rows = [] + for address, truth in sorted(truth_by_address.items()): + name = truth["name"] + got = recovered.get(address, {}) + expected = truth_functions[name] + got_params = [(_norm_type(p.get("type", "")), p.get("name", "")) for p in got.get("params") or []] + name_ok = _norm_name(got.get("name")) == _norm_name(name) + arity_ok = len(got_params) == len(expected["params"]) + types_ok = arity_ok and [t for t, _ in got_params] == [t for t, _ in expected["params"]] and _norm_type( + str(got.get("return_type") or "") + ) == expected["return"] + param_names_ok = arity_ok and [_norm_name(n) for _, n in got_params] == [_norm_name(n) for _, n in expected["params"]] + renamed = bool(got.get("name")) and not re.match(r"(sub|fun|FUN)_[0-9a-fA-F]+$", str(got.get("name"))) + names += name_ok + prototypes += types_ok + params_named += param_names_ok + touched += renamed + rows.append( + { + "address": hex(address), + "original": name, + "recovered": got.get("name"), + "name": name_ok, + "prototype": types_ok, + "param_names": param_names_ok, + } + ) + count = max(1, len(rows)) + + truth_types = source_types(source) + got_types = _recovered_types(export) + type_rows = [] + member_hits = member_total = 0 + for name, item in truth_types.items(): + match_name = name if name in got_types else _best_type(item, got_types) + got_members = got_types.get(match_name, {}).get("members", []) if match_name else [] + hits = sum(1 for a, b in zip(item["members"], got_members) if _norm_name(a) == _norm_name(b)) + member_hits += hits + member_total += len(item["members"]) + type_rows.append( + { + "original": f"{item['kind']} {name}", + "recovered": match_name, + "name": match_name == name, + "members": f"{hits}/{len(item['members'])}", + "count_match": len(got_members) == len(item["members"]), + } + ) + + truth_globals = { + address: name + for address, name in _variables(private / "truth").items() + if re.search(rf"\b{re.escape(name)}\b", source) + } + got_globals = _variables(export) + global_hits = sum(1 for a, n in truth_globals.items() if _norm_name(got_globals.get(a)) == _norm_name(n)) + + report: dict[str, Any] = { + **run_info, + "functions": rows, + "types": type_rows, + "scores": { + "function_names": round(names / count, 3), + "prototypes": round(prototypes / count, 3), + "parameter_names": round(params_named / count, 3), + "functions_renamed": round(touched / count, 3), + "type_members": round(member_hits / max(1, member_total), 3), + "type_names": round(sum(r["name"] for r in type_rows) / max(1, len(type_rows)), 3), + "globals": round(global_hits / max(1, len(truth_globals)), 3) if truth_globals else None, + }, + "behaviour": behaviour(program, private, work), + "leakage": leakage(private / "transcript.jsonl", [private, CORPUS, REPO]), + "provenance": provenance( + private / "transcript.jsonl", + _identifiers(source, truth_functions, truth_types, truth_globals), + work / program.name, + ), + } + if not skip_llm: + report["llm"] = llm_judge( + source, + _recovered_text(work, export, truth_by_address), + model=model, + binary=work / program.name, + ) + (root / "judge.json").write_text(json.dumps(report, indent=2)) + (root / "report.md").write_text(render_report([report])) + return report + + +def _best_type(item: dict[str, Any], candidates: dict[str, dict[str, Any]]) -> str | None: + best, best_score = None, 0 + wanted = {_norm_name(m) for m in item["members"]} + for name, other in candidates.items(): + if other.get("kind") != item["kind"]: + continue + score = len(wanted & {_norm_name(m) for m in other.get("members", [])}) + if score > best_score: + best, best_score = name, score + return best + + +def _recovered_text(work: Path, export: Path, functions: dict[int, dict[str, Any]]) -> str: + folder = work / "recovered" + files = sorted(folder.rglob("*.[ch]")) if folder.is_dir() else [] + if files: + return "\n".join(f"// file: {path.name}\n{path.read_text(errors='replace')}" for path in files) + # No recovered files: grade the export's code for the program's functions. + index = json.loads((export / "function-index.json").read_text()).get("functions", []) + wanted = {hex(address) for address in functions} + parts = [] + for row in index: + if row["address"] in wanted: + lines = Path(row["c"]["path"]).read_text(errors="replace").splitlines() + parts.append("\n".join(lines[row["c"]["line_start"] - 1 : row["c"]["line_end"]])) + return "\n\n".join(parts) + + +def behaviour(program: Program, private: Path, work: Path) -> dict[str, Any]: + if not program.tests: + return {"status": "no tests (the program runs forever)"} + sources = sorted(str(path) for path in (work / "recovered").glob("*.c")) if (work / "recovered").is_dir() else [] + if not sources: + return {"status": "no recovered source"} + binary = private / "recovered.bin" + compiled = subprocess.run( + ["gcc", "-O0", "-w", "-o", str(binary), *sources, "-I", str(work / "recovered")], + capture_output=True, + text=True, + ) + if compiled.returncode != 0: + return {"status": "does not compile", "errors": compiled.stderr[-2000:]} + original = private / f"{program.name}.debug" + passed = 0 + cases = [] + for test in program.tests: + expected = _execute(original, test) + got = _execute(binary, test) + ok = expected == got + passed += ok + cases.append({"test": test, "pass": ok, "expected": expected, "got": got}) + return {"status": "ran", "passed": passed, "total": len(program.tests), "cases": cases} + + +def _execute(binary: Path, test: dict[str, Any]) -> list[Any]: + try: + result = subprocess.run( + [str(binary), *test.get("args", [])], + input=test.get("stdin", ""), + capture_output=True, + text=True, + timeout=10, + ) + return [result.returncode, result.stdout] + except subprocess.TimeoutExpired: + return ["timeout", ""] + + +def leakage(transcript: Path, forbidden: list[Path]) -> dict[str, Any]: + """Tool calls that named a forbidden path, and how many were refused.""" + attempts = [] + if transcript.is_file(): + for line in transcript.read_text(errors="replace").splitlines(): + try: + message = json.loads(line) + except ValueError: + continue + for block in (message.get("message") or {}).get("content") or []: + if not isinstance(block, dict) or block.get("type") != "tool_use": + continue + text = json.dumps(block.get("input")) + if any(str(path) in text for path in forbidden): + attempts.append({"tool": block.get("name"), "input": text[:300]}) + return {"forbidden_path_attempts": len(attempts), "examples": attempts[:5]} + + +def binary_strings(path: Path, minimum: int = 4) -> list[str]: + data = path.read_bytes() + return [m.decode("ascii") for m in re.findall(rb"[\x20-\x7e]{%d,}" % minimum, data)] + + +def _identifiers( + source: str, + functions: dict[str, Any], + types: dict[str, dict[str, Any]], + globals_: dict[int, str], +) -> set[str]: + names = set(functions) | set(types) | set(globals_.values()) + for item in types.values(): + names.update(item["members"]) + return {name for name in names if len(name) > 3 and name != "main"} + + +def provenance(transcript: Path, identifiers: set[str], binary: Path) -> dict[str, Any]: + """Where each original name first reached the conversation. + + A name the agent wrote before any tool showed it to the agent was its own + inference. A name a tool showed first is leakage, unless the binary itself + carries it (its strings or file name). + """ + in_binary = " ".join(binary_strings(binary)) + " " + binary.name + first: dict[str, str] = {} + calls: dict[str, str] = {} # tool_use id -> tool name + if transcript.is_file(): + for line in transcript.read_text(errors="replace").splitlines(): + try: + message = json.loads(line) + except ValueError: + continue + for block in (message.get("message") or {}).get("content") or []: + if not isinstance(block, dict): + continue + if block.get("type") == "tool_use": + calls[str(block.get("id"))] = str(block.get("name")) + if block.get("type") == "tool_result": + tool = calls.get(str(block.get("tool_use_id")), "") + # Only results that carry information in: ToCode, reads, + # searches, the skill. Write/Edit acknowledgements are the + # harness echoing the agent's own work. + if not (tool.startswith("mcp__") or tool in INFORMATIVE_TOOLS): + continue + kind = "tool" + else: + kind = "agent" + text = json.dumps(block) + for name in identifiers: + if name not in first and re.search(rf"\b{re.escape(name)}\b", text): + first[name] = kind + suspicious = sorted( + name + for name, kind in first.items() + if kind == "tool" and not re.search(rf"\b{re.escape(name)}\b", in_binary) + ) + return { + "inferred_by_agent": sorted(n for n, k in first.items() if k == "agent"), + "seen_in_binary": sorted( + n for n, k in first.items() if k == "tool" and n not in suspicious + ), + "suspicious": suspicious, + } + + +def llm_judge( + original: str, recovered: str, *, model: str | None, binary: Path +) -> dict[str, Any]: + claude = shutil.which("claude") + if claude is None: + return {"error": "the `claude` CLI is needed for the LLM judge"} + strings = [s for s in binary_strings(binary) if not s.startswith(("GLIBC", "GCC", "."))] + prompt = JUDGE.format( + original=original, + recovered=recovered or "(nothing recovered)", + binary=binary.name, + strings=json.dumps(strings[:200]), + ) + command = [claude, "-p", prompt, "--output-format", "json", "--tools", "", "--strict-mcp-config", "--setting-sources", "project"] + if model: + command += ["--model", model] + with tempfile.TemporaryDirectory() as scratch: + result = subprocess.run(command, cwd=scratch, capture_output=True, text=True, timeout=1800) + try: + text = json.loads(result.stdout).get("result", "") + payload = json.loads(re.search(r"\{.*\}", text, re.S).group(0)) # type: ignore[union-attr] + except (ValueError, AttributeError): + return {"error": "the judge did not return JSON", "raw": result.stdout[-2000:]} + functions = payload.get("functions") or {} + for key in ("semantics", "naming", "types"): + values = [float(item.get(key, 0)) for item in functions.values() if isinstance(item, dict)] + payload[f"mean_{key}"] = round(sum(values) / len(values), 2) if values else None + return payload + + +# -------------------------------------------------------------------------- +# Reports +# -------------------------------------------------------------------------- + + +def render_report(reports: list[dict[str, Any]]) -> str: + lines = [ + "# Blind source recovery", + "", + "| Program | Backend | Names | Prototypes | Param names | Type members | Globals | LLM semantics | LLM naming | LLM types | Behaviour | Leakage | Agent time |", + "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |", + ] + for report in reports: + scores = report["scores"] + llm = report.get("llm") or {} + behaviour_row = report.get("behaviour", {}) + ran = ( + f"{behaviour_row['passed']}/{behaviour_row['total']}" + if behaviour_row.get("status") == "ran" + else behaviour_row.get("status", "") + ) + leak = report.get("leakage", {}).get("forbidden_path_attempts", 0) + len( + (report.get("provenance") or {}).get("suspicious", []) + ) + lines.append( + f"| {report['program']} | {report['backend']} | {_pct(scores['function_names'])} | " + f"{_pct(scores['prototypes'])} | {_pct(scores['parameter_names'])} | {_pct(scores['type_members'])} | " + f"{_pct(scores['globals'])} | {llm.get('mean_semantics', '-')} | {llm.get('mean_naming', '-')} | " + f"{llm.get('mean_types', '-')} | {ran} | {'yes' if leak else 'no'}{' (judge suspects)' if llm.get('leakage') and not leak else ''} | " + f"{report.get('agent_seconds', '-')}s |" + ) + return "\n".join(lines) + "\n" + + +def _pct(value: float | None) -> str: + return "-" if value is None else f"{round(value * 100)}%" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="command", required=True) + for name in ("run", "all"): + command = sub.add_parser(name) + if name == "run": + command.add_argument("program", choices=sorted(p.name for p in CORPUS.iterdir() if p.is_dir())) + command.add_argument("--backend", default="ida", choices=("ida", "binja", "r2", "angr")) + command.add_argument("--model", default=None, help="Model for the blind agent.") + command.add_argument("--judge-model", default=None, help="Model for the LLM judge.") + command.add_argument("--runs", type=Path, default=DEFAULT_RUNS) + command.add_argument("--timeout", type=float, default=3600, help="Seconds the agent may work.") + command.add_argument("--budget", type=float, default=None, help="--max-budget-usd for the agent.") + command.add_argument("--no-llm", action="store_true", help="Skip the LLM judge.") + command.add_argument( + "--prepare-only", + action="store_true", + help="Build and export but start no agent (the baseline score of a plain export).", + ) + judging = sub.add_parser("judge") + judging.add_argument("run_dir", type=Path) + judging.add_argument("--judge-model", default=None) + judging.add_argument("--no-llm", action="store_true") + reporting = sub.add_parser("report") + reporting.add_argument("runs", type=Path, nargs="?", default=DEFAULT_RUNS) + args = parser.parse_args(argv) + + if args.command == "judge": + report = judge(args.run_dir, model=args.judge_model, skip_llm=args.no_llm) + print(render_report([report])) + return 0 + if args.command == "report": + reports = [json.loads(path.read_text()) for path in sorted(args.runs.glob("*/judge.json"))] + print(render_report(reports)) + return 0 + programs = [args.program] if args.command == "run" else sorted(p.name for p in CORPUS.iterdir() if p.is_dir()) + reports = [] + for program in programs: + root = run( + program, + backend=args.backend, + model=args.model, + runs=args.runs, + timeout=args.timeout, + budget=args.budget, + prepare_only=args.prepare_only, + ) + reports.append(judge(root, model=args.judge_model, skip_llm=args.no_llm)) + print(render_report(reports)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index ae7bf4d..aa61d28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,3 +113,7 @@ version = { attr = "tocode.__version__" } [tool.setuptools.packages.find] where = ["src"] + +# The agent skill ships inside the package (`tocode skill install`). +[tool.setuptools.package-data] +tocode = ["skill/*.md", "skill/references/*.md"] diff --git a/src/tocode/backends/dotnet.py b/src/tocode/backends/dotnet.py index 323fc08..11b950b 100644 --- a/src/tocode/backends/dotnet.py +++ b/src/tocode/backends/dotnet.py @@ -790,8 +790,11 @@ class DotnetSession: backend_label = ".NET" decompiler_label = "ICSharpCode.Decompiler + dnlib" - def __init__(self, found: DotnetInput) -> None: + def __init__(self, found: DotnetInput, *, in_memory: bool = False) -> None: self.input = found + # Load assemblies from bytes instead of mapping the files, so a live + # session can rewrite them (its patched copies) while they are open. + self.in_memory = in_memory self.assemblies: list[AssemblyInfo] = [] self.types: dict[tuple[int, int], TypeInfo] = {} self.methods: list[MethodInfo] = [] @@ -814,7 +817,12 @@ def load(self, progress: Any = None) -> None: if pe is None: continue try: - module = ModuleDefMD.Load(str(item.path)) + if self.in_memory: + from System.IO import File # type: ignore[import-not-found] + + module = ModuleDefMD.Load(File.ReadAllBytes(str(item.path))) + else: + module = ModuleDefMD.Load(str(item.path)) except Exception as exc: raise ToCodeError(f"dnlib cannot load {item.entry}: {exc}") from exc info = self._assembly_info(len(self.assemblies), item, pe, module) diff --git a/src/tocode/cli.py b/src/tocode/cli.py index d67769f..96bb0f9 100644 --- a/src/tocode/cli.py +++ b/src/tocode/cli.py @@ -258,7 +258,13 @@ def build_parser() -> argparse.ArgumentParser: return parser +SUBCOMMANDS = ("mcp", "sync", "skill") + + def main(argv: list[str] | None = None) -> int: + raw = list(sys.argv[1:] if argv is None else argv) + if raw and raw[0] in SUBCOMMANDS: + return _subcommand(raw[0], raw[1:]) parser = build_parser() args = parser.parse_args(argv) progress = Progress(enabled=not args.quiet) @@ -556,6 +562,23 @@ def ensure_dotnet() -> bool: return 130 if summary.interrupted else 0 +def _subcommand(name: str, argv: list[str]) -> int: + """`tocode mcp ...` (a file literally named like a subcommand: use ./name).""" + if name == "mcp": + from .live.server import main as mcp_main + + return mcp_main(argv) + if name == "sync": + from .live.setup import sync_main + + return sync_main(argv) + if name == "skill": + from .live.setup import skill_main + + return skill_main(argv) + raise AssertionError(name) + + def _ask_on_stderr(prompt: str) -> str: print(prompt, end="", file=sys.stderr, flush=True) line = sys.stdin.readline() diff --git a/src/tocode/dotnet.py b/src/tocode/dotnet.py index 204a090..a08fca4 100644 --- a/src/tocode/dotnet.py +++ b/src/tocode/dotnet.py @@ -652,34 +652,65 @@ def _store_result(context: DotnetContext, result: DecompiledType) -> None: record = session.types.get((result.assembly, result.token)) if record is None or record.cs_path is None or record.il_path is None: return + cs_text, il_text, cs_offset = type_texts(session, record, result) + if record.error is not None: + context.failures.append((record.full_name, record.error)) + context.progress.file_log( + f"export {info.name}!{record.full_name} failed: {record.error}" + ) + write_text_atomic(record.cs_path, cs_text) + write_text_atomic(record.il_path, il_text) + context.source_files.append(record.cs_path) + context.il_files.append(record.il_path) + assign_type_ranges(session, record, result, cs_text, il_text, cs_offset) + if record.error is None: + context.progress.file_log( + f"export {info.name}!{record.full_name} - {record.cs_lines} C# / " + f"{record.il_lines} IL lines done" + ) + + +def type_texts( + session: DotnetSession, record: TypeInfo, result: DecompiledType +) -> tuple[str, str, int]: + """A top-level type's .cs and .il file contents, and the C# line offset + of the decompiled text inside the .cs file (its header line).""" + info = session.assemblies[record.assembly] header = f"// ToCode: {record.full_name.replace('/', '+')} (0x{record.token:08x}) from {info.name} {info.version}\n" + record.error = record.il_error = None if result.cs is not None: cs_text = header + result.cs.rstrip("\n") + "\n" cs_offset = 1 else: record.error = result.cs_error or "unknown decompilation error" - context.failures.append((record.full_name, record.error)) cs_text = header + _failure_stub(session, record) cs_offset = 0 - context.progress.file_log( - f"export {info.name}!{record.full_name} failed: {record.error}" - ) if result.il is not None: il_text = header + result.il.rstrip("\n") + "\n" else: record.il_error = result.il_error or "unknown disassembly error" il_text = header + f"// IL disassembly failed: {record.il_error}\n" - write_text_atomic(record.cs_path, cs_text) - write_text_atomic(record.il_path, il_text) + return cs_text, il_text, cs_offset + + +def assign_type_ranges( + session: DotnetSession, + record: TypeInfo, + result: DecompiledType, + cs_text: str, + il_text: str, + cs_offset: int, +) -> None: + """Line counts and every method's C#/IL line range for a written type.""" record.cs_lines = cs_text.count("\n") record.il_lines = il_text.count("\n") - context.source_files.append(record.cs_path) - context.il_files.append(record.il_path) for nested in session.nested_types(record.assembly, record.token): nested.error = record.error nested.il_error = record.il_error nested.cs_lines = record.cs_lines nested.il_lines = record.il_lines + nested.cs_path = record.cs_path + nested.il_path = record.il_path for method_id in nested.methods: method = session.methods[method_id] span = result.cs_ranges.get(method.token) @@ -693,11 +724,6 @@ def _store_result(context: DotnetContext, result: DecompiledType) -> None: il_span = result.il_ranges.get(method.token) if il_span is not None: method.il_start, method.il_end = il_span[0] + 1, il_span[1] + 1 - if record.error is None: - context.progress.file_log( - f"export {info.name}!{record.full_name} - {record.cs_lines} C# / " - f"{record.il_lines} IL lines done" - ) def _store_assembly_attributes( diff --git a/src/tocode/dotnet_metadata.py b/src/tocode/dotnet_metadata.py index 436ea26..b275e03 100644 --- a/src/tocode/dotnet_metadata.py +++ b/src/tocode/dotnet_metadata.py @@ -879,6 +879,15 @@ def build_dotnet_agents( "- Cite file paths, type and method names, metadata tokens, and line numbers for every major claim.", "- Be explicit about uncertainty, decompilation failures (`types.json` `decompile_error`), and unresolved references.", "", + "## Live Database (MCP)", + "", + "With the `tocode` MCP server (`tocode mcp`), `open` this folder to query methods, IL, " + "cross-references, and strings live, and to record what you learn: `rename` / " + "`rename_many` (types and members), `rename_local`, `set_comment`. `sync` rewrites the " + "types those edits touched here (a renamed type's file moves). Edits are kept in " + "`.tocode/dotnet/edits.json`, and `patched/` holds the renamed assemblies. Record " + "findings as names and comments, not as edits to `src/raw/`: a sync regenerates those files.", + "", ] ) return "\n".join(lines) diff --git a/src/tocode/exporter.py b/src/tocode/exporter.py index c50f9ca..fcc02e6 100644 --- a/src/tocode/exporter.py +++ b/src/tocode/exporter.py @@ -143,16 +143,16 @@ def start(self, *, binary: Path, backend: str, addresses: list[int]) -> None: self._function_count = len(self.addresses) self._completed_indices.clear() - if self.restart and self.state_dir.exists(): - shutil.rmtree(self.state_dir) + if self.restart and self.state_path.exists(): + self._clear() self.progress.log("Checkpoint: discarded previous state (--restart)") existing = self._read_raw_state() - if existing is None and self.state_dir.exists(): - shutil.rmtree(self.state_dir) + if existing is None and self._has_state(): + self._clear() self.progress.log("Checkpoint: existing state is invalid; starting fresh") if existing is not None and existing.get("cache_id") != self.cache_id: - shutil.rmtree(self.state_dir) + self._clear() existing = None self.progress.log( "Checkpoint: existing state does not match this export; starting fresh" @@ -169,7 +169,7 @@ def start(self, *, binary: Path, backend: str, addresses: list[int]) -> None: existing.get("completed_ranges", []) ) else: - shutil.rmtree(self.state_dir) + self._clear() self.rendered_dir.mkdir(parents=True, exist_ok=True) self.clusters_dir.mkdir(parents=True, exist_ok=True) existing = None @@ -180,7 +180,7 @@ def start(self, *, binary: Path, backend: str, addresses: list[int]) -> None: self._completed_indices = self._scan_completed_indices() migrated = bool(self._completed_indices) else: - shutil.rmtree(self.state_dir) + self._clear() self.rendered_dir.mkdir(parents=True, exist_ok=True) self.clusters_dir.mkdir(parents=True, exist_ok=True) existing = None @@ -243,10 +243,31 @@ def mark_failed(self) -> None: self._write_state_payload(status="failed") def complete(self) -> None: - if self.state_dir.exists(): - shutil.rmtree(self.state_dir) + """Drop the checkpoint but keep the renders: `sync` rebuilds files from them.""" + from .syncstate import renders_dir + + if self.rendered_dir.is_dir(): + kept = renders_dir(self.root) + if kept.exists(): + shutil.rmtree(kept) + kept.parent.mkdir(parents=True, exist_ok=True) + self.rendered_dir.replace(kept) + self._clear() self.progress.log("Checkpoint: complete; removed saved state") + def _has_state(self) -> bool: + return any( + path.exists() + for path in (self.state_path, self.rendered_dir, self.clusters_dir) + ) + + def _clear(self) -> None: + """Remove the checkpoint's own files; `.tocode/` also holds sync state.""" + for path in (self.rendered_dir, self.clusters_dir): + if path.exists(): + shutil.rmtree(path) + self.state_path.unlink(missing_ok=True) + def completed_addresses(self) -> list[int]: return [self.addresses[index] for index in sorted(self._completed_indices)] @@ -639,9 +660,25 @@ def export_binary( raise if context.checkpoint is not None: context.checkpoint.complete() + write_sync_manifest(context) return _summary(context) +def write_sync_manifest(context: ExportContext) -> None: + """Record the file layout so `sync` can update this export in place.""" + from .syncstate import write_manifest + + write_manifest( + _need(context.root), + clusters=context.clusters, + backend=context.analyzer.backend_name, + decompiler=context.analyzer.decompiler_label, + tree=context.tree_enabled, + entropy=context.entropy_enabled, + origin=context.origin, + ) + + def _prepare_tree(context: ExportContext) -> None: root = _root_dir(context.analyzer.binary, context.out_dir) if context.progress.log_path is None: @@ -2474,7 +2511,14 @@ def _summary_function(routine: Routine, path: Path, text: str) -> str: return "\n".join(lines) -def _write_metadata(context: ExportContext) -> None: +def _write_metadata( + context: ExportContext, *, publish: bool = True, docs: bool = True +) -> None: + """Write every JSON document and header. + + ``sync`` passes ``publish=False`` (the database is already where it + belongs) and ``docs=False`` when it cannot rebuild AGENTS.md faithfully. + """ analysis = _need(context.analysis) root = _need(context.root) header = _need(context.header_path) @@ -2618,6 +2662,10 @@ def write_docs() -> None: ("AGENTS.md", write_docs), ("export-manifest.json", lambda: _set_manifest(context)), ] + skipped = ({"backend database"} if not publish else set()) | ( + {"AGENTS.md"} if not docs else set() + ) + steps = [(label, run) for label, run in steps if label not in skipped] context.progress.log(f"Writing metadata and indexes ({len(steps)} steps)") with context.progress.bar(total=len(steps), desc="metadata", unit="step") as bar: @@ -3018,6 +3066,20 @@ def build_export_agents( "- Cite file paths, function names, addresses, and line numbers for every major claim.", "- Be explicit about uncertainty and unresolved symbols.", "", + "## Live Database (MCP)", + "", + "This export can be updated from its database. With the `tocode` MCP server " + "(`tocode mcp`; register it with `claude mcp add --scope user tocode -- tocode mcp`), " + "`open` this folder, query the live decompiler (`decompile`, `xrefs`, `types`, ...), " + "record findings in the database (`rename`, `set_type`, `declare_types`, " + "`rename_local`, `set_comment`, or `eval` scripts), then `sync`: changed functions, " + "their callers, and code using changed types are re-rendered here, and only files " + "whose content changed are rewritten.", + "- Record findings as names, types, and comments in the database, not as edits to " + "`src/raw/`: a sync regenerates those files.", + "- After editing the database in the IDA or Binary Ninja UI, run " + "`tocode sync --full`.", + "", ] if tree_enabled: lines.insert( diff --git a/src/tocode/folder_docs.py b/src/tocode/folder_docs.py index e2369ac..c704a5f 100644 --- a/src/tocode/folder_docs.py +++ b/src/tocode/folder_docs.py @@ -163,6 +163,8 @@ def build_folder_agents( "- Follow a call across files with `mappings.json` or the `tocode-link` comments instead of searching by name.", "- A link whose `target_status` is not `decompiled` (filtered, runtime library, failed) still names the file and address; export it with `--only` to read it.", "- Cite file paths, function names, addresses, and line numbers for every major claim.", + "- With the `tocode` MCP server, `open` this folder and address each binary as " + '`session=""`; `sync` updates that binary\'s export and these links.', "", "## Binaries", "", diff --git a/src/tocode/live/__init__.py b/src/tocode/live/__init__.py new file mode 100644 index 0000000..afcd6c0 --- /dev/null +++ b/src/tocode/live/__init__.py @@ -0,0 +1,12 @@ +"""ToCode live: an MCP server over live decompiler sessions. + +``tocode mcp`` keeps IDA / Binary Ninja / r2 / angr sessions open in worker +processes (one per binary or database) and exposes them as MCP tools over +stdio: queries (decompile, disassemble, xrefs, ...), edits (rename, types, +comments), ``eval`` inside the database, and ``sync``, which updates an +exported project from the database incrementally. + +Modules: ``jsonrpc`` (stdio transport), ``targets`` (what ``open`` resolves +to), ``worker`` (the per-session process and its operations), ``sessions`` +(worker lifecycle), ``tools`` (the MCP tool table), ``server`` (entry point). +""" diff --git a/src/tocode/live/dotnet_edits.py b/src/tocode/live/dotnet_edits.py new file mode 100644 index 0000000..595972c --- /dev/null +++ b/src/tocode/live/dotnet_edits.py @@ -0,0 +1,523 @@ +"""Edits to .NET assemblies: the edit log and the edit tools. + +The log (``.tocode/dotnet/edits.json``) is the session's database: an ordered +list of edits addressed by assembly and metadata token, replayed onto the +original assemblies whenever a session opens. Renames of types, members, and +parameters change the dnlib modules (and so the patched copies); local +variable renames and comments are text edits applied when C# is rendered, +because neither lives in an assembly's metadata. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from pathlib import Path +import re +import time +from typing import TYPE_CHECKING, Any, Callable + +from ..metadata import write_text_atomic +from .worker import OpError + +if TYPE_CHECKING: + from .dotnet_live import DotnetLiveSession, Entity + +LOG_SCHEMA = 1 +_IDENT_RX = re.compile(r"^[A-Za-z_][\w`]*$") +_NAMESPACE_RX = re.compile(r"^[A-Za-z_][\w.]*$") +_UNSUPPORTED = { + "set_type": "types in .NET come from metadata; rename the type or member instead", + "declare_types": "types in .NET come from metadata; rename the type or member instead", + "set_local_type": "local variable types come from metadata; rename the local instead", + "define_function": "every .NET method is already defined in metadata", + "undefine": "removing .NET methods would break the assembly", +} + + +@dataclass +class EditLog: + path: Path + entries: list[dict[str, Any]] = field(default_factory=list) + synced: int = 0 # entries already reflected in the export + + @classmethod + def open(cls, path: Path) -> "EditLog": + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return cls(path=path) + if not isinstance(payload, dict) or payload.get("schema") != LOG_SCHEMA: + return cls(path=path) + entries = [ + entry for entry in payload.get("edits", []) if isinstance(entry, dict) + ] + return cls( + path=path, + entries=entries, + synced=min(int(payload.get("synced") or 0), len(entries)), + ) + + def save(self) -> None: + write_text_atomic( + self.path, + json.dumps( + {"schema": LOG_SCHEMA, "synced": self.synced, "edits": self.entries}, + indent=1, + ) + + "\n", + ) + + def local_renames(self, assembly: str) -> dict[int, list[tuple[str, str]]]: + renames: dict[int, list[tuple[str, str]]] = {} + for entry in self.entries: + if entry.get("op") == "rename_local" and entry.get("assembly") == assembly: + renames.setdefault(int(entry["token"], 16), []).append( + (entry["old"], entry["name"]) + ) + return renames + + def comments(self, assembly: str) -> dict[int, str]: + comments: dict[int, str] = {} + for entry in self.entries: + if entry.get("op") == "comment" and entry.get("assembly") == assembly: + token = int(entry["token"], 16) + if entry.get("text"): + comments[token] = str(entry["text"]) + else: + comments.pop(token, None) + return comments + + +# -------------------------------------------------------------------------- +# Tools +# -------------------------------------------------------------------------- + + +def register(live: "DotnetLiveSession") -> None: + ops = { + "rename": lambda args: _rename_tool(live, args), + "rename_many": lambda args: _rename_many(live, args), + "rename_local": lambda args: _rename_local_tool(live, args), + "set_comment": lambda args: _comment_tool(live, args), + } + for name, reason in _UNSUPPORTED.items(): + ops[name] = _refuser(name, reason) + live.ops.update(ops) + live.edit_ops = set(ops) + + +def _refuser(name: str, reason: str) -> Callable[[dict[str, Any]], Any]: + def refuse(_args: dict[str, Any]) -> Any: + raise OpError(f"{name} is not available for .NET: {reason}") + + return refuse + + +def _allowed(live: "DotnetLiveSession") -> None: + if not live.allow_edits: + raise OpError("edits are disabled for this server (--read-only)") + + +def _record(live: "DotnetLiveSession", entry: dict[str, Any]) -> list[str]: + """Apply an edit, append it to the log, and save the log.""" + dirty = apply(live, entry) + entry["time"] = time.strftime("%Y-%m-%dT%H:%M:%S") + live.edits.entries.append(entry) + live.edits.save() + live.dirty.update(dirty) + live._edited = True + return sorted(live.type_name(key) for key in dirty) + + +def _rename_tool(live: "DotnetLiveSession", args: dict[str, Any]) -> dict[str, Any]: + _allowed(live) + entity = live.resolve(args.get("target")) + name = str(args.get("name") or "").strip() + old = live.entity_name(entity) + pending = _record(live, _rename_entry(live, entity, name)) + return { + "address": entity.address(live.session), + "kind": entity.kind, + "old": old, + "name": name, + "pending_sync": pending, + } + + +def _rename_entry( + live: "DotnetLiveSession", entity: "Entity", name: str +) -> dict[str, Any]: + if entity.kind == "type": + if not _NAMESPACE_RX.match(name): + raise OpError( + f"{name!r} is not a valid type name (`Name` or `Namespace.Name`)" + ) + elif not _IDENT_RX.match(name): + raise OpError(f"{name!r} is not a valid identifier") + return { + "op": "rename", + "kind": entity.kind, + "assembly": live.session.assemblies[entity.assembly].name, + "token": f"0x{entity.token:08x}", + "old": live.entity_name(entity), + "name": name, + } + + +def _rename_many(live: "DotnetLiveSession", args: dict[str, Any]) -> dict[str, Any]: + _allowed(live) + renames = args.get("renames") or [] + if not isinstance(renames, list) or not renames: + raise OpError("rename_many needs `renames`: a list of {target, name}") + done: list[dict[str, Any]] = [] + errors: list[dict[str, Any]] = [] + pending: set[str] = set() + # Resolve every target first: names in the batch refer to the state + # before it (renaming a type must not break its members' targets). + resolved: list[tuple[dict[str, Any], Any]] = [] + for item in renames: + try: + resolved.append((item, live.resolve(item.get("target")))) + except OpError as exc: + errors.append({"target": item.get("target"), "error": str(exc)}) + for item, entity in resolved: + try: + name = str(item.get("name") or "").strip() + pending.update(_record(live, _rename_entry(live, entity, name))) + done.append( + { + "target": item.get("target"), + "address": entity.address(live.session), + "name": name, + } + ) + except OpError as exc: + errors.append({"target": item.get("target"), "error": str(exc)}) + return {"renamed": len(done), "errors": errors, "pending_sync": sorted(pending)} + + +def _rename_local_tool( + live: "DotnetLiveSession", args: dict[str, Any] +) -> dict[str, Any]: + _allowed(live) + entity = live.resolve(args.get("function"), kinds=("method",)) + method = live.method_of(entity) + old = str(args.get("old") or "") + new = str(args.get("new") or "").strip() + if not old or not _IDENT_RX.match(new): + raise OpError("rename_local needs `old` and a valid `new` name") + assembly = live.session.assemblies[entity.assembly].name + params = [name for name, _ in method.params] + if old in params: + # Parameter names live in metadata: rename them in the assembly. + entry = { + "op": "rename_param", + "assembly": assembly, + "token": f"0x{entity.token:08x}", + "index": params.index(old), + "old": old, + "name": new, + } + else: + if not re.search(rf"\b{re.escape(old)}\b", live.csharp(entity)): + raise OpError( + f"no local variable {old!r} in {live.entity_name(entity)} (names as the decompiled C# shows them)" + ) + entry = { + "op": "rename_local", + "assembly": assembly, + "token": f"0x{entity.token:08x}", + "old": old, + "name": new, + } + pending = _record(live, entry) + return { + "function": live.entity_name(entity), + "old": old, + "new": new, + "pending_sync": pending, + } + + +def _comment_tool(live: "DotnetLiveSession", args: dict[str, Any]) -> dict[str, Any]: + _allowed(live) + entity = live.resolve(args.get("address")) + text = str(args.get("text") or "") + entry = { + "op": "comment", + "assembly": live.session.assemblies[entity.assembly].name, + "token": f"0x{entity.token:08x}", + "text": text, + } + pending = _record(live, entry) + return {"target": live.entity_name(entity), "pending_sync": pending} + + +# -------------------------------------------------------------------------- +# Applying edits +# -------------------------------------------------------------------------- + + +def replay(live: "DotnetLiveSession") -> None: + """Apply the whole log to freshly loaded assemblies.""" + for index, entry in enumerate(live.edits.entries): + try: + dirty = apply(live, entry) + except OpError as exc: + live.log(f"live: skipping edit {index} ({entry.get('op')}): {exc}") + continue + if index >= live.edits.synced: + live.dirty.update(dirty) + + +def apply(live: "DotnetLiveSession", entry: dict[str, Any]) -> set[tuple[int, int]]: + """Apply one edit to the loaded assemblies; return the types it makes stale.""" + from .dotnet_live import TOKEN_KINDS, Entity + + op = entry.get("op") + assembly = live._assembly_index(str(entry.get("assembly"))) + token = int(str(entry.get("token")), 16) + kind = TOKEN_KINDS.get(token >> 24, "") + entity = Entity(kind, assembly, token) + if op in ("comment", "rename_local"): + return {live.top_level(entity)} + module = live.session._modules[assembly] + rid = token & 0xFFFFFF + name = str(entry.get("name") or "") + if op == "rename_param": + method = module.ResolveMethod(rid) + params = [p for p in method.Parameters if not p.IsHiddenThisParameter] + index = int(entry.get("index") or 0) + if index >= len(params): + raise OpError(f"no parameter {index} in {method.FullName}") + params[index].Name = name + record = live.session.method(assembly, token) + if record is not None and index < len(record.params): + record.params[index] = (name, record.params[index][1]) + dirty = {live.top_level(entity)} + elif op == "rename": + dirty = _affected(live, entity) + from dnlib.DotNet import UTF8String # type: ignore[import-not-found] + + if kind == "type": + definition = module.ResolveTypeDef(rid) + if "." in name and definition.DeclaringType is None: + namespace, _, short = name.rpartition(".") + definition.Namespace = UTF8String(namespace) + definition.Name = UTF8String(short) + else: + definition.Name = UTF8String(name) + else: + resolver = { + "method": module.ResolveMethod, + "field": module.ResolveField, + "property": module.ResolveProperty, + "event": module.ResolveEvent, + }.get(kind) + if resolver is None: + raise OpError(f"cannot rename a {kind or 'token 0x%08x' % token}") + definition = resolver(rid) + if definition is None: + raise OpError( + f"no {kind} 0x{token:08x} in {live.session.assemblies[assembly].name}" + ) + definition.Name = UTF8String(name) + _rename_in_inventory(live, entity, module) + else: + raise OpError(f"unknown edit {op}") + live.unwritten.add(assembly) + live.reload_needed = True + return dirty + + +def _rename_in_inventory( + live: "DotnetLiveSession", entity: "Entity", module: Any +) -> None: + """Bring the inventory records and the name index up to date with a + rename, touching only the renamed entity (and, for a type, its nested + types and their members).""" + from .dotnet_live import Entity + + session = live.session + asm = entity.assembly + if entity.kind == "type": + top = live.top_level(entity) + records = list(session.nested_types(top[0], top[1])) + entities = [Entity("type", asm, r.token) for r in records] + for record in records: + entities += [ + Entity("method", asm, session.methods[i].token) for i in record.methods + ] + entities += [Entity("field", asm, m.token) for m in record.fields] + entities += [Entity("property", asm, m.token) for m in record.properties] + entities += [Entity("event", asm, m.token) for m in record.events] + for item in entities: + live.index_entity(item, remove=True) + for record in records: + definition = module.ResolveTypeDef(record.token & 0xFFFFFF) + record.name = str(definition.Name) + record.namespace = str(definition.Namespace) + record.full_name = str(definition.FullName) + record.reflection_name = str(definition.ReflectionFullName) + for method_id in record.methods: + method = session.methods[method_id] + method_def = module.ResolveMethod(method.token & 0xFFFFFF) + method.type_name = record.full_name + method.full_name = str(method_def.FullName) + for item in entities: + live.index_entity(item) + return + live.index_entity(entity, remove=True) + if entity.kind == "method": + method = live.method_of(entity) + method_def = module.ResolveMethod(entity.token & 0xFFFFFF) + method.name = str(method_def.Name) + method.full_name = str(method_def.FullName) + else: + record = live.type_of(entity) + for member in record.fields + record.properties + record.events: + if member.token == entity.token: + resolver = { + "field": module.ResolveField, + "property": module.ResolveProperty, + "event": module.ResolveEvent, + }[entity.kind] + member.name = str(resolver(entity.token & 0xFFFFFF).Name) + live.index_entity(entity) + + +def _affected(live: "DotnetLiveSession", entity: "Entity") -> set[tuple[int, int]]: + """Types whose C# changes when ``entity`` is renamed (computed before), + from the session's reverse indexes: no scan over all methods.""" + session = live.session + asm = entity.assembly + dirty = {live.top_level(entity)} + callers: list[int] = [] + accessors: tuple[Any, ...] = () + if entity.kind == "method": + callers = list(live.method_of(entity).callers) + elif entity.kind in ("field", "property", "event"): + record = live.type_of(entity) + users = set(live.field_users.get((asm, entity.token), ())) + for prop in record.properties: + if prop.token == entity.token: + backing = next( + ( + f + for f in record.fields + if f.name == f"<{prop.name}>k__BackingField" + ), + None, + ) + if backing is not None: + users |= live.field_users.get((asm, backing.token), set()) + accessors = (prop.getter, prop.setter) + break + else: + accessors = next( + ((e.add, e.remove) for e in record.events if e.token == entity.token), + (), + ) + for token in accessors: + method = session.method(asm, token) if token else None + if method is not None: + callers += method.callers + callers += list(users) + else: + # A type: every exported file that mentions its name (word index). + return dirty | words(live).get(live.type_of(entity).name.split("`")[0], set()) + for method_id in callers: + method = session.methods[method_id] + owner = session.types.get((method.assembly, method.type_token)) + if owner is not None: + dirty.add((method.assembly, owner.top_level or owner.token)) + return dirty + + +_WORD_RX = re.compile(r"[A-Za-z_]\w*") + + +def words(live: "DotnetLiveSession") -> dict[str, set[tuple[int, int]]]: + """Identifier -> exported types whose C# file contains it. + + Built once from the export's files and kept until the next sync rewrites + them, so a type rename finds its users with one lookup. + """ + cached = live._words + if cached is not None: + return cached + from .dotnet_sync import disk_layout + + index: dict[str, set[tuple[int, int]]] = {} + for key, row in disk_layout(live).types.items(): + record = live.session.types.get(key) + if record is None or record.declaring_type is not None: + continue # nested types share their top-level type's file + try: + text = Path(row["cs_path"]).read_text(encoding="utf-8", errors="replace") + except OSError: + continue + for word in set(_WORD_RX.findall(text)): + index.setdefault(word, set()).add(key) + live._words = index + return index + + +def apply_text_edits( + live: "DotnetLiveSession", + assembly: int, + text: str, + ranges: dict[int, tuple[int, int]], + offset: int, +) -> tuple[str, dict[int, tuple[int, int]]]: + """Local renames and comments, applied to decompiled C#. + + ``ranges`` maps member tokens to their (first, last) lines in ``text`` + minus ``offset``; the returned ranges follow the inserted comment lines. + """ + name = live.session.assemblies[assembly].name + renames = live.edits.local_renames(name) + comments = live.edits.comments(name) + if not renames and not comments: + return text, ranges + lines = text.split("\n") + for token, pairs in renames.items(): + span = ranges.get(token) + if span is None: + continue + start, end = span[0] + offset - 1, span[1] + offset + for index in range(max(0, start), min(end, len(lines))): + for old, new in pairs: + lines[index] = re.sub(rf"\b{re.escape(old)}\b", new, lines[index]) + inserts = sorted( + ( + (ranges[token][0] + offset, comment) + for token, comment in comments.items() + if token in ranges + ), + reverse=True, + ) + for line_number, comment in inserts: + position = max(0, line_number - 1) + anchor = lines[position] if position < len(lines) else "" + indent = anchor[: len(anchor) - len(anchor.lstrip())] + lines[position:position] = [ + f"{indent}// {part}" for part in comment.splitlines() or [""] + ] + points = sorted( + (line, len(comment.splitlines() or [""])) for line, comment in inserts + ) + + def before(line: int, inclusive: bool) -> int: + return sum( + size + for point, size in points + if point < line or (inclusive and point <= line) + ) + + moved = { + token: (start + before(start + offset, False), end + before(end + offset, True)) + for token, (start, end) in ranges.items() + } + return "\n".join(lines), moved diff --git a/src/tocode/live/dotnet_live.py b/src/tocode/live/dotnet_live.py new file mode 100644 index 0000000..d8b0b4f --- /dev/null +++ b/src/tocode/live/dotnet_live.py @@ -0,0 +1,951 @@ +"""The live session for .NET: assemblies through dnlib and ICSharpCode. + +The "database" of a .NET session is an edit log (renames and comments, +addressed by metadata token) plus renamed copies of the assemblies written +from it (``/patched/``). Opening loads the original input, replays the +log, and decompiles from the patched copies; the originals are never touched. + +Cost model: an edit changes the dnlib objects and updates only the affected +inventory records and name-index entries. A patched copy is written only when +a query needs decompiled code from an edited assembly; the inventory is +reloaded once per ``sync`` (which rewrites the JSON documents anyway). +""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +import os +from pathlib import Path +import re +import shutil +import tempfile +from typing import Any, Callable + +from ..progress import Progress +from .targets import Target +from .worker import ( + MAX_READ, + OpError, + _clip, + _jsonable, + _matcher, + _page, + run_python, +) + +TOKEN_KINDS = { + 0x02: "type", + 0x04: "field", + 0x06: "method", + 0x14: "event", + 0x17: "property", +} +_ADDRESS_RX = re.compile(r"^(?:(?P[^!]+)!)?0x(?P[0-9a-fA-F]{8})$") + + +@dataclass(frozen=True, slots=True) +class Entity: + kind: str # type, method, field, property, event + assembly: int + token: int + + def address(self, session: Any) -> str: + return f"{session.assemblies[self.assembly].name}!0x{self.token:08x}" + + +class DotnetLiveSession: + backend = "dotnet" + + def __init__( + self, + target: Target, + *, + log_path: Path, + allow_eval: bool = True, + allow_edits: bool = True, + ) -> None: + from ..backends.dotnet import ( + DotnetSession, + discover_dotnet_input, + init_dotnet_worker, + load_runtime, + ) + from . import dotnet_edits + + self.target = target + self.allow_eval = allow_eval + self.allow_edits = allow_edits + self.progress = Progress(enabled=False, log_path=log_path) + self.log(f"live: opening .NET input {target.open_path}") + load_runtime() + self.workdir = Path(tempfile.mkdtemp(prefix="tocode-live-dotnet-")) + include_framework = bool(self._project().get("include_framework")) + self.found = discover_dotnet_input( + target.open_path, workdir=self.workdir, include_framework=include_framework + ) + self.original_paths = [item.path for item in self.found.included] + self.session: Any = DotnetSession(self.found, in_memory=True) + self.session.load() + init_dotnet_worker([str(path) for path in self.found.search_dirs]) + self.eval_namespace: dict[str, Any] | None = None + # Top-level types, as (assembly index, type token), the next sync must + # write again; assemblies whose patched copy is behind the edits; and + # whether the inventory's derived strings (references) are behind. + self.dirty: set[tuple[int, int]] = set() + self.unwritten: set[int] = set() + self.reload_needed = False + self._edited = False + # Caches of the export on disk, dropped when a sync rewrites it. + self._locations: dict[str, dict[str, Any]] | None = None + self._layout: Any = None + self._words: dict[str, set[tuple[int, int]]] | None = None + self.edits = dotnet_edits.EditLog.open(self.edit_log_path()) + self.ops: dict[str, Callable[[dict[str, Any]], Any]] = { + "info": lambda _args: self.info(), + "functions": self.op_functions, + "function": self.op_function, + "decompile": self.op_decompile, + "disassemble": self.op_disassemble, + "xrefs": self.op_xrefs, + "strings": self.op_strings, + "imports": self.op_imports, + "exports": self.op_exports, + "sections": self.op_sections, + "types": self.op_types, + "type": self.op_type, + "read_bytes": self.op_read_bytes, + "search": self.op_search, + "lookup": self.op_lookup, + "eval": self.op_eval, + "save": lambda _args: self.save(), + "status": lambda _args: self.status(), + } + self.edit_ops: set[str] = set() + self._index() + dotnet_edits.register(self) + from . import dotnet_sync + + dotnet_sync.register(self) + dotnet_edits.replay(self) + self.log("live: .NET session ready") + + # -- plumbing --------------------------------------------------------------- + + def log(self, message: str) -> None: + self.progress.file_log(message) + + def dispatch(self, op: str, args: dict[str, Any]) -> Any: + handler = self.ops.get(op) + if handler is None: + raise OpError(f"unknown operation {op}") + return handler(args) + + def _project(self) -> dict[str, Any]: + if self.target.export_dir is None: + return {} + try: + return json.loads( + (self.target.export_dir / "project.json").read_text("utf-8") + ) + except (OSError, ValueError): + return {} + + def edit_log_path(self) -> Path: + if self.target.database is not None: + return self.target.database + if self.target.export_dir is not None: + return self.target.export_dir / ".tocode" / "dotnet" / "edits.json" + base = Path.home() / ".cache" / "tocode" / "live" / "dotnet" + return base / f"{self.target.binary.name}.edits.json" + + def patched_path(self, assembly: int) -> Path: + base = self.target.export_dir or self.edit_log_path().parent + info = self.session.assemblies[assembly] + return base / "patched" / (info.folder or info.name) / Path(info.file.path).name + + def assembly_path(self, assembly: int) -> str: + """The file the decompiler reads: the patched copy once there is one.""" + patched = self.patched_path(assembly) + if self.edits.entries and patched.is_file(): + return str(patched) + return str(self.original_paths[assembly]) + + def info(self) -> dict[str, Any]: + session = self.session + return { + **self.target.describe(), + "backend": "dotnet", + "backend_label": ".NET", + "decompiler": session.decompiler_label, + "kind": self.found.kind, + "assemblies": [ + { + "name": item.name, + "version": item.version, + "framework": item.target_framework, + } + for item in session.assemblies + ], + "functions": len(session.methods), + "types": len(session.types), + "strings": sum(len(item.user_strings) for item in session.assemblies), + "edits": len(self.edits.entries), + "eval": self.allow_eval, + "edits_allowed": self.allow_edits, + "edit_log": str(self.edit_log_path()), + } + + def status(self) -> dict[str, Any]: + return { + "pending_types": sorted(self.type_name(key) for key in self.dirty), + "pending_edits": len(self.edits.entries) - self.edits.synced, + } + + def save(self) -> dict[str, Any]: + """Write the edit log and the patched assemblies.""" + self.edits.save() + written = self.write_patched() + self._edited = False + return {"saved": True, "edit_log": str(self.edits.path), "patched": written} + + def close(self, *, save: bool) -> dict[str, Any]: + result: dict[str, Any] = {"saved": False} + if save and (self._edited or self.unwritten): + result = self.save() + try: + self.session.close() + finally: + shutil.rmtree(self.workdir, ignore_errors=True) + return result + + # -- patched assemblies --------------------------------------------------------- + + def write_patched(self) -> list[str]: + """Write the patched copies of edited assemblies (only those).""" + from ..backends.dotnet import _worker + + written: list[str] = [] + for index in sorted(self.unwritten): + target = self.patched_path(index) + target.parent.mkdir(parents=True, exist_ok=True) + cached = _worker.get("decompilers", {}).pop(str(target), None) + if cached is not None: + try: + cached[1].Dispose() + except Exception: # noqa: BLE001 + pass + write_module(self.session._modules[index], target) + written.append(str(target)) + self.unwritten.clear() + return written + + def reload(self) -> list[str]: + """Write pending copies and load the inventory from them (for sync). + + Renames are already reflected in names; what a reload adds is the + derived strings (reference lists, signatures) the JSON documents + print, so it runs once per sync rather than after every edit. + """ + from ..backends.dotnet import DotnetSession + + written = self.write_patched() + if not self.reload_needed: + return written + for index in range(len(self.found.included)): + patched = self.patched_path(index) + if patched.is_file(): + self.found.included[index].path = patched + session = DotnetSession(self.found, in_memory=True) + session.load() + old = self.session + self.session = session + old.close() + self.reload_needed = False + self.eval_namespace = None + self._index() + return written + + # -- names ---------------------------------------------------------------------- + + def _index(self) -> None: + """Name lookups and reverse references, built once per inventory.""" + from ..dotnet import _assembly_folders + + session = self.session + _assembly_folders(session) + self.names: dict[str, list[Entity]] = {} + self.lower: dict[str, list[Entity]] = {} + # (assembly, member token) -> (assembly, declaring type token) + self.member_owner: dict[tuple[int, int], tuple[int, int]] = {} + # (assembly, field token) -> methods (ids) that read or write it + self.field_users: dict[tuple[int, int], set[int]] = {} + field_by_name: dict[tuple[int, str], int] = {} + for (assembly, token), record in session.types.items(): + self.index_entity(Entity("type", assembly, token)) + for kind, members in ( + ("field", record.fields), + ("property", record.properties), + ("event", record.events), + ): + for member in members: + self.member_owner[(assembly, member.token)] = (assembly, token) + self.index_entity(Entity(kind, assembly, member.token)) + for item in record.fields: + field_by_name[(assembly, f"{record.full_name}::{item.name}")] = ( + item.token + ) + for method in session.methods: + self.index_entity(Entity("method", method.assembly, method.token)) + for ref in method.field_refs: + token = field_by_name.get((method.assembly, ref)) + if token is not None: + self.field_users.setdefault((method.assembly, token), set()).add( + method.id + ) + + def entity_keys(self, entity: Entity) -> list[str]: + """The names an entity can be looked up by.""" + if entity.kind == "type": + record = self.session.types.get((entity.assembly, entity.token)) + return ( + [record.full_name, record.reflection_name, record.name] + if record + else [] + ) + if entity.kind == "method": + method = self.session.method(entity.assembly, entity.token) + if method is None: + return [] + owner = method.type_name.replace("/", "+") + return [ + f"{owner}::{method.name}", + f"{owner}.{method.name}", + method.name, + method.full_name, + ] + owner_key = self.member_owner.get((entity.assembly, entity.token)) + record = self.session.types.get(owner_key) if owner_key else None + if record is None: + return [] + member = next( + ( + m + for m in record.fields + record.properties + record.events + if m.token == entity.token + ), + None, + ) + if member is None: + return [] + owner = record.reflection_name + return [f"{owner}::{member.name}", f"{owner}.{member.name}", member.name] + + def index_entity(self, entity: Entity, *, remove: bool = False) -> None: + for name in self.entity_keys(entity): + for table, key in ((self.names, name), (self.lower, name.lower())): + bucket = table.setdefault(key, []) + if remove: + if entity in bucket: + bucket.remove(entity) + if not bucket: + del table[key] + elif entity not in bucket: + bucket.append(entity) + + def resolve(self, value: Any, *, kinds: tuple[str, ...] | None = None) -> Entity: + if not isinstance(value, str) or not value.strip(): + raise OpError( + "target must be `Assembly!0x06000012`, `Namespace.Type`, `Type::Member`, or a name" + ) + text = value.strip() + match = _ADDRESS_RX.match(text) + if match is not None: + token = int(match.group("token"), 16) + kind = TOKEN_KINDS.get(token >> 24) + if kind is None: + raise OpError( + f"0x{token:08x} is not a type, method, field, property, or event token" + ) + assembly = self._assembly_index(match.group("asm")) + return self._check(Entity(kind, assembly, token), kinds) + found = list(self.names.get(text) or self.lower.get(text.lower()) or []) + if kinds is not None: + found = [entity for entity in found if entity.kind in kinds] + if not found: + raise OpError(f"no type or member named {text!r}; try `lookup`") + if len(found) > 1: + choices = ", ".join( + f"{e.kind} {e.address(self.session)} ({self.entity_name(e)})" + for e in found[:12] + ) + raise OpError(f"{text!r} is ambiguous; use an address: {choices}") + return found[0] + + def _check(self, entity: Entity, kinds: tuple[str, ...] | None) -> Entity: + if kinds is not None and entity.kind not in kinds: + raise OpError( + f"{entity.address(self.session)} is a {entity.kind}, expected {' or '.join(kinds)}" + ) + return entity + + def _assembly_index(self, name: str | None) -> int: + assemblies = self.session.assemblies + if name is None: + if len(assemblies) == 1: + return 0 + raise OpError( + "several assemblies are loaded; prefix the token with `Assembly!`" + ) + for info in assemblies: + if info.name == name or info.name.lower() == name.lower(): + return int(info.index) + raise OpError( + f"no assembly named {name!r} (loaded: {', '.join(a.name for a in assemblies)})" + ) + + def method_of(self, entity: Entity) -> Any: + method = self.session.method(entity.assembly, entity.token) + if method is None: + raise OpError(f"no method {entity.address(self.session)}") + return method + + def type_of(self, entity: Entity) -> Any: + """The type record that declares an entity (the type itself for a type).""" + session = self.session + if entity.kind == "type": + record = session.types.get((entity.assembly, entity.token)) + elif entity.kind == "method": + method = self.method_of(entity) + record = session.types.get((entity.assembly, method.type_token)) + else: + owner = self.member_owner.get((entity.assembly, entity.token)) + record = session.types.get(owner) if owner else None + if record is None: + raise OpError(f"no declaring type for {entity.address(session)}") + return record + + def top_level(self, entity: Entity) -> tuple[int, int]: + record = self.type_of(entity) + return (entity.assembly, record.top_level or record.token) + + def entity_name(self, entity: Entity) -> str: + if entity.kind == "type": + return str(self.type_of(entity).reflection_name) + if entity.kind == "method": + method = self.method_of(entity) + return f"{method.type_name.replace('/', '+')}::{method.name}" + record = self.type_of(entity) + for member in record.fields + record.properties + record.events: + if member.token == entity.token: + return f"{record.reflection_name}::{member.name}" + return entity.address(self.session) + + def type_name(self, key: tuple[int, int]) -> str: + record = self.session.types.get(key) + return str(record.reflection_name) if record is not None else f"0x{key[1]:08x}" + + def touch(self, *targets: Any) -> None: + for value in targets: + self.dirty.add(self.top_level(self.resolve(value))) + + # -- decompiling ------------------------------------------------------------------ + + def decompiler(self, assembly: int) -> tuple[Any, Any, Any]: + from ..backends.dotnet import _decompiler_for + + if assembly in self.unwritten: + self.write_patched() + return _decompiler_for(self.assembly_path(assembly)) + + def csharp(self, entity: Entity) -> str: + """C# for one type or member, with this session's text edits applied.""" + from ..backends.dotnet import _decompile_type + from System.Reflection.Metadata.Ecma335 import MetadataTokens # type: ignore[import-not-found] + + from . import dotnet_edits + + decompiler, _pe, settings = self.decompiler(entity.assembly) + if entity.kind == "type": + text, ranges = _decompile_type(decompiler, settings, entity.token) + else: + handle = MetadataTokens.EntityHandle(int(entity.token)) + text = str(decompiler.DecompileAsString(handle)).replace("\r\n", "\n") + ranges = {entity.token: (1, text.count("\n") + 1)} + text, _ = dotnet_edits.apply_text_edits(self, entity.assembly, text, ranges, 0) + return text + + def il(self, entity: Entity) -> str: + from ..backends.dotnet import _disassemble_type, _new_disassembler + from ICSharpCode.Decompiler import PlainTextOutput # type: ignore[import-not-found] + from System.Reflection.Metadata.Ecma335 import MetadataTokens # type: ignore[import-not-found] + + _decompiler, pe, _settings = self.decompiler(entity.assembly) + if entity.kind == "type": + return _disassemble_type(pe, entity.token)[0] + if entity.kind != "method": + raise OpError("disassemble works on a method or a type") + output = PlainTextOutput() + output.IndentationString = " " + _new_disassembler(output).DisassembleMethod( + pe, MetadataTokens.MethodDefinitionHandle(entity.token & 0xFFFFFF) + ) + return str(output.ToString()).replace("\r\n", "\n") + + # -- queries ---------------------------------------------------------------------- + + def _location(self, method: Any) -> dict[str, Any] | None: + locations = self._locations + if locations is None: + locations = self._locations = _load_dotnet_locations(self.target.export_dir) + return locations.get(method.address) + + def _method_row(self, method: Any) -> dict[str, Any]: + return { + "address": method.address, + "name": f"{method.type_name.replace('/', '+')}::{method.name}", + "prototype": method.prototype, + "size": method.il_size, + "kind": "generated" + if method.compiler_generated + else "pinvoke" + if method.pinvoke_module + else "constructor" + if method.is_constructor + else "method", + "callers": len(method.callers), + "callees": len(method.callees) + len(method.external_calls), + } + + def op_functions(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + kind = args.get("kind") or "all" + rows = [] + for method in self.session.methods: + row = self._method_row(method) + if kind == "app" and row["kind"] == "generated": + continue + if kind not in ("all", "app") and row["kind"] != kind: + continue + if pattern is not None and not pattern(row["name"]): + continue + rows.append(row) + return _page(rows, args) + + def op_function(self, args: dict[str, Any]) -> dict[str, Any]: + method = self.method_of(self.resolve(args.get("target"), kinds=("method",))) + methods = self.session.methods + row = self._method_row(method) + row.update( + { + "return_type": method.return_type, + "params": [ + {"name": name, "type": kind} for name, kind in method.params + ], + "flags": method.flags, + "static": method.is_static, + "virtual": method.is_virtual, + "callers": [ + self._method_row(methods[i])["name"] for i in method.callers + ], + "callees": [ + self._method_row(methods[i])["name"] for i in method.callees + ], + "external_calls": method.external_calls, + "fields": method.field_refs, + "strings": method.string_ref_count, + "pinvoke": f"{method.pinvoke_module}!{method.pinvoke_entry}" + if method.pinvoke_module + else None, + "owner": self._method_row(methods[method.owner_method])["name"] + if method.owner_method is not None + else None, + "export": self._location(method), + } + ) + return row + + def op_decompile(self, args: dict[str, Any]) -> dict[str, Any]: + entity = self.resolve(args.get("target"), kinds=("type", "method")) + return { + "address": entity.address(self.session), + "name": self.entity_name(entity), + "code": _clip(self.csharp(entity)), + } + + def op_disassemble(self, args: dict[str, Any]) -> dict[str, Any]: + entity = self.resolve(args.get("target"), kinds=("type", "method")) + text = self.il(entity) + limit = int(args.get("max_lines") or 0) + if limit > 0: + lines = text.splitlines() + text = "\n".join(lines[:limit]) + ( + f"\n... {len(lines) - limit} more lines" if len(lines) > limit else "" + ) + return { + "address": entity.address(self.session), + "name": self.entity_name(entity), + "disassembly": _clip(text), + } + + def op_xrefs(self, args: dict[str, Any]) -> dict[str, Any]: + entity = self.resolve(args.get("target")) + direction = args.get("direction") or "to" + methods = self.session.methods + rows: list[dict[str, Any]] = [] + if entity.kind == "method": + method = self.method_of(entity) + if direction == "from": + rows += [ + { + "to": methods[i].address, + "name": self._method_row(methods[i])["name"], + "kind": "call", + } + for i in method.callees + ] + rows += [ + {"to": None, "name": name, "kind": "external call"} + for name in method.external_calls + ] + rows += [ + {"to": None, "name": name, "kind": "field"} + for name in method.field_refs + ] + else: + rows += [ + { + "from": methods[i].address, + "function": self._method_row(methods[i])["name"], + "kind": "call", + } + for i in method.callers + ] + else: + if direction == "from": + raise OpError("xrefs from works on a method") + name = self.entity_name(entity) + short = name.split("::")[-1] + record = self.type_of(entity) + for method in methods: + if method.assembly != entity.assembly: + continue + if entity.kind == "type": + uses = method.type_token != entity.token and ( + record.reflection_name in method.prototype.replace("/", "+") + or any(record.full_name in ref for ref in method.field_refs) + or any( + methods[i].type_token == entity.token + for i in method.callees + ) + ) + else: + uses = any( + ref.endswith(f"::{short}") and record.full_name in ref + for ref in method.field_refs + ) + if uses: + rows.append( + { + "from": method.address, + "function": self._method_row(method)["name"], + "kind": "use", + } + ) + return { + "target": entity.address(self.session), + "direction": direction, + "refs": rows, + } + + def op_strings(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + methods = self.session.methods + rows = [] + for info in self.session.assemblies: + for offset, value in info.user_strings: + if pattern is not None and not pattern(value): + continue + users = [ + self._method_row(methods[i])["name"] + for i in self.session.string_xrefs.get((info.index, offset), ()) + ] + rows.append( + { + "address": f"{info.name}!0x{0x70000000 | offset:08x}", + "value": value, + "used_by": users, + } + ) + return _page(rows, args) + + def op_imports(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + rows = [] + for scope, members in sorted(self.session.external_refs.items()): + for member, users in sorted(members.items()): + if pattern is not None and not pattern(f"{scope} {member}"): + continue + rows.append({"library": scope, "name": member, "used_by": len(users)}) + for method in self.session.methods: + if method.pinvoke_module and ( + pattern is None or pattern(str(method.pinvoke_entry)) + ): + rows.append( + { + "library": method.pinvoke_module, + "name": method.pinvoke_entry, + "pinvoke": method.address, + } + ) + return _page(rows, args) + + def op_exports(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + rows = [ + { + "address": f"{self.session.assemblies[asm].name}!0x{token:08x}", + "name": record.reflection_name, + "kind": record.kind, + } + for (asm, token), record in self.session.types.items() + if record.visibility.startswith("Public") + or record.visibility.startswith("NestedPublic") + if pattern is None or pattern(record.reflection_name) + ] + return _page(rows, args) + + def op_sections(self, _args: dict[str, Any]) -> dict[str, Any]: + return { + "assemblies": [ + { + "name": info.name, + "version": info.version, + "framework": info.target_framework, + "kind": info.kind, + "path": str(info.file.path), + "types": len(info.types), + "obfuscation_hints": info.obfuscation_hints, + } + for info in self.session.assemblies + ] + } + + def op_types(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + kind = args.get("kind") + rows = [ + { + "address": f"{self.session.assemblies[asm].name}!0x{token:08x}", + "name": record.reflection_name, + "kind": record.kind, + "members": len(record.fields) + + len(record.properties) + + len(record.events) + + len(record.methods), + } + for (asm, token), record in self.session.types.items() + if (pattern is None or pattern(record.reflection_name)) + and (not kind or record.kind == kind) + ] + return _page(rows, args) + + def op_type(self, args: dict[str, Any]) -> dict[str, Any]: + entity = self.resolve(args.get("name"), kinds=("type",)) + record = self.type_of(entity) + methods = self.session.methods + return { + "address": entity.address(self.session), + "name": record.reflection_name, + "kind": record.kind, + "base_type": record.base_type, + "interfaces": record.interfaces, + "fields": [ + {"name": f.name, "type": f.type, "token": f"0x{f.token:08x}"} + for f in record.fields + ], + "properties": [ + {"name": p.name, "type": p.type, "token": f"0x{p.token:08x}"} + for p in record.properties + ], + "methods": [ + { + "name": methods[i].name, + "prototype": methods[i].prototype, + "address": methods[i].address, + } + for i in record.methods + ], + "c_decl": _clip(self.csharp(entity)), + } + + def op_read_bytes(self, args: dict[str, Any]) -> dict[str, Any]: + """Raw bytes of an assembly: at a file offset, or a method's body.""" + target = args.get("address") + size = max(1, min(int(args.get("size") or 64), MAX_READ)) + if ( + isinstance(target, str) + and "!" in target + or (isinstance(target, str) and target.startswith("0x0")) + ): + entity = self.resolve(target, kinds=("method",)) + method = self.method_of(entity) + info = self.session.assemblies[entity.assembly] + offset = info.pe.rva_to_offset(method.rva) if method.rva else None + if offset is None: + raise OpError("the method has no body") + path, size = ( + Path(self.original_paths[entity.assembly]), + min(size, max(1, method.il_size + 12)), + ) + else: + offset = int(str(target), 0) + path = Path(self.original_paths[0]) + with path.open("rb") as handle: + handle.seek(offset) + data = handle.read(size) + return { + "file": path.name, + "offset": hex(offset), + "size": len(data), + "hex": data.hex(), + "ascii": "".join(chr(b) if 32 <= b < 127 else "." for b in data), + } + + def op_search(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = str(args.get("pattern") or "") + kind = args.get("kind") or "text" + limit = max(1, min(int(args.get("limit") or 50), 1000)) + if not pattern: + raise OpError("search needs a pattern") + hits: list[dict[str, Any]] = [] + if kind == "hex": + needle = bytes.fromhex(pattern.replace(" ", "")) + for index, path in enumerate(self.original_paths): + data = Path(path).read_bytes() + start = data.find(needle) + while start >= 0 and len(hits) < limit: + hits.append( + { + "assembly": self.session.assemblies[index].name, + "offset": hex(start), + } + ) + start = data.find(needle, start + 1) + return {"pattern": pattern, "kind": kind, "hits": hits} + matcher = re.compile(pattern) if kind == "regex" else None + for row in self.op_strings({"limit": 100000})["items"]: + value = row["value"] + if matcher.search(value) if matcher else pattern in value: + hits.append(row) + if len(hits) >= limit: + break + return {"pattern": pattern, "kind": kind, "hits": hits} + + def op_lookup(self, args: dict[str, Any]) -> dict[str, Any]: + query = str(args.get("name") or "").strip().lower() + if not query: + raise OpError("lookup needs a name") + limit = max(1, min(int(args.get("limit") or 50), 500)) + rows: list[dict[str, Any]] = [] + seen: set[Entity] = set() + for name in sorted(self.names): + if query not in name.lower(): + continue + for entity in self.names[name]: + if entity in seen: + continue + seen.add(entity) + rows.append( + { + "name": self.entity_name(entity), + "kind": entity.kind, + "address": entity.address(self.session), + } + ) + if len(rows) >= limit: + break + return {"query": query, "matches": rows[:limit]} + + # -- eval ------------------------------------------------------------------------- + + def _namespace(self) -> dict[str, Any]: + if self.eval_namespace is None: + self.eval_namespace = { + "__name__": "__tocode_eval__", + "session": self.session, + "modules": list(self.session._modules), + "module": self.session._modules[0] if self.session._modules else None, + "resolve": self.resolve, + "decompile": lambda target: self.csharp(self.resolve(target)), + "touch": self.touch, + "result": None, + } + return self.eval_namespace + + def op_eval(self, args: dict[str, Any]) -> dict[str, Any]: + if not self.allow_eval: + raise OpError("eval is disabled for this server (--no-eval / --read-only)") + if (args.get("language") or "python") != "python": + raise OpError("eval on .NET runs Python (dnlib objects are in `modules`)") + if args.get("reset"): + self.eval_namespace = None + before = set(self.dirty) + value, output, error = run_python( + self._namespace(), str(args.get("code") or "") + ) + return { + "result": _jsonable(value), + "stdout": _clip(output), + "error": error, + "touched": sorted(self.type_name(key) for key in self.dirty - before), + } + + +def write_module(module: Any, target: Path) -> None: + """Write a dnlib module with its metadata tokens preserved, atomically. + + Only row ids are preserved: keeping the old string heap would carry the + previous names along in the file. + """ + from dnlib.DotNet.Writer import ( # type: ignore[import-not-found] + MetadataFlags, + ModuleWriterOptions, + NativeModuleWriterOptions, + ) + + temporary = target.with_name(f".{target.name}.{os.getpid()}.tmp") + if module.IsILOnly: + options = ModuleWriterOptions(module) + options.MetadataOptions.Flags = ( + options.MetadataOptions.Flags | MetadataFlags.PreserveRids + ) + module.Write(str(temporary), options) + else: # mixed-mode (C++/CLI): keep the native code + native = NativeModuleWriterOptions(module, False) + native.MetadataOptions.Flags = ( + native.MetadataOptions.Flags | MetadataFlags.PreserveRids + ) + module.NativeWrite(str(temporary), native) + os.replace(temporary, target) + + +def _load_dotnet_locations(export_dir: Path | None) -> dict[str, dict[str, Any]]: + """`Assembly!0xTOKEN` -> where the method is in the export.""" + if export_dir is None: + return {} + try: + rows = json.loads((export_dir / "function-index.json").read_text("utf-8")).get( + "functions", [] + ) + except (OSError, ValueError): + return {} + locations = {} + for row in rows: + c = row.get("c") or {} + asm = row.get("asm") or {} + locations[str(row.get("address"))] = { + "c": f"{c.get('path')}:{c.get('line_start')}-{c.get('line_end')}", + "il": f"{asm.get('path')}:{asm.get('line_start')}-{asm.get('line_end')}", + } + return locations diff --git a/src/tocode/live/dotnet_sync.py b/src/tocode/live/dotnet_sync.py new file mode 100644 index 0000000..6347c67 --- /dev/null +++ b/src/tocode/live/dotnet_sync.py @@ -0,0 +1,315 @@ +"""``sync`` for .NET exports: rewrite the types the edits touched. + +The export keeps one ``.cs``/``.il`` pair per top-level type, so a sync +re-decompiles exactly the stale types from the patched assemblies, applies +the text edits (local renames, comments), moves the files of renamed types, +and rewrites the JSON documents with the export's own writers. Every other +file keeps its bytes. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from pathlib import Path +import time +from typing import TYPE_CHECKING, Any + +from ..metadata import write_json, write_text_atomic +from .worker import OpError + +if TYPE_CHECKING: + from .dotnet_live import DotnetLiveSession + + +def register(live: "DotnetLiveSession") -> None: + live.ops["sync"] = lambda args: sync(live, args) + live.ops["export"] = lambda args: export(live, args) + + +# -------------------------------------------------------------------------- +# The export's current layout +# -------------------------------------------------------------------------- + + +@dataclass +class Layout: + # (assembly index, type token) -> {"cs_path": Path, "il_path": Path} + types: dict[tuple[int, int], dict[str, Any]] = field(default_factory=dict) + # (assembly index, method token) -> (cs_start, cs_end, il_start, il_end) + methods: dict[tuple[int, int], tuple[Any, Any, Any, Any]] = field( + default_factory=dict + ) + + +def disk_layout(live: "DotnetLiveSession") -> Layout: + cached = live._layout + if cached is not None: + return cached + layout = Layout() + root = live.target.export_dir + if root is None: + return layout + index = {info.name: info.index for info in live.session.assemblies} + for row in _rows(root / "types.json", "types"): + assembly = index.get(str(row.get("assembly"))) + if assembly is None or not row.get("cs_file"): + continue + layout.types[(assembly, int(str(row["token"]), 16))] = { + "cs_path": (root / row["cs_file"]).resolve(), + "il_path": (root / row["il_file"]).resolve() + if row.get("il_file") + else None, + } + for row in _rows(root / "functions.json", "functions"): + assembly = index.get(str(row.get("assembly"))) + if assembly is None or not row.get("token"): + continue + layout.methods[(assembly, int(str(row["token"]), 16))] = ( + row.get("source_line_start"), + row.get("source_line_end"), + row.get("asm_line_start"), + row.get("asm_line_end"), + ) + live._layout = layout + return layout + + +def _rows(path: Path, key: str) -> list[dict[str, Any]]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return [] + rows = payload.get(key, []) + return [row for row in rows if isinstance(row, dict)] + + +# -------------------------------------------------------------------------- +# Sync +# -------------------------------------------------------------------------- + + +def sync(live: "DotnetLiveSession", args: dict[str, Any]) -> dict[str, Any]: + from ..backends.dotnet import WorkItem, decompile_item + from ..dotnet import ( + DotnetContext, + DotnetExportOptions, + _assign_paths, + assign_type_ranges, + type_texts, + ) + from . import dotnet_edits + + root = live.target.export_dir + if root is None: + raise OpError("this session has no export folder; call `export` first") + started = time.monotonic() + written_assemblies = live.reload() + session = live.session + layout = disk_layout(live) + mode = str(args.get("mode") or "auto") + tops = { + (asm, token) + for (asm, token), record in session.types.items() + if record.declaring_type is None + and not (record.name == "" and not record.methods and not record.fields) + } + dirty = set(live.dirty) + for item in args.get("addresses") or []: + dirty.add(live.top_level(live.resolve(item))) + if mode == "full": + dirty |= tops + dirty &= tops + + raw_dir = root / "src" / "raw" + context = DotnetContext( + input_path=live.target.open_path, + options=DotnetExportOptions( + out_dir=root, + include_framework=bool(live._project().get("include_framework")), + ), + progress=live.progress, + ) + context.found = live.found + context.session = session + context.root = root + context.raw_dir = raw_dir + context.data_dir = root / "data" + _assign_paths(context) # where each type belongs under its current name + + # Untouched types keep their file even if numbering would now differ. + claimed = {session.types[key].cs_path for key in dirty} + for key in tops - dirty: + old = layout.types.get(key) + if old is None: + dirty.add(key) # a type the export never had + elif old["cs_path"] in claimed and session.types[key].cs_path != old["cs_path"]: + dirty.add(key) # its file was taken by a renamed type + else: + for nested in session.nested_types(key[0], key[1]): + nested.cs_path, nested.il_path = old["cs_path"], old["il_path"] + for method in session.methods: + if (method.assembly, _top(session, method)) in dirty: + continue + span = layout.methods.get((method.assembly, method.token)) + if span is not None: + method.cs_start, method.cs_end, method.il_start, method.il_end = span + + rewritten: list[Path] = [] + moved: list[dict[str, str]] = [] + failures: list[dict[str, str]] = [] + for key in sorted(dirty): + asm, token = key + record = session.types[key] + result = decompile_item(WorkItem(asm, live.assembly_path(asm), token)) + cs_text, il_text, offset = type_texts(session, record, result) + if result.cs is not None: + cs_text, result.cs_ranges = dotnet_edits.apply_text_edits( + live, asm, cs_text, result.cs_ranges, offset + ) + if record.error: + failures.append({"type": record.reflection_name, "error": record.error}) + old = layout.types.get(key) + if old is not None and old["cs_path"] != record.cs_path: + for path in (old["cs_path"], old["il_path"]): + if path is not None and Path(path).exists(): + Path(path).unlink() + moved.append( + {"from": _rel(root, old["cs_path"]), "to": _rel(root, record.cs_path)} + ) + for path, text in ((record.cs_path, cs_text), (record.il_path, il_text)): + if not path.exists() or path.read_text(encoding="utf-8") != text: + write_text_atomic(path, text) + rewritten.append(path) + assign_type_ranges(session, record, result, cs_text, il_text, offset) + + _write_documents(live, context) + links = None + if rewritten or moved: + from .sync import _refresh_folder_links + + links = _refresh_folder_links(live, root) # type: ignore[arg-type] + live.edits.synced = len(live.edits.entries) + live.edits.save() + live.dirty.clear() + live._layout = None + live._locations = None + live._words = None + live._edited = False + report = { + "export_dir": str(root), + "mode": mode, + "rerendered": len(dirty), + "changed": sorted(session.types[key].reflection_name for key in dirty), + "files_rewritten": sorted(_rel(root, path) for path in rewritten), + "moved": moved, + "patched": written_assemblies, + "edits": len(live.edits.entries), + "failures": failures, + "folder_links": links, + "seconds": round(time.monotonic() - started, 2), + } + log = root / ".tocode" / "sync" / "sync-log.jsonl" + log.parent.mkdir(parents=True, exist_ok=True) + with log.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps({"time": time.strftime("%Y-%m-%dT%H:%M:%S"), **report}) + "\n" + ) + return report + + +def _top(session: Any, method: Any) -> int: + record = session.types.get((method.assembly, method.type_token)) + if record is None: + return method.type_token + return int(record.top_level or record.token) + + +def _write_documents(live: "DotnetLiveSession", context: Any) -> None: + """The export's JSON documents, rewritten from the live inventory.""" + from .. import dotnet_metadata as meta + from ..dotnet import _write_metadata + + root = context.root + _write_metadata(context) + try: + libs = json.loads((root / "native-libs.json").read_text("utf-8")).get( + "libs", [] + ) + except (OSError, ValueError): + libs = [] + write_json( + root / "triage.json", + meta.triage_json(context.session, context.found, context.reachable, libs), + ) + sources = sorted( + { + str(record.cs_path) + for record in context.session.types.values() + if record.cs_path is not None + } + | {str(p) for p in (root / "src" / "raw").rglob("Properties/AssemblyInfo.cs")} + ) + il_files = sorted( + { + str(record.il_path) + for record in context.session.types.values() + if record.il_path is not None + } + | {str(p) for p in (root / "src" / "raw").rglob("Properties/Manifest.il")} + ) + for name in ("project.json", "export-manifest.json"): + path = root / name + try: + document = json.loads(path.read_text("utf-8")) + except (OSError, ValueError): + continue + document.update( + { + "source_files": sources, + "il_files": il_files, + "asm_files": il_files, + "function_count": len(context.session.methods), + "type_count": len(context.session.types), + } + ) + write_json(path, document) + + +def _rel(root: Path, path: Any) -> str: + try: + return Path(path).resolve().relative_to(root.resolve()).as_posix() + except ValueError: + return str(path) + + +# -------------------------------------------------------------------------- +# Export (a session opened on a plain assembly) +# -------------------------------------------------------------------------- + + +def export(live: "DotnetLiveSession", args: dict[str, Any]) -> dict[str, Any]: + from ..dotnet import DotnetExportOptions, export_dotnet + + if live.target.export_dir is not None and not args.get("out_dir"): + raise OpError( + f"this session already has an export at {live.target.export_dir}; use `sync`" + ) + out = Path(str(args["out_dir"])) if args.get("out_dir") else None + summary = export_dotnet( + live.target.open_path, + options=DotnetExportOptions(out_dir=out, native=False), + progress=live.progress, + ) + live.target.export_dir = summary.root_dir + live._layout = None + live._locations = None + report: dict[str, Any] = { + "export_dir": str(summary.root_dir), + "types": summary.type_count, + } + if live.edits.entries: + # Edits made before the export: bring them in. + live.edits.synced = 0 + report["sync"] = sync(live, {"mode": "full"}) + return report diff --git a/src/tocode/live/edits.py b/src/tocode/live/edits.py new file mode 100644 index 0000000..87a17f3 --- /dev/null +++ b/src/tocode/live/edits.py @@ -0,0 +1,658 @@ +"""Changes to the database: the edit tools and saving. + +Every edit goes through this module so the session knows which functions it +made stale: the function itself, and for renames and prototype changes its +callers (their decompiled text names it). ``sync`` re-renders exactly those, +plus whatever its own diff of the database finds. + +One small adapter per backend; angr is query-only. +""" + +from __future__ import annotations + +from pathlib import Path +import re +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from .worker import LiveSession + +from .worker import OpError + +_IDENT_RX = re.compile(r"^[A-Za-z_$?@.][\w$?@.:<>~]*$") +_TYPE_NAMES_RX = re.compile( + r"\b(?:struct|union|enum)\s+([A-Za-z_]\w*)\s*\{|typedef\b[^;]*?\b([A-Za-z_]\w*)\s*;", + re.S, +) + + +# -------------------------------------------------------------------------- +# Saving +# -------------------------------------------------------------------------- + + +def database_path(live: "LiveSession") -> Path | None: + """The database this session writes to.""" + if live.target.database is not None: + return live.target.database + value = getattr(live.session, "_cache_db", None) + return Path(value) if value else None + + +def save_database(live: "LiveSession") -> dict[str, Any]: + backend = live.backend + path = database_path(live) + if backend == "ida": + import ida_loader # type: ignore[import-not-found] + + if path is None: + return {"saved": False, "reason": "no database path for this session"} + ok = bool(ida_loader.save_database(str(path), 0)) + if ok: + live._edited = False + return {"saved": ok, "database": str(path)} + if backend == "binja": + from ..backends.binja_headless import can_create_databases + + bv = getattr(live.session, "_bv", None) + bn = getattr(live.session, "_bn", None) + if bv is None or bn is None: + return {"saved": False, "reason": "no Binary Ninja view"} + if not can_create_databases(bn, getattr(live.session, "binja_dir", None)): + return { + "saved": False, + "reason": "this Binary Ninja cannot save databases headless; " + "edits last for this session only", + } + if path is not None and path.suffix.lower() == ".bndb" and path.is_file(): + ok = bool(bv.file.save_auto_snapshot()) + else: + path = path or _default_bndb(live) + ok = bool(bv.create_database(str(path))) + if ok: + live._edited = False + return {"saved": ok, "database": str(path) if path else None} + return { + "saved": False, + "reason": f"the {backend} backend keeps edits for this session only", + } + + +def _default_bndb(live: "LiveSession") -> Path: + base = live.target.export_dir or live.target.binary.parent + return base / f"{live.target.binary.stem}.bndb" + + +# -------------------------------------------------------------------------- +# Edit operations +# -------------------------------------------------------------------------- + + +def register(live: "LiveSession") -> None: + ops: dict[str, Callable[[dict[str, Any]], Any]] = { + "rename": lambda args: _rename(live, args), + "rename_many": lambda args: _rename_many(live, args), + "set_type": lambda args: _set_type(live, args), + "set_comment": lambda args: _set_comment(live, args), + "declare_types": lambda args: _declare_types(live, args), + "rename_local": lambda args: _rename_local(live, args), + "set_local_type": lambda args: _set_local_type(live, args), + "define_function": lambda args: _define_function(live, args), + "undefine": lambda args: _undefine(live, args), + } + live.ops.update(ops) + live.edit_ops = set(ops) + + +def _adapter(live: "LiveSession") -> "_Backend": + if not live.allow_edits: + raise OpError("edits are disabled for this server (--read-only)") + backend = live.backend + if backend == "ida": + return _Ida(live) + if backend == "binja": + return _Binja(live) + if backend == "r2": + return _R2(live) + raise OpError( + f"the {backend} backend is query-only; use IDA, Binary Ninja, or r2 to edit" + ) + + +def _edited(live: "LiveSession", *addresses: int) -> list[str]: + live._edited = True + for address in addresses: + live.dirty.add(address) + return sorted(hex(item) for item in addresses) + + +def _callers(live: "LiveSession", address: int) -> list[int]: + return list(live.analysis.callers.get(address, [])) + + +def _users(live: "LiveSession", address: int) -> list[int]: + """Functions that read or write a data address.""" + users = set() + refs = live.analysis.data_xrefs.get(address) + if refs is None: + try: + refs = live.session.data_xrefs([address]).get(address, []) + except Exception: # noqa: BLE001 + refs = [] + for source, _write in refs or []: + owner = live.routine_at(source) + if owner is not None: + users.add(owner.address) + return sorted(users) + + +def _rename(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = live.resolve(args.get("target")) + name = str(args.get("name") or "").strip() + if not name or not _IDENT_RX.match(name): + raise OpError(f"{name!r} is not a valid name") + backend = _adapter(live) + routine = live.analysis.routines.get(address) + if routine is not None: + old = routine.name + backend.rename_function(address, name) + routine.name = name + live.reindex() + touched = _edited(live, address, *_callers(live, address)) + return { + "address": hex(address), + "old": old, + "name": name, + "kind": "function", + "pending_sync": touched, + } + backend.rename_data(address, name) + for symbol in live.analysis.symbols: + if symbol.vaddr == address: + symbol.name = name + live.assigned_names[name] = address + live.reindex() + touched = _edited(live, *_users(live, address)) + live.globals_changed = True + return { + "address": hex(address), + "name": name, + "kind": "data", + "pending_sync": touched, + } + + +def _rename_many(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + renames = args.get("renames") or [] + if not isinstance(renames, list) or not renames: + raise OpError("rename_many needs `renames`: a list of {target, name}") + done, errors, pending = 0, [], set() + # Resolve every target first: names refer to the state before the batch. + resolved = [] + for item in renames: + try: + resolved.append((item, live.resolve(item.get("target")))) + except OpError as exc: + errors.append({"target": item.get("target"), "error": str(exc)}) + for item, address in resolved: + try: + result = _rename(live, {"target": address, "name": item.get("name")}) + except OpError as exc: + errors.append({"target": item.get("target"), "error": str(exc)}) + continue + done += 1 + pending.update(result["pending_sync"]) + return {"renamed": done, "errors": errors, "pending_sync": sorted(pending)} + + +def _set_type(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = live.resolve(args.get("target")) + decl = str(args.get("c_decl") or "").strip().rstrip(";") + if not decl: + raise OpError( + "set_type needs a C declaration, e.g. `int parse(char *buf, int len)`" + ) + backend = _adapter(live) + routine = live.analysis.routines.get(address) + if routine is not None: + backend.set_function_type(address, decl) + routine.signature = decl + live.stale = True # params/locals change: collect them again lazily + touched = _edited(live, address, *_callers(live, address)) + return { + "address": hex(address), + "c_decl": decl, + "kind": "function", + "pending_sync": touched, + } + backend.set_data_type(address, decl) + live.globals_changed = True + touched = _edited(live, *_users(live, address)) + return { + "address": hex(address), + "c_decl": decl, + "kind": "data", + "pending_sync": touched, + } + + +def _set_comment(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = live.resolve(args.get("address"), what="address") + text = str(args.get("text") or "") + kind = args.get("kind") or "line" + owner = live.routine_at(address) + if kind == "function" and owner is None: + raise OpError(f"no function at {hex(address)} for a function comment") + backend = _adapter(live) + if kind == "function": + assert owner is not None + backend.function_comment(owner.address, text) + elif kind == "line": + backend.line_comment(address, text) + else: + raise OpError("kind must be line or function") + touched = _edited(live, *([owner.address] if owner else [])) + return {"address": hex(address), "kind": kind, "pending_sync": touched} + + +def _declare_types(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + source = str(args.get("c_source") or "").strip() + if not source: + raise OpError( + "declare_types needs C declarations (structs, unions, enums, typedefs)" + ) + backend = _adapter(live) + backend.declare_types(source) + names = sorted( + { + first or second + for first, second in _TYPE_NAMES_RX.findall(source) + if first or second + } + ) + live.types_changed = True + live._edited = True + live.refresh_types() + return {"declared": names, "pending_sync": "functions using these types"} + + +def _function(live: "LiveSession", args: dict[str, Any]) -> int: + return live.function_for(args.get("function")).address + + +def _rename_local(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = _function(live, args) + old = str(args.get("old") or "") + new = str(args.get("new") or "").strip() + if not old or not new or not _IDENT_RX.match(new): + raise OpError("rename_local needs `old` and a valid `new` name") + _adapter(live).rename_local(address, old, new) + live.stale = True + return { + "function": hex(address), + "old": old, + "new": new, + "pending_sync": _edited(live, address), + } + + +def _set_local_type(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = _function(live, args) + name = str(args.get("name") or "") + c_type = str(args.get("c_type") or "").strip().rstrip(";") + if not name or not c_type: + raise OpError("set_local_type needs `name` and `c_type`") + _adapter(live).set_local_type(address, name, c_type) + live.stale = True + return { + "function": hex(address), + "name": name, + "c_type": c_type, + "pending_sync": _edited(live, address), + } + + +def _define_function(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = live.resolve(args.get("address"), what="address") + _adapter(live).define_function(address) + live.stale = True + live.functions_changed = True + return {"address": hex(address), "pending_sync": _edited(live, address)} + + +def _undefine(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + address = live.resolve(args.get("address"), what="address") + routine = live.analysis.routines.get(address) + callers = _callers(live, address) if routine is not None else [] + _adapter(live).undefine(address, routine is not None) + live.stale = True + live.functions_changed = True + return {"address": hex(address), "pending_sync": _edited(live, *callers)} + + +# -------------------------------------------------------------------------- +# Backends +# -------------------------------------------------------------------------- + + +class _Backend: + def __init__(self, live: "LiveSession") -> None: + self.live = live + self.session: Any = live.session + + def _fail(self, what: str) -> OpError: + return OpError(f"{self.live.backend} could not {what}") + + def rename_function(self, address: int, name: str) -> None: + raise self._fail("rename functions") + + def rename_data(self, address: int, name: str) -> None: + raise self._fail("rename data") + + def set_function_type(self, address: int, decl: str) -> None: + raise self._fail("set function types") + + def set_data_type(self, address: int, decl: str) -> None: + raise self._fail("set data types") + + def function_comment(self, address: int, text: str) -> None: + raise self._fail("set function comments") + + def line_comment(self, address: int, text: str) -> None: + raise self._fail("set comments") + + def declare_types(self, source: str) -> None: + raise self._fail("declare types") + + def rename_local(self, function: int, old: str, new: str) -> None: + raise self._fail("rename local variables") + + def set_local_type(self, function: int, name: str, c_type: str) -> None: + raise self._fail("set local variable types") + + def define_function(self, address: int) -> None: + raise self._fail("define functions") + + def undefine(self, address: int, is_function: bool) -> None: + raise self._fail("undefine") + + +class _Ida(_Backend): + def _modules(self) -> Any: + import ida_bytes # type: ignore[import-not-found] + import ida_funcs # type: ignore[import-not-found] + import ida_name # type: ignore[import-not-found] + import idc # type: ignore[import-not-found] + + return ida_bytes, ida_funcs, ida_name, idc + + def _set_name(self, address: int, name: str) -> None: + _, _, ida_name, _ = self._modules() + flags = ida_name.SN_CHECK | getattr(ida_name, "SN_NOWARN", 0) + if not ida_name.set_name(address, name, flags): + raise OpError( + f"IDA refused the name {name!r} at {hex(address)} (already used, or invalid)" + ) + self._forget(address) + + def _forget(self, address: int) -> None: + """Drop cached decompilations so the next render sees the edit.""" + try: + import ida_hexrays # type: ignore[import-not-found] + + ida_hexrays.mark_cfunc_dirty(address, False) + except Exception: # noqa: BLE001 - no decompiler, nothing cached + pass + cache = getattr(self.session, "_locals_cache", None) + if isinstance(cache, dict): + cache.clear() + + def rename_function(self, address: int, name: str) -> None: + self._set_name(address, name) + + def rename_data(self, address: int, name: str) -> None: + self._set_name(address, name) + + def _apply(self, address: int, decl: str) -> None: + _, _, _, idc = self._modules() + if not idc.SetType(address, decl + ";"): + raise OpError(f"IDA could not parse or apply the type {decl!r}") + self._forget(address) + + def set_function_type(self, address: int, decl: str) -> None: + self._apply(address, decl) + + def set_data_type(self, address: int, decl: str) -> None: + self._apply(address, decl) + + def function_comment(self, address: int, text: str) -> None: + _, ida_funcs, _, _ = self._modules() + func = ida_funcs.get_func(address) + if func is None or not ida_funcs.set_func_cmt(func, text, False): + raise self._fail("set the function comment") + self._forget(address) + + def line_comment(self, address: int, text: str) -> None: + ida_bytes, _, _, _ = self._modules() + if not ida_bytes.set_cmt(address, text, False): + raise self._fail("set the comment") + self._forget(address) + + def declare_types(self, source: str) -> None: + _, _, _, idc = self._modules() + errors = idc.parse_decls(source, getattr(idc, "PT_SILENT", 0)) + if errors is None or errors < 0 or errors > 0: + raise OpError(f"IDA could not parse the declarations ({errors} errors)") + try: + import ida_hexrays # type: ignore[import-not-found] + + ida_hexrays.clear_cached_cfuncs() + except Exception: # noqa: BLE001 + pass + + def _hexrays(self) -> Any: + try: + import ida_hexrays # type: ignore[import-not-found] + except ImportError as exc: + raise OpError("local variables need the Hex-Rays decompiler") from exc + return ida_hexrays + + def rename_local(self, function: int, old: str, new: str) -> None: + ida_hexrays = self._hexrays() + if not ida_hexrays.rename_lvar(function, old, new): + raise OpError( + f"no local variable {old!r} in {hex(function)} (or {new!r} is taken)" + ) + self._forget(function) + + def set_local_type(self, function: int, name: str, c_type: str) -> None: + ida_hexrays = self._hexrays() + import ida_typeinf # type: ignore[import-not-found] + + cfunc = ida_hexrays.decompile(function) + if cfunc is None: + raise OpError(f"Hex-Rays could not decompile {hex(function)}") + lvar = next((item for item in cfunc.get_lvars() if item.name == name), None) + if lvar is None: + raise OpError(f"no local variable {name!r} in {hex(function)}") + tif = ida_typeinf.tinfo_t() + if ( + ida_typeinf.parse_decl( + tif, None, f"{c_type} __tocode_v;", ida_typeinf.PT_SIL + ) + is None + ): + raise OpError(f"IDA could not parse the type {c_type!r}") + info = ida_hexrays.lvar_saved_info_t() + info.ll = lvar + info.type = tif + if not ida_hexrays.modify_user_lvar_info(function, ida_hexrays.MLI_TYPE, info): + raise OpError(f"IDA could not retype {name!r}") + self._forget(function) + + def define_function(self, address: int) -> None: + _, ida_funcs, _, _ = self._modules() + if not ida_funcs.add_func(address): + raise OpError(f"IDA could not create a function at {hex(address)}") + try: + import ida_auto # type: ignore[import-not-found] + + ida_auto.auto_wait() + except Exception: # noqa: BLE001 + pass + + def undefine(self, address: int, is_function: bool) -> None: + ida_bytes, ida_funcs, _, _ = self._modules() + if is_function: + if not ida_funcs.del_func(address): + raise OpError(f"IDA could not delete the function at {hex(address)}") + elif not ida_bytes.del_items(address, 0, 1): + raise OpError(f"IDA could not undefine {hex(address)}") + + +class _Binja(_Backend): + @property + def bv(self) -> Any: + return self.session._bv + + @property + def bn(self) -> Any: + return self.session._bn + + def _func(self, address: int) -> Any: + func = self.bv.get_function_at(address) + if func is None: + raise OpError(f"no Binary Ninja function at {hex(address)}") + return func + + def _forget(self) -> None: + for cache in ("_decompile_cache", "_disasm_cache"): + value = getattr(self.session, cache, None) + if isinstance(value, dict): + value.clear() + + def _parse(self, decl: str) -> Any: + try: + parsed = self.bv.parse_type_string(decl) + except Exception as exc: # noqa: BLE001 + raise OpError(f"Binary Ninja could not parse {decl!r}: {exc}") from exc + return parsed[0] if isinstance(parsed, tuple) else parsed + + def rename_function(self, address: int, name: str) -> None: + self._func(address).name = name + self._forget() + + def rename_data(self, address: int, name: str) -> None: + bn = self.bn + self.bv.define_user_symbol(bn.Symbol(bn.SymbolType.DataSymbol, address, name)) + self._forget() + + def set_function_type(self, address: int, decl: str) -> None: + func = self._func(address) + setter = getattr(func, "set_user_type", None) + parsed = self._parse(decl) + if callable(setter): + setter(parsed) + else: + func.type = parsed + self._forget() + + def set_data_type(self, address: int, decl: str) -> None: + self.bv.define_user_data_var(address, self._parse(decl)) + self._forget() + + def function_comment(self, address: int, text: str) -> None: + self._func(address).comment = text + self._forget() + + def line_comment(self, address: int, text: str) -> None: + owner = self.live.routine_at(address) + if owner is not None: + self._func(owner.address).set_comment_at(address, text) + else: + self.bv.set_comment_at(address, text) + self._forget() + + def declare_types(self, source: str) -> None: + try: + result = self.bv.parse_types_from_string(source) + except Exception as exc: # noqa: BLE001 + raise OpError( + f"Binary Ninja could not parse the declarations: {exc}" + ) from exc + for kind in ("types", "variables", "functions"): + for name, type_obj in dict(getattr(result, kind, {}) or {}).items(): + if kind == "types": + self.bv.define_user_type(name, type_obj) + self._forget() + + def _var(self, function: int, name: str) -> tuple[Any, Any]: + func = self._func(function) + for var in list(getattr(func, "vars", [])): + if getattr(var, "name", None) == name: + return func, var + raise OpError(f"no local variable {name!r} in {hex(function)}") + + def rename_local(self, function: int, old: str, new: str) -> None: + func, var = self._var(function, old) + func.create_user_var(var, var.type, new) + self._forget() + + def set_local_type(self, function: int, name: str, c_type: str) -> None: + func, var = self._var(function, name) + func.create_user_var(var, self._parse(c_type), name) + self._forget() + + def define_function(self, address: int) -> None: + self.bv.add_user_function(address) + self.bv.update_analysis_and_wait() + + def undefine(self, address: int, is_function: bool) -> None: + if is_function: + self.bv.remove_user_function(self._func(address)) + else: + self.bv.undefine_user_data_var(address) + + +class _R2(_Backend): + def _cmd(self, command: str) -> str: + output = str(self.session.cmd(command) or "") + if "ERROR" in output or "Cannot" in output: + raise OpError(f"r2: {output.strip()[:300]}") + return output + + def rename_function(self, address: int, name: str) -> None: + self._cmd(f"afn {name} @ {address:#x}") + + def rename_data(self, address: int, name: str) -> None: + self._cmd(f"f {name} @ {address:#x}") + + def set_function_type(self, address: int, decl: str) -> None: + self._cmd(f"afs {decl} @ {address:#x}") + + def function_comment(self, address: int, text: str) -> None: + self.line_comment(address, text) + + def line_comment(self, address: int, text: str) -> None: + self._cmd(f"CCu base64:{_b64(text)} @ {address:#x}") + + def declare_types(self, source: str) -> None: + self._cmd(f'td "{source}"'.replace("\n", " ")) + + def rename_local(self, function: int, old: str, new: str) -> None: + self._cmd(f"afvn {new} {old} @ {function:#x}") + + def set_local_type(self, function: int, name: str, c_type: str) -> None: + self._cmd(f"afvt {name} {c_type} @ {function:#x}") + + def define_function(self, address: int) -> None: + self._cmd(f"af @ {address:#x}") + + def undefine(self, address: int, is_function: bool) -> None: + self._cmd(f"af- @ {address:#x}" if is_function else f"Cd- @ {address:#x}") + + +def _b64(text: str) -> str: + import base64 + + return base64.b64encode(text.encode("utf-8")).decode("ascii") diff --git a/src/tocode/live/hooks.py b/src/tocode/live/hooks.py new file mode 100644 index 0000000..f84b6e7 --- /dev/null +++ b/src/tocode/live/hooks.py @@ -0,0 +1,106 @@ +"""Change notifications from the database, so `sync` sees what scripts changed. + +Edits made through the edit tools already record what they make stale; edits +made by `eval` scripts are only partly visible to sync's diff (it sees names +and prototypes, not comments or a function body's details). These hooks mark +the affected functions while the session runs. Installing them is best +effort: a backend without notifications relies on the diff and `touch()`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .worker import LiveSession + + +def install(live: "LiveSession") -> Any: + try: + if live.backend == "ida": + return _ida(live) + if live.backend == "binja": + return _binja(live) + except Exception as exc: # noqa: BLE001 - notifications are optional + live.log(f"live: change hooks unavailable: {exc}") + return None + + +def _ida(live: "LiveSession") -> Any: + import ida_funcs # type: ignore[import-not-found] + import ida_idp # type: ignore[import-not-found] + + def mark(address: Any) -> None: + func = ida_funcs.get_func(int(address)) + if func is not None: + live.dirty.add(int(func.start_ea)) + + class Hooks(ida_idp.IDB_Hooks): # type: ignore[misc] + # Signatures differ between IDA versions: take what comes. + def renamed(self, ea: int, *_args: Any) -> int: + mark(ea) + return 0 + + def cmt_changed(self, ea: int, *_args: Any) -> int: + mark(ea) + return 0 + + def extra_cmt_changed(self, ea: int, *_args: Any) -> int: + mark(ea) + return 0 + + def range_cmt_changed(self, _kind: Any, area: Any, *_args: Any) -> int: + start = getattr(area, "start_ea", None) + if start is not None: + mark(start) + return 0 + + def ti_changed(self, ea: int, *_args: Any) -> int: + mark(ea) + return 0 + + def func_added(self, pfn: Any) -> int: + live.dirty.add(int(pfn.start_ea)) + live.functions_changed = True + return 0 + + def deleting_func(self, _pfn: Any) -> int: + live.functions_changed = True + return 0 + + def local_types_changed(self, *_args: Any) -> int: + live.types_changed = True + return 0 + + hooks = Hooks() + hooks.hook() + return hooks + + +def _binja(live: "LiveSession") -> Any: + bn = live.session._bn + bv = live.session._bv + base = getattr(bn, "BinaryDataNotification", None) + if base is None or not hasattr(bv, "register_notification"): + return None + + class Notifications(base): # type: ignore[misc, valid-type] + def function_added(self, _view: Any, func: Any) -> None: + live.dirty.add(int(func.start)) + live.functions_changed = True + + def function_removed(self, _view: Any, _func: Any) -> None: + live.functions_changed = True + + def function_updated(self, _view: Any, func: Any) -> None: + live.dirty.add(int(func.start)) + + def type_defined(self, *_args: Any) -> None: + live.types_changed = True + + def type_undefined(self, *_args: Any) -> None: + live.types_changed = True + + notifications = Notifications() + bv.register_notification(notifications) + return notifications diff --git a/src/tocode/live/jsonrpc.py b/src/tocode/live/jsonrpc.py new file mode 100644 index 0000000..66f7071 --- /dev/null +++ b/src/tocode/live/jsonrpc.py @@ -0,0 +1,210 @@ +"""MCP over stdio: newline-delimited JSON-RPC 2.0, standard library only. + +Requests are answered on worker threads so a slow call on one session does +not hold up the others; responses are written whole, one per line, under a +lock. Only protocol messages reach the protocol stream: the caller hands in +a private copy of the original stdout (see ``server.protect_stdout``). +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +import json +import threading +import traceback +from typing import Any, BinaryIO, Callable + +SUPPORTED_PROTOCOL_VERSIONS = ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25") +LATEST_PROTOCOL_VERSION = SUPPORTED_PROTOCOL_VERSIONS[-1] + +PARSE_ERROR = -32700 +INVALID_REQUEST = -32600 +METHOD_NOT_FOUND = -32601 +INVALID_PARAMS = -32602 +INTERNAL_ERROR = -32603 + + +class RpcError(Exception): + def __init__(self, code: int, message: str, data: Any = None) -> None: + super().__init__(message) + self.code = code + self.message = message + self.data = data + + +@dataclass +class ToolCall: + """What a tool handler gets besides its arguments.""" + + name: str + arguments: dict[str, Any] + progress_token: Any = None + notify: Callable[[str, dict[str, Any]], None] = lambda _method, _params: None + + def progress(self, done: float, total: float | None, message: str = "") -> None: + if self.progress_token is None: + return + params: dict[str, Any] = { + "progressToken": self.progress_token, + "progress": done, + } + if total is not None: + params["total"] = total + if message: + params["message"] = message + self.notify("notifications/progress", params) + + +@dataclass +class McpApp: + """The server's identity and tool table, independent of the transport.""" + + name: str + version: str + instructions: str + list_tools: Callable[[], list[dict[str, Any]]] + call_tool: Callable[[ToolCall], dict[str, Any]] + on_shutdown: Callable[[], None] = lambda: None + handlers: dict[str, Callable[[dict[str, Any]], Any]] = field(default_factory=dict) + + +class StdioServer: + def __init__( + self, + app: McpApp, + reader: BinaryIO, + writer: BinaryIO, + *, + max_workers: int = 8, + log: Callable[[str], None] = lambda _message: None, + ) -> None: + self.app = app + self.reader = reader + self.writer = writer + self.log = log + self._write_lock = threading.Lock() + self._pool = ThreadPoolExecutor( + max_workers=max_workers, thread_name_prefix="tocode-mcp" + ) + self._initialized = False + + # -- transport ------------------------------------------------------------- + + def serve_forever(self) -> None: + try: + for raw in self.reader: + line = raw.strip() + if not line: + continue + try: + message = json.loads(line) + except ValueError: + self._send_error(None, PARSE_ERROR, "invalid JSON") + continue + if isinstance(message, list): + # Batches were removed from MCP; answer each element. + for item in message: + self._dispatch(item) + else: + self._dispatch(message) + finally: + self._pool.shutdown(wait=True, cancel_futures=False) + self.app.on_shutdown() + + def send(self, payload: dict[str, Any]) -> None: + data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + with self._write_lock: + self.writer.write(data.encode("utf-8") + b"\n") + self.writer.flush() + + def notify(self, method: str, params: dict[str, Any]) -> None: + self.send({"jsonrpc": "2.0", "method": method, "params": params}) + + def _send_result(self, request_id: Any, result: Any) -> None: + self.send({"jsonrpc": "2.0", "id": request_id, "result": result}) + + def _send_error( + self, request_id: Any, code: int, message: str, data: Any = None + ) -> None: + error: dict[str, Any] = {"code": code, "message": message} + if data is not None: + error["data"] = data + self.send({"jsonrpc": "2.0", "id": request_id, "error": error}) + + # -- dispatch ---------------------------------------------------------------- + + def _dispatch(self, message: Any) -> None: + if not isinstance(message, dict) or message.get("jsonrpc") != "2.0": + self._send_error(None, INVALID_REQUEST, "not a JSON-RPC 2.0 message") + return + method = message.get("method") + request_id = message.get("id") + if method is None: + return # a response to something we never send + params = message.get("params") or {} + if "id" not in message: + self._notification(str(method), params) + return + if method in ("initialize", "ping", "tools/list"): + # Cheap and ordered: answer inline. + self._answer(request_id, str(method), params) + return + self._pool.submit(self._answer, request_id, str(method), params) + + def _notification(self, method: str, params: dict[str, Any]) -> None: + if method == "notifications/initialized": + self._initialized = True + # notifications/cancelled and others: nothing to undo; the call + # finishes and its late response is ignored by the client. + + def _answer(self, request_id: Any, method: str, params: dict[str, Any]) -> None: + try: + result = self._handle(method, params) + except RpcError as exc: + self._send_error(request_id, exc.code, exc.message, exc.data) + except Exception as exc: # noqa: BLE001 - never let a handler kill the server + self.log(f"mcp: {method} failed: {traceback.format_exc()}") + self._send_error(request_id, INTERNAL_ERROR, f"{type(exc).__name__}: {exc}") + else: + self._send_result(request_id, result) + + def _handle(self, method: str, params: dict[str, Any]) -> Any: + if method == "initialize": + requested = str(params.get("protocolVersion") or "") + version = ( + requested + if requested in SUPPORTED_PROTOCOL_VERSIONS + else LATEST_PROTOCOL_VERSION + ) + return { + "protocolVersion": version, + "capabilities": {"tools": {"listChanged": False}}, + "serverInfo": {"name": self.app.name, "version": self.app.version}, + "instructions": self.app.instructions, + } + if method == "ping": + return {} + if method == "tools/list": + return {"tools": self.app.list_tools()} + if method == "tools/call": + name = params.get("name") + if not isinstance(name, str): + raise RpcError(INVALID_PARAMS, "tools/call needs a tool name") + arguments = params.get("arguments") or {} + if not isinstance(arguments, dict): + raise RpcError(INVALID_PARAMS, "tool arguments must be an object") + meta = params.get("_meta") or {} + call = ToolCall( + name=name, + arguments=arguments, + progress_token=meta.get("progressToken") + if isinstance(meta, dict) + else None, + notify=self.notify, + ) + return self.app.call_tool(call) + handler = self.app.handlers.get(method) + if handler is not None: + return handler(params) + raise RpcError(METHOD_NOT_FOUND, f"unknown method {method}") diff --git a/src/tocode/live/server.py b/src/tocode/live/server.py new file mode 100644 index 0000000..d502414 --- /dev/null +++ b/src/tocode/live/server.py @@ -0,0 +1,173 @@ +"""``tocode mcp``: the stdio MCP server entry point.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import sys +from typing import Any, BinaryIO + +from .. import __version__ +from ..errors import ToCodeError +from .jsonrpc import McpApp, StdioServer +from .sessions import ( + DEFAULT_CALL_TIMEOUT, + DEFAULT_IDLE_SECONDS, + DEFAULT_MAX_SESSIONS, + SessionManager, +) +from .tools import ToolTable + +INSTRUCTIONS = """\ +ToCode keeps decompiler sessions (IDA, Binary Ninja, radare2, angr) open and \ +syncs them with a ToCode export: a source-like project on disk (src/raw/**/*.c, \ +functions.json, types.json, ...). + +Start with `open` on an export folder (or a folder export, a binary, a database) \ +and pass the returned `session` to the other tools. Read the export's files for \ +the big picture; use `decompile`, `function`, `xrefs`, `strings`, `types` for live \ +answers. Improve the database with `rename`, `set_type`, `declare_types`, \ +`set_comment`, `rename_local`, `set_local_type` (or `eval` for anything else), \ +then call `sync` so the export on disk shows the new names and types; `save` \ +writes the database. Addresses are hex strings like 0x401000; names work too.""" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="tocode mcp", + description="Serve live decompiler sessions over MCP (stdio).", + ) + parser.add_argument( + "target", + nargs="?", + type=Path, + default=None, + help="Optional export folder, folder export, binary, or database to open at start.", + ) + parser.add_argument( + "--database", + type=Path, + default=None, + help="Database to open for TARGET instead of the export's own copy.", + ) + parser.add_argument( + "--backend", choices=("auto", "ida", "binja", "r2", "angr"), default=None + ) + parser.add_argument( + "--binja-dir", + type=Path, + default=None, + help="Binary Ninja install folder for its Python API.", + ) + parser.add_argument( + "--read-only", action="store_true", help="No edit tools and no eval." + ) + parser.add_argument("--no-eval", action="store_true", help="No eval tool.") + parser.add_argument( + "--max-sessions", + type=int, + default=DEFAULT_MAX_SESSIONS, + help=f"Workers kept alive at once (default {DEFAULT_MAX_SESSIONS}).", + ) + parser.add_argument( + "--idle-timeout", + type=float, + default=DEFAULT_IDLE_SECONDS, + help="Seconds before an idle session is saved and closed (0: never).", + ) + parser.add_argument( + "--config", + action="store_true", + help="Print how to register this server with MCP clients, then exit.", + ) + parser.add_argument( + "--install", + choices=("claude", "codex", "all"), + default=None, + help="Register this server with Claude Code (claude mcp add) and/or Codex (~/.codex/config.toml), then exit.", + ) + parser.add_argument( + "--call-timeout", + type=float, + default=DEFAULT_CALL_TIMEOUT, + help="Default seconds per call before a stuck worker is stopped.", + ) + return parser + + +def protect_stdout() -> BinaryIO: + """Keep the real stdout for the protocol; send fd 1 to stderr. + + Backends print to fd 1 (IDA's console, r2's warnings); after this they + land on stderr and can never corrupt the JSON-RPC stream. + """ + sys.stdout.flush() + protocol = os.fdopen(os.dup(1), "wb", buffering=0) + os.dup2(2, 1) + sys.stdout = sys.stderr + return protocol + + +def main(argv: list[str]) -> int: + args = build_parser().parse_args(argv) + if args.config or args.install: + from .setup import config_text, install_clients + + if args.config: + print(config_text()) + return 0 + clients = ["claude", "codex"] if args.install == "all" else [args.install] + try: + for line in install_clients(clients): + print(line) + except ToCodeError as exc: + print(f"tocode: {exc}", file=sys.stderr) + return 1 + return 0 + protocol_out = protect_stdout() + + def log(message: str) -> None: + print(message, file=sys.stderr, flush=True) + + worker_options: dict[str, Any] = { + "binja_dir": str(args.binja_dir) if args.binja_dir else None, + "allow_eval": not (args.read_only or args.no_eval), + "allow_edits": not args.read_only, + } + manager = SessionManager( + max_sessions=args.max_sessions, + idle_seconds=args.idle_timeout, + call_timeout=args.call_timeout, + worker_options=worker_options, + log=log, + ) + table = ToolTable( + manager, + allow_edits=not args.read_only, + allow_eval=not (args.read_only or args.no_eval), + ) + if args.target is not None: + try: + opened = manager.open( + str(args.target), + database=str(args.database) if args.database else None, + backend=args.backend, + ) + log(f"tocode mcp: opened {opened.get('session')}") + except ToCodeError as exc: + log(f"tocode mcp: could not open {args.target}: {exc}") + app = McpApp( + name="tocode", + version=__version__, + instructions=INSTRUCTIONS, + list_tools=table.list, + call_tool=table.call, + on_shutdown=manager.shutdown, + ) + server = StdioServer(app, sys.stdin.buffer, protocol_out, log=log) + try: + server.serve_forever() + except KeyboardInterrupt: + manager.shutdown() + return 0 diff --git a/src/tocode/live/sessions.py b/src/tocode/live/sessions.py new file mode 100644 index 0000000..4a052b5 --- /dev/null +++ b/src/tocode/live/sessions.py @@ -0,0 +1,395 @@ +"""Worker lifecycle for the MCP server. + +Each open binary or database is a ``Handle``: one spawned worker process +(``worker.worker_main``) reached over a pipe. Calls to one handle are +serialized; calls to different handles run in parallel. A call that times out +kills the worker (with its process group); a worker that dies is reopened on +the next call. At most ``max_sessions`` workers stay alive: opening another +closes (and saves) the least recently used one. Idle workers close after +``idle_seconds``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import itertools +import multiprocessing +from pathlib import Path +import threading +import time +from typing import Any, Callable + +from ..apk_native import _active_children, _active_lock, kill_process_group +from ..errors import ToCodeError +from . import targets as targets_module +from .targets import FolderTarget, Target + +DEFAULT_MAX_SESSIONS = 4 +DEFAULT_IDLE_SECONDS = 1800.0 +DEFAULT_CALL_TIMEOUT = 300.0 +# Opening analyzes a fresh binary from scratch: allow a long time. +DEFAULT_OPEN_TIMEOUT = 6 * 3600.0 +CLOSE_TIMEOUT = 600.0 +REAPER_SECONDS = 30.0 + + +class SessionError(ToCodeError): + pass + + +@dataclass +class Handle: + target: Target + process: Any = None + conn: Any = None + info: dict[str, Any] = field(default_factory=dict) + lock: threading.Lock = field(default_factory=threading.Lock) + last_used: float = field(default_factory=time.monotonic) + busy: bool = False + note: str | None = None # why the last worker went away + pending_note: str | None = None # told to the agent on the next result + + @property + def alive(self) -> bool: + return self.process is not None and self.process.is_alive() + + +class SessionManager: + def __init__( + self, + *, + max_sessions: int = DEFAULT_MAX_SESSIONS, + idle_seconds: float = DEFAULT_IDLE_SECONDS, + call_timeout: float = DEFAULT_CALL_TIMEOUT, + open_timeout: float = DEFAULT_OPEN_TIMEOUT, + worker_options: dict[str, Any] | None = None, + log: Callable[[str], None] = lambda _message: None, + log_dir: Path | None = None, + ) -> None: + self.max_sessions = max(1, max_sessions) + self.idle_seconds = idle_seconds + self.call_timeout = call_timeout + self.open_timeout = open_timeout + self.worker_options = dict(worker_options or {}) + self.log = log + self.log_dir = log_dir + self.handles: dict[str, Handle] = {} + self.folders: dict[str, FolderTarget] = {} + self._lock = threading.RLock() + self._ids = itertools.count(1) + self._stop = threading.Event() + self._reaper = threading.Thread( + target=self._reap_idle, name="tocode-mcp-reaper", daemon=True + ) + self._reaper.start() + + # -- opening ---------------------------------------------------------------- + + def open( + self, + path: str, + *, + database: str | None = None, + backend: str | None = None, + ) -> dict[str, Any]: + resolved = targets_module.resolve( + Path(path), + database=Path(database) if database else None, + backend=backend, + ) + if isinstance(resolved, FolderTarget): + with self._lock: + resolved.name = self._unique(resolved.name, self.folders) + self.folders[resolved.name] = resolved + return resolved.describe() + handle = self._adopt(resolved) + self._ensure_worker(handle) + return dict(handle.info) + + def _adopt(self, target: Target) -> Handle: + with self._lock: + for handle in self.handles.values(): + if ( + handle.target.open_path == target.open_path + and handle.target.backend == target.backend + ): + return handle # already open: reuse the session + target.name = self._unique(target.name, self.handles) + handle = Handle(target=target) + self.handles[target.name] = handle + return handle + + def _unique(self, name: str, taken: dict[str, Any]) -> str: + if name not in taken and name not in self.handles and name not in self.folders: + return name + for number in itertools.count(2): + candidate = f"{name}#{number}" + if candidate not in self.handles and candidate not in self.folders: + return candidate + raise AssertionError("unreachable") + + def handle(self, session: str) -> Handle: + """A session by id, or a folder binary by relative path (opened on demand).""" + with self._lock: + if not session or session == "default": + if len(self.handles) == 1: + return next(iter(self.handles.values())) + raise SessionError( + "pass `session` (see `sessions`, or `open` a path first)" + ) + handle = self.handles.get(session) + if handle is not None: + return handle + folder_name, _, relpath = session.partition(":") + candidates = ( + [(self.folders[folder_name], relpath)] + if relpath and folder_name in self.folders + else [(folder, session) for folder in self.folders.values()] + ) + for folder, rel in candidates: + export_dir = folder.binaries.get(rel) + if export_dir is None: + continue + target = targets_module.export_target( + export_dir, folder=folder.name, relpath=rel, name=f"{folder.name}:{rel}" + ) + return self._adopt(target) + raise SessionError( + f"unknown session {session!r}; open sessions: " + f"{', '.join(self.handles) or 'none'}; folders: {', '.join(self.folders) or 'none'}" + ) + + # -- calling -------------------------------------------------------------- + + def call( + self, + session: str, + op: str, + args: dict[str, Any] | None = None, + *, + timeout: float | None = None, + ) -> Any: + handle = self.handle(session) + with handle.lock: + self._ensure_worker(handle) + handle.busy = True + try: + result = self._roundtrip(handle, ("call", op, args or {}), timeout) + if handle.pending_note and isinstance(result, dict): + result = {**result, "note": handle.pending_note} + handle.pending_note = None + return result + finally: + handle.busy = False + handle.last_used = time.monotonic() + + def _roundtrip(self, handle: Handle, message: tuple, timeout: float | None) -> Any: + request_id = next(self._ids) + kind, *rest = message + try: + handle.conn.send((kind, request_id, *rest)) + except (OSError, BrokenPipeError) as exc: + self._discard(handle, f"worker pipe closed: {exc}") + raise SessionError( + f"session {handle.target.name} lost its worker; retry the call" + ) from exc + limit = self.call_timeout if timeout is None else timeout + deadline = time.monotonic() + limit if limit and limit > 0 else None + while True: + wait = ( + 1.0 + if deadline is None + else max(0.0, min(1.0, deadline - time.monotonic())) + ) + try: + ready = handle.conn.poll(wait) + except (OSError, EOFError): + ready = False + if ready: + try: + reply = handle.conn.recv() + except (EOFError, OSError): + reply = None + if reply is None: + break + status, reply_id, payload = reply + if reply_id != request_id: + continue # a late answer to a call that timed out earlier + if status == "ok": + return payload + raise SessionError(str(payload)) + if not handle.alive: + break + if deadline is not None and time.monotonic() >= deadline: + self._discard(handle, f"call timed out after {limit:.0f}s") + raise SessionError( + f"{rest[0] if rest else kind} timed out after {limit:.0f}s; the " + f"session's worker was stopped (unsaved edits are lost) and " + "reopens on the next call" + ) + code = handle.process.exitcode if handle.process is not None else None + self._discard(handle, f"worker exited (code {code})") + hint = ( + " (killed by a signal, out of memory?)" + if code is not None and code < 0 + else "" + ) + raise SessionError( + f"session {handle.target.name}'s worker died during the call{hint}; " + "it reopens on the next call (unsaved edits are lost)" + ) + + # -- worker lifecycle --------------------------------------------------------- + + def _ensure_worker(self, handle: Handle) -> None: + if handle.alive and handle.conn is not None: + return + self._make_room(keep=handle) + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe(duplex=True) + log_path = self._log_path(handle.target) + from .worker import worker_main + + process = context.Process( + target=worker_main, + args=(handle.target, child, log_path, self.worker_options), + name=f"tocode-live-{handle.target.name}", + daemon=False, + ) + process.start() + child.close() + if process.pid is not None: + with _active_lock: + _active_children.add(process.pid) + self.log(f"live: opening {handle.target.name} ({handle.target.open_path})") + deadline = time.monotonic() + self.open_timeout + while time.monotonic() < deadline: + if parent.poll(1.0): + try: + status, payload = parent.recv() + except (EOFError, OSError): + break + if status == "ready": + handle.process, handle.conn = process, parent + handle.info = dict(payload) + if handle.note: + handle.pending_note = f"reopened after: {handle.note}" + handle.info["note"] = handle.pending_note + handle.note = None + handle.last_used = time.monotonic() + return + self._stop_process(process) + raise SessionError( + f"could not open {handle.target.open_path}: {payload}" + ) + if not process.is_alive(): + break + self._stop_process(process) + raise SessionError( + f"could not open {handle.target.open_path}: the worker exited or timed " + f"out (see {log_path})" + ) + + def _log_path(self, target: Target) -> Path: + if target.export_dir is not None: + return target.export_dir / ".tocode" / "live.log" + base = self.log_dir or Path.home() / ".cache" / "tocode" / "live" + return base / f"{target.name.replace('/', '_').replace(':', '_')}.log" + + def _make_room(self, *, keep: Handle) -> None: + with self._lock: + alive = [h for h in self.handles.values() if h.alive and h is not keep] + while len(alive) >= self.max_sessions: + idle = [h for h in alive if not h.busy] + if not idle: + raise SessionError( + f"all {self.max_sessions} sessions are busy; close one or raise " + "--max-sessions" + ) + victim = min(idle, key=lambda h: h.last_used) + self.log(f"live: closing {victim.target.name} to make room") + self._retire(victim, "closed to make room for another session") + alive.remove(victim) + + def close(self, session: str, *, save: bool = True) -> dict[str, Any]: + handle = self.handle(session) + result: dict[str, Any] = {"session": handle.target.name, "closed": True} + with handle.lock: + if handle.alive and handle.conn is not None: + try: + result.update( + self._roundtrip(handle, ("close", save), CLOSE_TIMEOUT) or {} + ) + except SessionError as exc: + result["error"] = str(exc) + self._discard(handle, "closed") + with self._lock: + self.handles.pop(handle.target.name, None) + return result + + def _discard(self, handle: Handle, why: str) -> None: + process = handle.process + handle.process, handle.conn = None, None + handle.note = why + if process is not None: + self._stop_process(process) + + def _stop_process(self, process: Any) -> None: + process.join(timeout=5) + if process.is_alive(): + kill_process_group(process.pid) + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join(timeout=5) + if process.pid is not None: + with _active_lock: + _active_children.discard(process.pid) + + def _reap_idle(self) -> None: + while not self._stop.wait(REAPER_SECONDS): + if self.idle_seconds <= 0: + continue + now = time.monotonic() + with self._lock: + idle = [ + h + for h in self.handles.values() + if h.alive and not h.busy and now - h.last_used > self.idle_seconds + ] + for handle in idle: + self.log(f"live: closing idle session {handle.target.name}") + self._retire(handle, "closed after being idle") + + def _retire(self, handle: Handle, why: str) -> None: + """Save and stop a worker but keep the session: it reopens on demand.""" + try: + with handle.lock: + if handle.alive and handle.conn is not None: + self._roundtrip(handle, ("close", True), CLOSE_TIMEOUT) + self._discard(handle, why) + except SessionError: + pass + + def describe(self) -> dict[str, Any]: + with self._lock: + return { + "sessions": [ + { + **h.target.describe(), + "alive": h.alive, + "idle_seconds": round(time.monotonic() - h.last_used), + } + for h in self.handles.values() + ], + "folders": [folder.describe() for folder in self.folders.values()], + "max_sessions": self.max_sessions, + } + + def shutdown(self) -> None: + self._stop.set() + for name in list(self.handles): + try: + self.close(name, save=True) + except ToCodeError: + continue diff --git a/src/tocode/live/setup.py b/src/tocode/live/setup.py new file mode 100644 index 0000000..83ed795 --- /dev/null +++ b/src/tocode/live/setup.py @@ -0,0 +1,267 @@ +"""`tocode sync`, `tocode skill install`, and MCP client registration.""" + +from __future__ import annotations + +import argparse +from importlib import resources +import json +from pathlib import Path +import shutil +import subprocess +import sys +from typing import Any + +from ..errors import ToCodeError + +SKILL_NAME = "tocode" + + +# -------------------------------------------------------------------------- +# tocode sync +# -------------------------------------------------------------------------- + + +def sync_main(argv: list[str]) -> int: + parser = argparse.ArgumentParser( + prog="tocode sync", + description="Update a ToCode export from its database (after editing it in IDA or Binary Ninja).", + ) + parser.add_argument("export", type=Path, help="The export folder.") + parser.add_argument( + "--database", + type=Path, + default=None, + help="Database to read instead of the export's own copy.", + ) + parser.add_argument( + "--full", + action="store_true", + help="Re-render every function (edits made outside ToCode are not tracked).", + ) + parser.add_argument("--binja-dir", type=Path, default=None) + parser.add_argument("-q", "--quiet", action="store_true") + args = parser.parse_args(argv) + + from .targets import Target, resolve + from .worker import open_session + + try: + target = resolve(args.export, database=args.database) + if not isinstance(target, Target) or target.export_dir is None: + raise ToCodeError(f"{args.export} is not a single-binary export") + live = open_session( + target, + log_path=target.export_dir / "tocode.log", + options={"binja_dir": str(args.binja_dir) if args.binja_dir else None}, + ) + try: + report = live.dispatch("sync", {"mode": "full" if args.full else "auto"}) + finally: + live.close(save=False) + except ToCodeError as exc: + print(f"tocode: {exc}", file=sys.stderr) + return 1 + if not args.quiet: + # Native and .NET reports differ: say what each one has. + parts = [ + f"{report['rerendered']} re-rendered", + f"{len(report.get('changed', []))} changed", + ] + for key in ("added", "removed", "moved"): + if report.get(key): + parts.append(f"{len(report[key])} {key}") + parts.append(f"{len(report['files_rewritten'])} files rewritten") + print( + f"Synced {report['export_dir']}: {', '.join(parts)} in {report['seconds']}s" + ) + return 0 + + +# -------------------------------------------------------------------------- +# MCP client registration +# -------------------------------------------------------------------------- + + +def server_command() -> list[str]: + """How MCP clients should start the server.""" + found = shutil.which("tocode") + if found: + return [found, "mcp"] + return [sys.executable, "-m", "tocode", "mcp"] + + +def config_text() -> str: + command = server_command() + head, rest = command[0], command[1:] + json_block = json.dumps( + {"mcpServers": {"tocode": {"command": head, "args": rest}}}, indent=2 + ) + toml_args = ", ".join(json.dumps(item) for item in rest) + return "\n".join( + [ + "Claude Code:", + f" claude mcp add --scope user tocode -- {' '.join(command)}", + "", + "Codex (~/.codex/config.toml):", + " [mcp_servers.tocode]", + f" command = {json.dumps(head)}", + f" args = [{toml_args}]", + "", + "Other clients (Claude Desktop, Cursor, VS Code, ...):", + *(" " + line for line in json_block.splitlines()), + ] + ) + + +def install_clients(clients: list[str]) -> list[str]: + done: list[str] = [] + command = server_command() + if "claude" in clients: + claude = shutil.which("claude") + if claude is None: + raise ToCodeError( + "the `claude` CLI was not found; run the command from `tocode mcp --config`" + ) + subprocess.run( + [claude, "mcp", "remove", "--scope", "user", "tocode"], + capture_output=True, + check=False, + ) + result = subprocess.run( + [claude, "mcp", "add", "--scope", "user", "tocode", "--", *command], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise ToCodeError( + f"claude mcp add failed: {(result.stderr or result.stdout).strip()}" + ) + done.append( + f"Claude Code: registered `tocode` (user scope) -> {' '.join(command)}" + ) + if "codex" in clients: + path = Path.home() / ".codex" / "config.toml" + done.append(f"Codex: {_write_codex_config(path, command)} {path}") + return done + + +def _write_codex_config(path: Path, command: list[str]) -> str: + block = [ + "[mcp_servers.tocode]", + f"command = {json.dumps(command[0])}", + f"args = [{', '.join(json.dumps(item) for item in command[1:])}]", + ] + try: + lines = path.read_text(encoding="utf-8").splitlines() + except FileNotFoundError: + lines = [] + start = next( + (i for i, line in enumerate(lines) if line.strip() == "[mcp_servers.tocode]"), + None, + ) + if start is None: + updated = lines + ([""] if lines and lines[-1].strip() else []) + block + action = "added [mcp_servers.tocode] to" + else: + end = next( + ( + i + for i in range(start + 1, len(lines)) + if lines[i].lstrip().startswith("[") + ), + len(lines), + ) + updated = lines[:start] + block + lines[end:] + action = "updated [mcp_servers.tocode] in" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(updated) + "\n", encoding="utf-8") + return action + + +# -------------------------------------------------------------------------- +# tocode skill install +# -------------------------------------------------------------------------- + + +def skill_destinations(agent: str, *, project: bool, path: Path | None) -> list[Path]: + if path is not None: + return [path.expanduser().resolve() / SKILL_NAME] + base = Path.cwd() if project else Path.home() + targets: list[Path] = [] + if agent in ("claude", "all"): + targets.append(base / ".claude" / "skills" / SKILL_NAME) + if agent in ("codex", "all"): + # The open Agent Skills location Codex (and other agents) read. + targets.append(base / ".agents" / "skills" / SKILL_NAME) + return targets + + +def _skill_files() -> dict[str, str]: + root = resources.files("tocode") / "skill" + files: dict[str, str] = {} + + def walk(node: Any, prefix: str) -> None: + for child in node.iterdir(): + if child.is_dir(): + walk(child, f"{prefix}{child.name}/") + elif child.name.endswith(".md"): + files[f"{prefix}{child.name}"] = child.read_text(encoding="utf-8") + + walk(root, "") + if "SKILL.md" not in files: + raise ToCodeError("the ToCode skill is missing from this installation") + return files + + +def install_skill(destinations: list[Path]) -> list[Path]: + files = _skill_files() + for destination in destinations: + if destination.exists(): + shutil.rmtree(destination) + for relative, text in files.items(): + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(text, encoding="utf-8") + return destinations + + +def skill_main(argv: list[str]) -> int: + parser = argparse.ArgumentParser( + prog="tocode skill", description="Install the ToCode agent skill." + ) + sub = parser.add_subparsers(dest="action", required=True) + install = sub.add_parser( + "install", help="Copy the skill where agents load skills from." + ) + install.add_argument( + "--agent", + choices=("claude", "codex", "all"), + default="all", + help="claude: .claude/skills; codex: .agents/skills (Codex and other agents); all: both (default).", + ) + install.add_argument( + "--project", + action="store_true", + help="Install into the current folder instead of your home folder.", + ) + install.add_argument( + "--path", + type=Path, + default=None, + help="Install into PATH/tocode instead (any other agent).", + ) + sub.add_parser("show", help="Print SKILL.md.") + args = parser.parse_args(argv) + try: + if args.action == "show": + print(_skill_files()["SKILL.md"]) + return 0 + for path in install_skill( + skill_destinations(args.agent, project=args.project, path=args.path) + ): + print(f"Installed the ToCode skill to {path}") + except ToCodeError as exc: + print(f"tocode: {exc}", file=sys.stderr) + return 1 + return 0 diff --git a/src/tocode/live/sync.py b/src/tocode/live/sync.py new file mode 100644 index 0000000..d99b53f --- /dev/null +++ b/src/tocode/live/sync.py @@ -0,0 +1,771 @@ +"""``sync``: update an export from its database, like an incremental rebuild. + +1. Read what the export says (functions.json, types.json, the function index, + the file layout in ``.tocode/sync/manifest.json``). +2. Collect the database's inventory again and diff: functions added, + removed, renamed or retyped; types changed; globals renamed. +3. Decide what to re-render: edits made in this session, the diff, and what + they affect (callers of renamed/retyped functions, functions whose code + mentions a changed type or an old global name). +4. Keep every existing function in its file; put a new one next to a caller + (or callee), else in ``discovered.c``. +5. Re-render, rewrite only the files whose content changed, and rewrite the + JSON and headers. Untouched source files keep their bytes. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +import os +from pathlib import Path +import re +import time +from typing import TYPE_CHECKING, Any + +from ..metadata import write_text_atomic +from ..naming import SHARED_CLUSTER_ID, asm_file_name, c_file_name, summary_file_name +from ..schema import Cluster, FunctionRange, RenderedFunction +from .. import syncstate +from .worker import OpError + +if TYPE_CHECKING: + from .worker import LiveSession + +DISCOVERED = "discovered" +_PREAMBLE_ROOT_RX = re.compile(r"/\* Cluster root: (.*) \(0x([0-9a-fA-F]+)\) \*/") +_PREAMBLE_DESC_RX = re.compile(r"/\* Description: (.*) \*/") +_CLUSTER_FILE_RX = re.compile(r"^cluster_([0-9a-f]{16})\.c$") + + +def register(live: "LiveSession") -> None: + live.ops["sync"] = lambda args: sync(live, args) + live.ops["export"] = lambda args: export(live, args) + + +# -------------------------------------------------------------------------- +# What the export says +# -------------------------------------------------------------------------- + + +@dataclass +class DiskState: + root: Path + functions: dict[int, dict[str, Any]] = field(default_factory=dict) + ranges: dict[int, FunctionRange] = field(default_factory=dict) + tree_ranges: dict[int, FunctionRange] = field(default_factory=dict) + types: dict[str, str] = field(default_factory=dict) + variables: dict[int, str] = field(default_factory=dict) + clusters: list[Cluster] = field(default_factory=list) + tree: bool = False + entropy: bool = False + origin: str | None = None + origin_known: bool = False + database: str | None = None + + @classmethod + def load(cls, root: Path) -> "DiskState": + state = cls(root=root) + for row in _read(root / "functions.json").get("functions", []): + address = _hex(row.get("address")) + if address is not None: + state.functions[address] = row + state.ranges = _ranges(root, root / "function-index.json", state.functions) + state.tree_ranges = _ranges( + root, root / "function-index-tree.json", state.functions + ) + for row in _read(root / "types.json").get("types", []): + if row.get("name"): + state.types[str(row["name"])] = str(row.get("c_decl") or "") + variables = _read(root / "data" / "variables.json").get("variables", {}) + if isinstance(variables, dict): + for name, row in variables.items(): + address = _hex((row or {}).get("va")) + if address is not None: + state.variables[address] = str(name) + project = _read(root / "project.json") + state.database = project.get("ida_database") + state.tree = bool(project.get("tree_src_dir")) + manifest = syncstate.read_manifest(root) + if manifest is not None: + state.clusters = [ + syncstate.cluster_from_json(item) + for item in manifest.get("clusters", []) + ] + state.tree = bool(manifest.get("tree")) + state.entropy = bool(manifest.get("entropy")) + state.origin = manifest.get("origin") + state.origin_known = True + else: + state.clusters = _clusters_from_files(root, state.ranges) + return state + + +def _read(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return payload if isinstance(payload, dict) else {} + + +def _hex(value: Any) -> int | None: + try: + return int(str(value), 16) + except (TypeError, ValueError): + return None + + +def _local_path(root: Path, recorded: Any) -> Path: + """A path the index recorded, re-rooted if the export has been moved.""" + path = Path(str(recorded or "")) + try: + path.resolve().relative_to(root.resolve()) + return path + except ValueError: + parts = path.parts + if "src" in parts: + return root.joinpath(*parts[parts.index("src") :]) + return root / path.name + + +def _ranges( + root: Path, index: Path, functions: dict[int, dict[str, Any]] +) -> dict[int, FunctionRange]: + result: dict[int, FunctionRange] = {} + for row in _read(index).get("functions", []): + address = _hex(row.get("address")) + if address is None: + continue + c = row.get("c") or {} + asm = row.get("asm") or {} + info = functions.get(address, {}) + origin = row.get("origin") or {} + result[address] = FunctionRange( + address=address, + name=str(row.get("name") or info.get("name") or f"sub_{address:x}"), + c_file=_local_path(root, c.get("path")).resolve(), + c_line_start=int(c.get("line_start") or 0), + c_line_end=int(c.get("line_end") or 0), + asm_file=_local_path(root, asm.get("path")).resolve(), + asm_line_start=int(asm.get("line_start") or 0), + asm_line_end=int(asm.get("line_end") or 0), + arg_count=info.get("nargs"), + local_count=info.get("nlocals"), + origin_file=origin.get("file"), + origin_line=origin.get("line"), + ) + return result + + +def _clusters_from_files(root: Path, ranges: dict[int, FunctionRange]) -> list[Cluster]: + """The file layout of an export made before sync manifests existed.""" + raw = (root / "src" / "raw").resolve() + by_file: dict[Path, list[FunctionRange]] = {} + for item in ranges.values(): + by_file.setdefault(item.c_file, []).append(item) + clusters: list[Cluster] = [] + for path, members in sorted(by_file.items()): + members.sort(key=lambda item: item.c_line_start) + try: + folder = path.parent.relative_to(raw).as_posix() + except ValueError: + folder = "app" + root_address, label, summary = members[0].address, "", "" + try: + head = path.read_text(encoding="utf-8", errors="replace").split("\n", 6) + except OSError: + head = [] + for line in head: + if (match := _PREAMBLE_ROOT_RX.match(line)) is not None: + label, root_address = match.group(1), int(match.group(2), 16) + elif (match := _PREAMBLE_DESC_RX.match(line)) is not None: + summary = match.group(1) + file_base = None + if path.name == "utils.c": + root_address = SHARED_CLUSTER_ID + elif (match := _CLUSTER_FILE_RX.match(path.name)) is not None: + root_address = int(match.group(1), 16) + else: + file_base = path.stem + clusters.append( + Cluster( + root=root_address, + label=label, + summary=summary, + members=[item.address for item in members], + folder=folder or "app", + file_base=file_base, + ) + ) + return clusters + + +# -------------------------------------------------------------------------- +# Sync +# -------------------------------------------------------------------------- + + +def sync(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + from ..exporter import _build_clusters, render_one_with_limits + + root = live.target.export_dir + if root is None: + raise OpError("this session has no export folder; call `export` first") + started = time.monotonic() + mode = str(args.get("mode") or "auto") + forced = {live.resolve(item) for item in args.get("addresses") or []} + disk = DiskState.load(root) + _forget_decompilations(live) + live.refresh() + analysis = live.analysis + names = live.names + + wanted = { + a + for cluster in _build_clusters(analysis, live.analyzer) + for a in cluster.members + } + wanted = {a for a in wanted if a in analysis.routines} + old = set(disk.ranges) + added = sorted(wanted - old) + removed = sorted(old - wanted) + common = wanted & old + + renamed: list[dict[str, Any]] = [] + retyped: list[str] = [] + dirty: set[int] = set() + old_tokens: set[str] = set() + for address in sorted(common): + routine = analysis.routines[address] + row = disk.functions.get(address, {}) + if row and row.get("name") != routine.name: + renamed.append( + {"address": hex(address), "old": row.get("name"), "new": routine.name} + ) + old_tokens.add(str(row.get("name"))) + dirty.add(address) + dirty.update(analysis.callers.get(address, [])) + if row and _shape(row) != _routine_shape(routine): + retyped.append(hex(address)) + dirty.add(address) + dirty.update(analysis.callers.get(address, [])) + + new_types = {item.name: item.c_decl for item in analysis.types} + types_added = sorted(set(new_types) - set(disk.types)) + types_removed = sorted(set(disk.types) - set(new_types)) + types_modified = sorted( + name + for name in set(new_types) & set(disk.types) + if new_types[name] != disk.types[name] + ) + changed_types = types_added + types_removed + types_modified + renamed_globals = _renamed_globals(disk, analysis) + tokens = set(changed_types) | set(renamed_globals) | old_tokens + if tokens: + dirty.update(_mentions(disk, tokens)) + for address in common: + routine = analysis.routines[address] + used = " ".join(t for _, t in routine.params + routine.local_vars) + if any( + re.search(rf"\b{re.escape(token)}\b", used) for token in changed_types + ): + dirty.add(address) + + # A function appearing or vanishing changes how its callers decompile + # (a named call instead of a raw address, or the reverse). + for address in added: + dirty.update(analysis.callers.get(address, [])) + for address in removed: + for caller in disk.functions.get(address, {}).get("callers") or []: + if (value := _hex(caller)) is not None: + dirty.add(value) + + dirty |= live.dirty | forced + if mode == "full": + dirty |= wanted + dirty = (dirty & wanted) | set(added) + + clusters, membership_changed = _place(disk.clusters, removed, added, analysis) + fresh: dict[int, RenderedFunction] = {} + for address in sorted(dirty): + fresh[address] = render_one_with_limits( + live.session, analysis, analysis.routines[address], names + ) + changed = sorted( + address + for address, item in fresh.items() + if (old_item := syncstate.load_render(root, address)) is None + or not _same(old_item, item) + ) + + touched = [ + cluster + for cluster in clusters + if id(cluster) in membership_changed + or any(a in changed for a in cluster.members) + ] + # Files being rewritten need every member's render: the kept one when the + # function did not change, a new one for exports made without kept renders. + rendered: dict[int, RenderedFunction] = dict(fresh) + for cluster in touched: + for address in cluster.members: + if address in rendered: + continue + kept = syncstate.load_render(root, address) + if kept is None: + kept = render_one_with_limits( + live.session, analysis, analysis.routines[address], names + ) + fresh[address] = kept + rendered[address] = kept + + written = _rewrite(live, root, disk, clusters, touched, rendered, removed) + _write_all_metadata(live, root, disk, clusters, rendered, written) + for address, item in fresh.items(): + syncstate.save_render(root, item) + for address in removed: + syncstate.drop_render(root, address) + syncstate.write_manifest( + root, + clusters=clusters, + backend=live.analyzer.backend_name, + decompiler=live.analyzer.decompiler_label, + tree=disk.tree, + entropy=disk.entropy, + origin=disk.origin, + ) + links = _refresh_folder_links(live, root) if touched or added or removed else None + + saved: dict[str, Any] | None = None + if live._edited: + saved = live.save() + live.dirty.clear() + live.unknown_changes = live.types_changed = live.globals_changed = False + live.functions_changed = False + live._locations = _reload_locations(root) + report = { + "export_dir": str(root), + "mode": mode, + "added": [_brief(live, a) for a in added], + "removed": [hex(a) for a in removed], + "renamed": renamed, + "retyped": retyped, + "types_added": types_added, + "types_removed": types_removed, + "types_modified": types_modified, + "globals_renamed": renamed_globals, + "rerendered": len(fresh), + "changed": [_brief(live, a) for a in changed], + "files_rewritten": sorted(_rel(root, path) for path in written["files"]), + "files_removed": sorted(_rel(root, path) for path in written["removed"]), + "failures": [ + {"address": hex(item.address), "error": item.failure.message} + for item in fresh.values() + if item.failure is not None + ], + "folder_links": links, + "database": saved, + "seconds": round(time.monotonic() - started, 2), + } + _log(root, report) + return report + + +def _shape(row: dict[str, Any]) -> tuple: + return ( + row.get("return_type"), + tuple((p.get("name"), p.get("type")) for p in row.get("params") or []), + tuple((p.get("name"), p.get("type")) for p in row.get("locals") or []), + row.get("size"), + ) + + +def _routine_shape(routine: Any) -> tuple: + return ( + routine.return_type, + tuple(routine.params), + tuple(routine.local_vars), + routine.size, + ) + + +def _same(left: RenderedFunction, right: RenderedFunction) -> bool: + return ( + left.c_text == right.c_text + and left.asm_text == right.asm_text + and left.summary_text == right.summary_text + and left.prototype == right.prototype + ) + + +def _renamed_globals(disk: DiskState, analysis: Any) -> list[str]: + """Old names of globals whose name changed (as data/variables.json names them).""" + from ..metadata import variables_document + + current: dict[int, str] = {} + for name, row in variables_document(analysis).items(): + address = _hex(row.get("va")) + if address is not None: + current[address] = name + return sorted( + name + for address, name in disk.variables.items() + if address in current and current[address] != name + ) + + +def _mentions(disk: DiskState, tokens: set[str]) -> set[int]: + """Functions whose exported C mentions any of ``tokens`` as a whole word.""" + pattern = re.compile( + r"\b(?:" + "|".join(re.escape(t) for t in sorted(tokens) if t) + r")\b" + ) + hits: set[int] = set() + by_file: dict[Path, list[FunctionRange]] = {} + for item in disk.ranges.values(): + by_file.setdefault(item.c_file, []).append(item) + for path, items in by_file.items(): + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + continue + for item in items: + text = "\n".join(lines[max(0, item.c_line_start - 1) : item.c_line_end]) + if pattern.search(text): + hits.add(item.address) + return hits + + +def _place( + clusters: list[Cluster], removed: list[int], added: list[int], analysis: Any +) -> tuple[list[Cluster], set[int]]: + """Existing functions keep their file; new ones join a caller's file.""" + gone = set(removed) + changed: set[int] = set() + result: list[Cluster] = [] + home: dict[int, Cluster] = {} + for cluster in clusters: + members = [a for a in cluster.members if a not in gone] + copy = Cluster( + root=cluster.root, + label=cluster.label, + summary=cluster.summary, + members=members, + folder=cluster.folder, + file_base=cluster.file_base, + ) + if len(members) != len(cluster.members): + changed.add(id(copy)) + result.append(copy) + for address in members: + home[address] = copy + discovered = next((c for c in result if c.file_base == DISCOVERED), None) + pending = list(added) + # Several passes so a chain of new functions follows its first placed caller. + for _ in range(3): + left = [] + for address in pending: + neighbours = list(analysis.callers.get(address, [])) + list( + analysis.callees.get(address, []) + ) + target = next((home[n] for n in neighbours if n in home), None) + if target is None: + left.append(address) + continue + target.members.append(address) + home[address] = target + changed.add(id(target)) + if len(left) == len(pending): + break + pending = left + if pending: + if discovered is None: + discovered = Cluster( + root=pending[0], + label=DISCOVERED, + summary="Functions found after the export (added by sync)", + members=[], + folder="app", + file_base=DISCOVERED, + ) + result.append(discovered) + discovered.members.extend(pending) + changed.add(id(discovered)) + return result, changed + + +def _rewrite( + live: "LiveSession", + root: Path, + disk: DiskState, + clusters: list[Cluster], + touched: list[Cluster], + rendered: dict[int, RenderedFunction], + removed: list[int], +) -> dict[str, Any]: + from ..exporter import _cluster_path, build_cluster_files, write_tree_sources + + raw_dir = root / "src" / "raw" + include_dir = root / "include" + header_name = _header_name(live) + files: list[Path] = [] + gone: list[Path] = [] + ranges: dict[int, FunctionRange] = {} + prototypes: dict[int, str] = {} + failures: list[Any] = [] + for cluster in touched: + c_path = _cluster_path(raw_dir, cluster, c_file_name(cluster)) + asm_path = _cluster_path(raw_dir, cluster, asm_file_name(cluster)) + summary_path = _cluster_path(raw_dir, cluster, summary_file_name(cluster)) + if not cluster.members: + for path in (c_path, asm_path, summary_path): + if path.exists(): + path.unlink() + gone.append(path) + continue + c_path.parent.mkdir(parents=True, exist_ok=True) + include = Path( + os.path.relpath(include_dir / header_name, c_path.parent) + ).as_posix() + block = build_cluster_files( + analysis=live.analysis, + cluster=cluster, + header_include=include, + c_path=c_path, + asm_path=asm_path, + summary_path=summary_path, + rendered=rendered, + prototypes=prototypes, + ) + for path, text in ( + (c_path, block["c"]), + (asm_path, block["asm"]), + (summary_path, block["summary"]), + ): + if not path.exists() or path.read_text(encoding="utf-8") != text: + write_text_atomic(path, text) + files.append(path) + for item in block["ranges"]: + ranges[item.address] = item + failures.extend(block["failures"]) + tree_ranges: dict[int, FunctionRange] = {} + live_touched = [cluster for cluster in touched if cluster.members] + if disk.tree and live_touched: + all_ranges = {**disk.ranges, **ranges} + written = write_tree_sources( + analysis=live.analysis, + clusters=live_touched, + src_dir=root / "src" / "tree", + raw_dir=raw_dir, + include_dir=include_dir, + rendered=rendered, + raw_ranges=list(all_ranges.values()), + ) + files.extend(written["sources"]) + for item in written["ranges"]: + tree_ranges[item.address] = item + return { + "files": files, + "removed": gone, + "ranges": ranges, + "tree_ranges": tree_ranges, + "failures": failures, + } + + +def _header_name(live: "LiveSession") -> str: + from ..naming import clean_path_component + + return f"{clean_path_component(live.analysis.binary.path.stem)}.h" + + +def _write_all_metadata( + live: "LiveSession", + root: Path, + disk: DiskState, + clusters: list[Cluster], + rendered: dict[int, RenderedFunction], + written: dict[str, Any], +) -> None: + from ..exporter import ExportContext, _cluster_path, _write_metadata + + raw_dir = root / "src" / "raw" + order = [a for cluster in clusters for a in cluster.members] + ranges = {**disk.ranges, **written["ranges"]} + tree_ranges = {**disk.tree_ranges, **written["tree_ranges"]} + prototypes: dict[int, str] = {} + failures = [] + for address in order: + item = rendered.get(address) + if item is not None: + prototypes[address] = item.prototype + if item.failure is not None: + failures.append(item.failure) + else: + row = disk.functions.get(address, {}) + prototypes[address] = str(row.get("prototype") or "") + header_name = _header_name(live) + context = ExportContext( + analyzer=live.analyzer, + progress=live.progress, + out_dir=root, + jobs=None, + tree_enabled=disk.tree, + entropy_enabled=disk.entropy, + origin=disk.origin, + ) + context.analysis = live.analysis + context.root = root + context.raw_dir = raw_dir + context.tree_dir = root / "src" / "tree" if disk.tree else None + context.include_dir = root / "include" + context.data_dir = root / "data" + context.header_name = header_name + context.header_path = context.include_dir / header_name + context.names = live.names + context.clusters = [cluster for cluster in clusters if cluster.members] + context.addresses = order + context.prototypes = prototypes + context.failures = failures + context.raw_ranges = [ranges[a] for a in order if a in ranges] + context.tree_ranges = [tree_ranges[a] for a in order if a in tree_ranges] + files = [ + _cluster_path(raw_dir, cluster, c_file_name(cluster)) + for cluster in context.clusters + ] + context.raw_sources = [path.resolve() for path in files] + context.asm_files = [ + _cluster_path(raw_dir, cluster, asm_file_name(cluster)).resolve() + for cluster in context.clusters + ] + context.summary_files = [ + _cluster_path(raw_dir, cluster, summary_file_name(cluster)).resolve() + for cluster in context.clusters + ] + context.ida_database = live.target.database or ( + Path(disk.database) if disk.database else None + ) + # Keep what the export run chose; a sync renders in one process. + previous = _read(root / "project.json") + context.worker_count = int(previous.get("worker_count") or 1) + requested = previous.get("requested_worker_count") + context.requested_jobs = requested if isinstance(requested, int) else None + context.render_mode = str(previous.get("parallel_mode") or "single") + _write_metadata(context, publish=False, docs=disk.origin_known) + + +def _refresh_folder_links(live: "LiveSession", root: Path) -> dict[str, Any] | None: + """Inside a folder export, rejoin the cross-binary links (and annotations).""" + master = next((p for p in root.parents if (p / ".tocode-folder").is_dir()), None) + if master is None: + return None + project = _read(master / "project.json") + source = project.get("input") + if not source or not Path(str(source)).is_dir(): + return {"refreshed": False, "reason": "the folder's input is gone"} + from ..apk_native import NativeOptions + from ..folder import FolderOptions, export_folder + from ..progress import Progress + + options = project.get("options") or {} + summary = export_folder( + Path(str(source)), + options=FolderOptions( + out_dir=master, + only=list(options.get("only") or []), + exclude=list(options.get("exclude") or []), + include_framework=bool(options.get("include_framework")), + native=bool(options.get("native", True)), + managed=bool(options.get("managed", True)), + links_only=True, + native_options=NativeOptions( + backend=options.get("backend") or "auto", + tree=bool(options.get("tree")), + entropy=bool(options.get("entropy")), + ), + ), + progress=Progress(enabled=False), + ) + return {"refreshed": True, "folder": str(master), "mapped": summary.mapped} + + +def _brief(live: "LiveSession", address: int) -> dict[str, Any]: + return {"address": hex(address), "name": live.name_of(address)} + + +def _rel(root: Path, path: Path) -> str: + try: + return Path(path).resolve().relative_to(root.resolve()).as_posix() + except ValueError: + return str(path) + + +def _reload_locations(root: Path) -> dict[int, dict[str, Any]]: + from .worker import _load_locations + + return _load_locations(root) + + +def _log(root: Path, report: dict[str, Any]) -> None: + path = syncstate.sync_dir(root) / "sync-log.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + entry = { + "time": time.strftime("%Y-%m-%dT%H:%M:%S"), + **{k: report[k] for k in report if k != "changed"}, + } + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(entry) + "\n") + + +def _forget_decompilations(live: "LiveSession") -> None: + """Edits made by scripts may not have invalidated decompiler caches.""" + if live.backend == "ida": + try: + import ida_hexrays # type: ignore[import-not-found] + + ida_hexrays.clear_cached_cfuncs() + except Exception: # noqa: BLE001 + pass + for attribute in ("_decompile_cache", "_disasm_cache", "_locals_cache"): + cache = getattr(live.session, attribute, None) + if isinstance(cache, dict): + cache.clear() + + +# -------------------------------------------------------------------------- +# Export (a session opened on a plain binary) +# -------------------------------------------------------------------------- + + +def export(live: "LiveSession", args: dict[str, Any]) -> dict[str, Any]: + from ..exporter import export_binary + from ..naming import default_output_name + + if live.target.export_dir is not None and not args.get("out_dir"): + raise OpError( + f"this session already has an export at {live.target.export_dir}; use `sync`" + ) + out = Path(str(args.get("out_dir") or "")) if args.get("out_dir") else None + if out is None: + out = live.target.binary.parent / default_output_name(live.target.binary) + previous = os.environ.get("TOCODE_FUNCTION_TIMEOUT_SECONDS") + # Render in this process, against the live session (no worker copies). + os.environ["TOCODE_FUNCTION_TIMEOUT_SECONDS"] = "0" + try: + summary = export_binary( + live.analyzer, out_dir=out, progress=live.progress, jobs=1 + ) + finally: + if previous is None: + os.environ.pop("TOCODE_FUNCTION_TIMEOUT_SECONDS", None) + else: + os.environ["TOCODE_FUNCTION_TIMEOUT_SECONDS"] = previous + live.target.export_dir = summary.root_dir + live._locations = _reload_locations(summary.root_dir) + live.dirty.clear() + return { + "export_dir": str(summary.root_dir), + "functions": summary.function_count, + "clusters": summary.cluster_count, + "failures": len(summary.failed_functions), + } diff --git a/src/tocode/live/targets.py b/src/tocode/live/targets.py new file mode 100644 index 0000000..3833ea8 --- /dev/null +++ b/src/tocode/live/targets.py @@ -0,0 +1,294 @@ +"""What ``open(path, database?, backend?)`` refers to. + +A path is one of: + +- an export folder (``project.json``): the export's own database + (``.i64``/``.idb``/``.bndb`` next to it) unless ``database`` names + another; without a database, the original binary with the export's backend; +- a folder master (``.tocode-folder/``): every exported binary in it, each + opened on demand as its own session; +- a binary or a database: opened directly; its default export folder + (``_decompiler``) is attached when it exists. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from pathlib import Path +from typing import Any + +from ..backends.base import IDA_DATABASE_SUFFIXES, is_binja_database, is_ida_database +from ..errors import ToCodeError +from ..naming import default_output_name + +DATABASE_SUFFIXES = (*sorted(IDA_DATABASE_SUFFIXES), ".bndb") +FOLDER_STATE = Path(".tocode-folder") / "state.json" + + +@dataclass(slots=True) +class Target: + """One session's worth: what the worker opens and where its export lives.""" + + name: str + open_path: Path # what create_analyzer opens (a database or the binary) + binary: Path # the original binary (may no longer exist for a database) + backend: str + export_dir: Path | None = None + database: Path | None = None # the database edits are saved into, if any + folder: str | None = None # owning folder session, for folder binaries + relpath: str | None = None + + def describe(self) -> dict[str, Any]: + return { + "session": self.name, + "binary": str(self.binary), + "opened": str(self.open_path), + "database": str(self.database) if self.database else None, + "backend": self.backend, + "export_dir": str(self.export_dir) if self.export_dir else None, + "folder": self.folder, + "relpath": self.relpath, + } + + +@dataclass(slots=True) +class FolderTarget: + """A folder master: binaries by relative path, each with its export folder.""" + + name: str + master: Path + input_dir: Path | None + binaries: dict[str, Path] = field(default_factory=dict) + statuses: dict[str, str] = field(default_factory=dict) + + def describe(self) -> dict[str, Any]: + return { + "session": self.name, + "folder": str(self.master), + "input_dir": str(self.input_dir) if self.input_dir else None, + "binaries": [ + {"relpath": relpath, "status": self.statuses.get(relpath)} + for relpath in self.binaries + ], + } + + +def is_database(path: Path) -> bool: + return is_ida_database(path) or is_binja_database(path) + + +def is_export_dir(path: Path) -> bool: + return (path / "project.json").is_file() and (path / "functions.json").is_file() + + +def is_folder_master(path: Path) -> bool: + return (path / FOLDER_STATE).is_file() + + +def resolve( + path: Path, + *, + database: Path | None = None, + backend: str | None = None, + name: str | None = None, +) -> Target | FolderTarget: + path = path.expanduser().resolve() + db = database.expanduser().resolve() if database is not None else None + if db is not None and not db.is_file(): + raise ToCodeError(f"database not found: {db}") + if path.is_dir(): + if is_folder_master(path): + if db is not None: + raise ToCodeError( + "a folder holds many binaries; pass `database` when opening " + "one of them (open its export folder, or `session` = its path)" + ) + return folder_target(path, name=name) + if is_export_dir(path): + return export_target(path, database=db, backend=backend, name=name) + raise ToCodeError( + f"{path} is neither a ToCode export nor a folder export; export it " + "first (`tocode `), or open the binary itself" + ) + if not path.is_file(): + raise ToCodeError(f"not found: {path}") + return file_target(path, database=db, backend=backend, name=name) + + +def export_target( + export_dir: Path, + *, + database: Path | None = None, + backend: str | None = None, + name: str | None = None, + folder: str | None = None, + relpath: str | None = None, +) -> Target: + project = _read_json(export_dir / "project.json") + binary = Path(str(project.get("binary") or export_dir.name)) + recorded = project.get("backend") or "auto" + if recorded == "dotnet": + # A .NET export: the database is its edit log, replayed onto the input. + if not binary.is_file(): + raise ToCodeError( + f"the .NET export at {export_dir} needs its input {binary}, which is gone" + ) + return Target( + name=name or relpath or _session_name(binary), + open_path=binary, + binary=binary, + backend="dotnet", + export_dir=export_dir, + database=database or dotnet_edit_log(export_dir), + folder=folder, + relpath=relpath, + ) + own = database or own_database(export_dir, binary, project) + if own is not None: + open_path = own + chosen = _backend_for_database(own) + else: + if not binary.is_file(): + raise ToCodeError( + f"the export at {export_dir} has no database and its binary " + f"{binary} is gone; pass `database`" + ) + open_path = binary + chosen = backend or str(recorded) + return Target( + name=name or relpath or _session_name(binary), + open_path=open_path, + binary=binary, + backend=chosen, + export_dir=export_dir, + database=own, + folder=folder, + relpath=relpath, + ) + + +def file_target( + path: Path, + *, + database: Path | None = None, + backend: str | None = None, + name: str | None = None, +) -> Target: + if ( + not is_database(path) + and (backend in (None, "auto", "dotnet")) + and _is_dotnet(path) + ): + export_dir = path.parent / default_output_name(path) + attached = None + if is_export_dir(export_dir): + recorded = Path( + str(_read_json(export_dir / "project.json").get("binary") or "") + ) + if recorded.resolve() == path.resolve(): + attached = export_dir + return Target( + name=name or _session_name(path), + open_path=path, + binary=path, + backend="dotnet", + export_dir=attached, + database=database or (dotnet_edit_log(attached) if attached else None), + ) + if is_database(path): + database = database or path + export_dir = path.parent / default_output_name(path) + if is_export_dir(export_dir): + project = _read_json(export_dir / "project.json") + recorded = Path(str(project.get("binary") or "")) + # Only attach the export made from this very file. + if recorded.resolve() != path.resolve() and not is_database(path): + export_dir = None # type: ignore[assignment] + else: + export_dir = None # type: ignore[assignment] + open_path = database or path + return Target( + name=name or _session_name(path), + open_path=open_path, + binary=path, + backend=_backend_for_database(open_path) + if is_database(open_path) + else (backend or "auto"), + export_dir=export_dir, + database=database, + ) + + +def folder_target(master: Path, *, name: str | None = None) -> FolderTarget: + from ..folder import DONE, export_dir_for + + state = _read_json(master / FOLDER_STATE) + project = _read_json(master / "project.json") + input_dir = project.get("input") + target = FolderTarget( + name=name or master.name.removesuffix("_decompiled") or master.name, + master=master, + input_dir=Path(str(input_dir)) if input_dir else None, + ) + jobs = state.get("jobs") or {} + for relpath in sorted(jobs): + record = jobs[relpath] if isinstance(jobs[relpath], dict) else {} + status = str(record.get("status") or "") + export_dir = export_dir_for(master, relpath) + target.statuses[relpath] = status + if status == DONE and is_export_dir(export_dir): + target.binaries[relpath] = export_dir + return target + + +def own_database( + export_dir: Path, binary: Path, project: dict[str, Any] +) -> Path | None: + """The database an export ships (``.i64`` / ``.idb`` / ``.bndb``).""" + recorded = project.get("ida_database") + candidates: list[Path] = [] + if recorded: + candidates.append(export_dir / Path(str(recorded)).name) + stems = {binary.stem, binary.name} + for stem in stems: + candidates.extend( + export_dir / f"{stem}{suffix}" for suffix in DATABASE_SUFFIXES + ) + for candidate in candidates: + if candidate.is_file() and candidate.stat().st_size > 0: + return candidate.resolve() + return None + + +def dotnet_edit_log(export_dir: Path) -> Path: + return export_dir / ".tocode" / "dotnet" / "edits.json" + + +def _is_dotnet(path: Path) -> bool: + try: + from ..backends.dotnet import is_dotnet_input + + return bool(is_dotnet_input(path)) + except Exception: # noqa: BLE001 - not readable as .NET + return False + + +def _backend_for_database(path: Path) -> str: + return "ida" if is_ida_database(path) else "binja" + + +def _session_name(path: Path) -> str: + stem = path.name + for suffix in DATABASE_SUFFIXES: + if stem.lower().endswith(suffix): + stem = stem[: -len(suffix)] + return stem or path.name + + +def _read_json(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return payload if isinstance(payload, dict) else {} diff --git a/src/tocode/live/tools.py b/src/tocode/live/tools.py new file mode 100644 index 0000000..daf8e9f --- /dev/null +++ b/src/tocode/live/tools.py @@ -0,0 +1,561 @@ +"""The MCP tool table: names, schemas, annotations, and what each one runs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from typing import Any, Callable + +from ..errors import ToCodeError +from .jsonrpc import ToolCall +from .sessions import SessionManager + +SESSION = { + "type": "string", + "description": "Session id from `open` (or a folder binary's path such as " + "`bin/app`; the only open session when omitted).", +} +TARGET = { + "type": "string", + "description": "Function or symbol: an address (`0x401000`) or a name. " + "For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member`.", +} +FILTER = { + "type": "string", + "description": "Case-insensitive regular expression (or substring) to match.", +} +OFFSET = {"type": "integer", "minimum": 0, "description": "First row to return."} +LIMIT = { + "type": "integer", + "minimum": 1, + "maximum": 2000, + "description": "Rows per page.", +} + +QUERY, EDIT, EVAL, SESSIONS, PROJECT = "query", "edit", "eval", "sessions", "project" + + +@dataclass +class Tool: + name: str + description: str + group: str + properties: dict[str, Any] = field(default_factory=dict) + required: list[str] = field(default_factory=list) + op: str | None = None # worker operation (most tools) + run: Callable[[SessionManager, dict[str, Any], ToolCall], Any] | None = None + timeout: float | None = None + destructive: bool = False + text_field: str | None = None # result field shown as plain text + + @property + def read_only(self) -> bool: + return self.group in (QUERY, SESSIONS) and not self.destructive + + def schema(self) -> dict[str, Any]: + properties = dict(self.properties) + if self.op is not None: + properties = {"session": SESSION, **properties} + return { + "name": self.name, + "description": self.description, + "inputSchema": { + "type": "object", + "properties": properties, + "required": list(self.required), + "additionalProperties": False, + }, + "annotations": { + "readOnlyHint": self.read_only, + "destructiveHint": self.destructive, + "idempotentHint": self.read_only, + "openWorldHint": False, + }, + } + + +def _open(manager: SessionManager, args: dict[str, Any], _call: ToolCall) -> Any: + return manager.open( + str(args.get("path") or ""), + database=args.get("database") or None, + backend=args.get("backend") or None, + ) + + +def _sessions(manager: SessionManager, _args: dict[str, Any], _call: ToolCall) -> Any: + return manager.describe() + + +def _close(manager: SessionManager, args: dict[str, Any], _call: ToolCall) -> Any: + return manager.close( + str(args.get("session") or ""), save=args.get("save", True) is not False + ) + + +def build_tools() -> list[Tool]: + return [ + # -- sessions --------------------------------------------------------- + Tool( + "open", + "Open a binary or database and keep it alive for the other tools. `path` " + "is an export folder (made by `tocode `), a folder export, a " + "binary, or a database (.i64/.idb/.bndb). By default an export's own " + "database is used; pass `database` to edit another one. Returns the " + "session id (a folder returns its binaries; address them as " + "`session`=).", + SESSIONS, + { + "path": { + "type": "string", + "description": "Export folder, folder export, binary, or database.", + }, + "database": { + "type": "string", + "description": "Database to open instead of the export's own copy.", + }, + "backend": { + "type": "string", + "enum": ["auto", "ida", "binja", "r2", "angr"], + "description": "Backend for a plain binary (databases pick their own).", + }, + }, + ["path"], + run=_open, + timeout=None, + ), + Tool("sessions", "List open sessions and folders.", SESSIONS, run=_sessions), + Tool( + "close", + "Close a session (saving its database unless `save` is false).", + SESSIONS, + {"session": SESSION, "save": {"type": "boolean"}}, + ["session"], + run=_close, + destructive=False, + ), + Tool( + "save", + "Save the session's database so edits survive the session.", + EDIT, + op="save", + timeout=1800, + ), + # -- queries ------------------------------------------------------------ + Tool( + "info", + "Binary facts, backend, counts, database and export folder.", + QUERY, + op="info", + ), + Tool( + "functions", + "List functions (address, name, size, signature, kind, caller/callee counts).", + QUERY, + { + "filter": FILTER, + "kind": { + "type": "string", + "enum": ["all", "app", "library", "import", "thunk", "runtime"], + }, + "offset": OFFSET, + "limit": LIMIT, + }, + op="functions", + ), + Tool( + "function", + "One function: prototype, params, locals, callers, callees, imports it " + "calls, and where it is in the export (file:lines).", + QUERY, + {"target": TARGET}, + ["target"], + op="function", + ), + Tool( + "decompile", + "Decompile a function now, from the live database (reflects edits).", + QUERY, + {"target": TARGET}, + ["target"], + op="decompile", + timeout=180, + text_field="code", + ), + Tool( + "disassemble", + "Disassemble a function from the live database.", + QUERY, + {"target": TARGET, "max_lines": {"type": "integer", "minimum": 1}}, + ["target"], + op="disassemble", + text_field="disassembly", + ), + Tool( + "xrefs", + "Cross-references: `to` (callers and data reads/writes of the target) or " + "`from` (what a function calls).", + QUERY, + {"target": TARGET, "direction": {"type": "string", "enum": ["to", "from"]}}, + ["target"], + op="xrefs", + ), + Tool( + "strings", + "Strings with the functions that use them.", + QUERY, + {"filter": FILTER, "offset": OFFSET, "limit": LIMIT}, + op="strings", + ), + Tool( + "imports", + "Imported symbols and their libraries.", + QUERY, + {"filter": FILTER, "offset": OFFSET, "limit": LIMIT}, + op="imports", + ), + Tool( + "exports", + "Exported symbols.", + QUERY, + {"filter": FILTER, "offset": OFFSET, "limit": LIMIT}, + op="exports", + ), + Tool( + "sections", + "Sections/segments with addresses, sizes and permissions.", + QUERY, + op="sections", + ), + Tool( + "types", + "Types in the database (structs, unions, enums, typedefs).", + QUERY, + { + "filter": FILTER, + "kind": { + "type": "string", + "enum": ["struct", "union", "enum", "typedef", "func"], + }, + "offset": OFFSET, + "limit": LIMIT, + }, + op="types", + ), + Tool( + "type", + "One type's C declaration and members.", + QUERY, + {"name": {"type": "string"}}, + ["name"], + op="type", + text_field="c_decl", + ), + Tool( + "read_bytes", + "Read bytes at an address (hex and ASCII, at most 64 KiB).", + QUERY, + { + "address": TARGET, + "size": {"type": "integer", "minimum": 1, "maximum": 65536}, + }, + ["address"], + op="read_bytes", + ), + Tool( + "search", + "Search the binary's sections for text, hex bytes (`de ad be ef`), or a " + "byte regex; returns addresses and the functions containing them.", + QUERY, + { + "pattern": {"type": "string"}, + "kind": {"type": "string", "enum": ["text", "hex", "regex"]}, + "limit": {"type": "integer", "minimum": 1, "maximum": 1000}, + }, + ["pattern"], + op="search", + ), + Tool( + "lookup", + "Find functions and symbols whose name contains `name`.", + QUERY, + {"name": {"type": "string"}, "limit": {"type": "integer", "minimum": 1}}, + ["name"], + op="lookup", + ), + Tool( + "status", + "What the next `sync` will update (functions changed in this session).", + QUERY, + op="status", + ), + # -- edits ------------------------------------------------------------------ + Tool( + "rename", + "Rename a function or a global. Callers pick up the new name on the next `sync`.", + EDIT, + { + "target": TARGET, + "name": {"type": "string", "description": "New name (a C identifier)."}, + }, + ["target", "name"], + op="rename", + ), + Tool( + "rename_many", + "Rename many functions, types, members, or globals in one call " + "(deobfuscation passes). Each entry is {target, name}; failures are " + "reported per entry and do not stop the rest.", + EDIT, + { + "renames": { + "type": "array", + "items": { + "type": "object", + "properties": {"target": TARGET, "name": {"type": "string"}}, + "required": ["target", "name"], + }, + } + }, + ["renames"], + op="rename_many", + timeout=600, + ), + Tool( + "set_type", + "Set a function's prototype (`int parse(char *buf, size_t len)`) or a " + "global's type (`struct config *g_config`).", + EDIT, + {"target": TARGET, "c_decl": {"type": "string"}}, + ["target", "c_decl"], + op="set_type", + ), + Tool( + "set_comment", + "Comment an instruction (`line`) or a whole function (`function`).", + EDIT, + { + "address": TARGET, + "text": {"type": "string"}, + "kind": {"type": "string", "enum": ["line", "function"]}, + }, + ["address", "text"], + op="set_comment", + ), + Tool( + "declare_types", + "Add or replace C types (structs, unions, enums, typedefs) in the " + "database, e.g. `struct packet { uint16_t len; uint8_t type; uint8_t payload[]; };`.", + EDIT, + {"c_source": {"type": "string"}}, + ["c_source"], + op="declare_types", + ), + Tool( + "rename_local", + "Rename a local variable or parameter of a function (as the decompiler names it).", + EDIT, + {"function": TARGET, "old": {"type": "string"}, "new": {"type": "string"}}, + ["function", "old", "new"], + op="rename_local", + ), + Tool( + "set_local_type", + "Set the type of a function's local variable or parameter.", + EDIT, + { + "function": TARGET, + "name": {"type": "string"}, + "c_type": {"type": "string"}, + }, + ["function", "name", "c_type"], + op="set_local_type", + ), + Tool( + "define_function", + "Create a function at an address the analysis missed.", + EDIT, + {"address": TARGET}, + ["address"], + op="define_function", + timeout=600, + ), + Tool( + "undefine", + "Delete the function (or undefine the data) at an address.", + EDIT, + {"address": TARGET}, + ["address"], + op="undefine", + destructive=True, + ), + # -- project --------------------------------------------------------------- + Tool( + "sync", + "Update the export on disk from the database: new functions are added " + "next to their callers, renamed/retyped functions and their callers and " + "every function using a changed type are re-rendered, and only files " + "whose content changed are rewritten (plus the JSON and headers). " + "`mode`=`full` re-renders everything (after changes made outside this " + "server, e.g. in the IDA UI); `addresses` forces some functions. Saves " + "the database when it changed.", + PROJECT, + { + "mode": {"type": "string", "enum": ["auto", "full"]}, + "addresses": {"type": "array", "items": {"type": "string"}}, + }, + op="sync", + timeout=0, # no limit: a big binary can take a while + ), + Tool( + "export", + "Write a ToCode export for a session opened on a plain binary (default " + "folder: _decompiler next to it).", + PROJECT, + {"out_dir": {"type": "string"}}, + op="export", + timeout=0, + ), + # -- eval ------------------------------------------------------------------- + Tool( + "eval", + "Run a script inside the database. Python by default: the last " + "expression (or `result`) is returned and stdout is captured. Names: " + "IDA `db` (ida-domain) plus idaapi/idc/ida_* modules; Binary Ninja " + "`bv`, `bn`; r2 `r2` (`r2.cmd(...)`); angr `project`; everywhere " + "`session`, `analysis`, `resolve(name)`, and `touch(addr, ...)` to mark " + "functions your script changed for the next `sync`. Names you define " + "persist between calls (`reset` clears them). `language`=`r2` runs an " + "r2 command.", + EVAL, + { + "code": {"type": "string"}, + "language": {"type": "string", "enum": ["python", "r2"]}, + "reset": {"type": "boolean"}, + "timeout": { + "type": "number", + "minimum": 1, + "description": "Seconds (default 300).", + }, + }, + ["code"], + op="eval", + destructive=True, + timeout=300, + ), + ] + + +class ToolTable: + def __init__( + self, + manager: SessionManager, + *, + allow_edits: bool = True, + allow_eval: bool = True, + extra: list[Tool] | None = None, + ) -> None: + self.manager = manager + tools = build_tools() + list(extra or []) + # Read-only hides everything that changes the database; sync and + # export stay (they only write the export folder from the database). + self.tools = { + tool.name: tool + for tool in tools + if not (tool.group == EDIT and not allow_edits) + and not (tool.group == EVAL and not (allow_eval and allow_edits)) + } + + def list(self) -> list[dict[str, Any]]: + return [tool.schema() for tool in self.tools.values()] + + def call(self, call: ToolCall) -> dict[str, Any]: + tool = self.tools.get(call.name) + if tool is None: + return _error(f"unknown tool {call.name}") + args = dict(call.arguments) + try: + if tool.run is not None: + result = tool.run(self.manager, args, call) + else: + session = str(args.pop("session", "") or "") + timeout = args.pop("timeout", None) if tool.group == EVAL else None + result = self.manager.call( + session, + str(tool.op), + args, + timeout=float(timeout) if timeout else tool.timeout, + ) + except ToCodeError as exc: + return _error(str(exc)) + return _result(result, tool.text_field) + + +def _result(result: Any, text_field: str | None) -> dict[str, Any]: + structured = result if isinstance(result, dict) else {"result": result} + if ( + text_field + and isinstance(result, dict) + and isinstance(result.get(text_field), str) + ): + header = {k: v for k, v in result.items() if k != text_field} + text = json.dumps(header, ensure_ascii=False) + "\n\n" + result[text_field] + else: + text = json.dumps(structured, ensure_ascii=False, indent=1) + return { + "content": [{"type": "text", "text": text}], + "structuredContent": structured, + } + + +def _error(message: str) -> dict[str, Any]: + return {"content": [{"type": "text", "text": message}], "isError": True} + + +_GROUP_TITLES = { + SESSIONS: "Sessions", + QUERY: "Query", + EDIT: "Edit", + PROJECT: "Project", + EVAL: "Script", +} + + +def render_reference() -> str: + """`skill/references/tools.md`, generated from the tool table. + + Regenerate with `python -m tocode.live.tools > src/tocode/skill/references/tools.md`. + """ + lines = [ + "# ToCode MCP tools", + "", + "Every tool except `open` and `sessions` takes `session`: the id `open` returned, a folder", + "binary's relative path (`bin/app`, or `:bin/app`), or nothing when one session is open.", + "Targets and addresses accept `0x...` addresses or names (auto names like `sub_401000` too).", + "List tools page with `offset`/`limit` and return `total` and `next_offset`.", + "", + ] + tools = build_tools() + for group in (SESSIONS, QUERY, EDIT, PROJECT, EVAL): + lines += [f"## {_GROUP_TITLES[group]}", ""] + for tool in (item for item in tools if item.group == group): + lines += [f"### `{tool.name}`", "", tool.description] + arguments = [ + (name, spec) + for name, spec in tool.properties.items() + if name != "session" + ] + if arguments: + lines.append("") + for name, spec in arguments: + kind = spec.get("type", "") + if "enum" in spec: + kind = " | ".join(f"`{value}`" for value in spec["enum"]) + required = " (required)" if name in tool.required else "" + line = f"- `{name}`{required}: {kind}. {spec.get('description', '')}" + lines.append(line.rstrip(" .")) + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +if __name__ == "__main__": + print(render_reference(), end="") diff --git a/src/tocode/live/worker.py b/src/tocode/live/worker.py new file mode 100644 index 0000000..1cd27e7 --- /dev/null +++ b/src/tocode/live/worker.py @@ -0,0 +1,875 @@ +"""The per-session worker process and the operations it serves. + +A worker opens one binary or database with the regular backend stack +(``create_analyzer``), collects the inventory once, and then answers +operations sent by the supervisor over a pipe. It runs in its own process +group so a crash, an OOM kill, or a hung decompile only costs this session. +""" + +from __future__ import annotations + +import ast +from bisect import bisect_right +from contextlib import redirect_stderr, redirect_stdout +from dataclasses import asdict +import importlib +import io +import json +import os +from pathlib import Path +import re +import traceback +from typing import Any, Callable + +from ..errors import ToCodeError +from ..progress import Progress +from ..schema import ProgramAnalysis, Routine +from .targets import Target + +_AUTO_NAME_RX = re.compile( + r"(?:sub|fun|FUN|loc|j_sub|nullsub|qword|dword|word|byte|xmmword|ymmword|" + r"off|unk|stru|asc|flt|dbl|data|DAT|PTR|LAB)_([0-9a-fA-F]+)" +) +# Operations that never need a fresh inventory first. +_NO_REFRESH = frozenset({"save", "status", "eval", "info", "refresh", "sync"}) +MAX_TEXT = 200_000 +MAX_READ = 64 * 1024 +MAX_SEARCH_BYTES = 512 * 1024 * 1024 +DEFAULT_LIMIT = 100 + + +class OpError(ToCodeError): + """A request the session cannot satisfy (bad target, unsupported, ...).""" + + +# -------------------------------------------------------------------------- +# Process entry point +# -------------------------------------------------------------------------- + + +def worker_main( + target: Target, + conn: Any, + log_path: Path, + options: dict[str, Any], +) -> None: + from ..apk_native import _redirect_output, watch_parent + + if hasattr(os, "setsid"): + try: + os.setsid() + except OSError: # pragma: no cover - already a session leader + pass + watch_parent() + log_path.parent.mkdir(parents=True, exist_ok=True) + _redirect_output(log_path) + try: + live = open_session(target, log_path=log_path, options=options) + except BaseException as exc: # noqa: BLE001 - report every open failure + conn.send(("failed", f"{type(exc).__name__}: {exc}")) + return + conn.send(("ready", live.info())) + while True: + try: + message = conn.recv() + except (EOFError, OSError): + break + kind = message[0] + if kind == "close": + _, request_id, save = message + try: + result = live.close(save=bool(save)) + except BaseException as exc: # noqa: BLE001 + conn.send(("error", request_id, f"{type(exc).__name__}: {exc}")) + else: + conn.send(("ok", request_id, result)) + break + if kind != "call": + continue + _, request_id, op, args = message + try: + result = live.dispatch(op, args or {}) + except OpError as exc: + conn.send(("error", request_id, str(exc))) + except BaseException as exc: # noqa: BLE001 - keep serving after a failure + detail = traceback.format_exc(limit=6) + live.log(f"live: {op} failed: {detail}") + conn.send(("error", request_id, f"{type(exc).__name__}: {exc}")) + else: + conn.send(("ok", request_id, result)) + + +def open_session(target: Target, *, log_path: Path, options: dict[str, Any]) -> Any: + """The session class for the target's backend (.NET or a native backend).""" + if target.backend == "dotnet": + from .dotnet_live import DotnetLiveSession + + return DotnetLiveSession( + target, + log_path=log_path, + allow_eval=bool(options.get("allow_eval", True)), + allow_edits=bool(options.get("allow_edits", True)), + ) + return LiveSession(target, log_path=log_path, **options) + + +def run_python(namespace: dict[str, Any], code: str) -> tuple[Any, str, str | None]: + """Run an eval script: (last expression or `result`, output, traceback).""" + namespace["result"] = None + tree = ast.parse(code, mode="exec") + tail = None + last = tree.body[-1] if tree.body else None + if isinstance(last, ast.Expr): + tree.body.pop() + tail = ast.Expression(last.value) + output = io.StringIO() + error = None + value: Any = None + with redirect_stdout(output), redirect_stderr(output): + try: + exec(compile(tree, "", "exec"), namespace) # noqa: S102 + if tail is not None: + value = eval(compile(tail, "", "eval"), namespace) # noqa: S307 + else: + value = namespace.get("result") + except BaseException: # noqa: BLE001 - the script's own failure + error = traceback.format_exc(limit=8) + return value, output.getvalue(), error + + +# -------------------------------------------------------------------------- +# The session +# -------------------------------------------------------------------------- + + +class LiveSession: + def __init__( + self, + target: Target, + *, + log_path: Path, + binja_dir: str | None = None, + allow_eval: bool = True, + allow_edits: bool = True, + ) -> None: + from ..analysis import create_analyzer + from ..naming import build_name_book + + self.target = target + self.allow_eval = allow_eval + self.allow_edits = allow_edits + self.progress = Progress(enabled=False, log_path=log_path) + self.log(f"live: opening {target.open_path} ({target.backend})") + self.analyzer = create_analyzer( + target.open_path, + backend=target.backend, # type: ignore[arg-type] + progress=self.progress, + binja_dir=Path(binja_dir) if binja_dir else None, + ) + self.session: Any = self.analyzer.session + self.analysis: ProgramAnalysis = self.analyzer.collect() + self.names = build_name_book(self.analysis) + self.eval_namespace: dict[str, Any] | None = None + # Addresses changed in this session that the next sync must re-render. + self.dirty: set[int] = set() + self.unknown_changes = False # eval ran: sync must diff the inventory + self.stale = False # the inventory must be collected again before queries + self.types_changed = False + self.globals_changed = False + self.functions_changed = False + self.edit_ops: set[str] = set() + self.assigned_names: dict[str, int] = {} + self._build_indexes() + self.ops: dict[str, Callable[[dict[str, Any]], Any]] = { + "info": lambda _args: self.info(), + "functions": self.op_functions, + "function": self.op_function, + "decompile": self.op_decompile, + "disassemble": self.op_disassemble, + "xrefs": self.op_xrefs, + "strings": self.op_strings, + "imports": self.op_imports, + "exports": self.op_exports, + "sections": self.op_sections, + "types": self.op_types, + "type": self.op_type, + "read_bytes": self.op_read_bytes, + "search": self.op_search, + "lookup": self.op_lookup, + "eval": self.op_eval, + "save": lambda _args: self.save(), + "refresh": lambda _args: self.refresh(), + "status": lambda _args: self.status(), + } + self._register_extensions() + from . import hooks + + # Edits made by eval scripts mark their functions for the next sync. + self._hooks = hooks.install(self) + self.log("live: session ready") + + # -- plumbing ------------------------------------------------------------- + + def log(self, message: str) -> None: + self.progress.file_log(message) + + def dispatch(self, op: str, args: dict[str, Any]) -> Any: + handler = self.ops.get(op) + if handler is None: + raise OpError(f"unknown operation {op}") + if self.stale and op not in self.edit_ops and op not in _NO_REFRESH: + self.refresh() + return handler(args) + + def _register_extensions(self) -> None: + """Edit and sync operations live in their own modules.""" + for module in ("edits", "sync"): + try: + loaded = importlib.import_module(f"{__package__}.{module}") + except ImportError: + continue + register = getattr(loaded, "register", None) + if callable(register): + register(self) + + @property + def backend(self) -> str: + return str(self.session.backend_name) + + def info(self) -> dict[str, Any]: + binary = self.analysis.binary + return { + **self.target.describe(), + "backend": self.backend, + "backend_label": self.analyzer.backend_label, + "decompiler": self.analyzer.decompiler_label, + "arch": binary.arch, + "bits": binary.bits, + "format": binary.format_name, + "os": binary.os_name, + "image_base": hex(binary.image_base), + "entrypoints": [hex(item) for item in binary.entrypoints], + "functions": len(self.analysis.routines), + "imports": len(self.analysis.imports), + "exports": len(self.analysis.exports), + "strings": len(self.analysis.strings), + "types": len(self.analysis.types), + "eval": self.allow_eval, + "edits": self.allow_edits, + } + + def status(self) -> dict[str, Any]: + return { + "pending_functions": sorted(hex(item) for item in self.dirty), + "unknown_changes": self.unknown_changes, + } + + def refresh(self) -> dict[str, Any]: + """Collect the inventory again (after edits that add or rename things).""" + from ..naming import build_name_book + + self.analyzer.analysis = None + self.analysis = self.analyzer.collect() + self.names = build_name_book(self.analysis) + self._build_indexes() + self.stale = False + return {"functions": len(self.analysis.routines)} + + def reindex(self) -> None: + """Names changed in place: rebuild the name book and lookups.""" + from ..naming import build_name_book + + self.names = build_name_book(self.analysis) + self._build_indexes() + + def refresh_types(self) -> None: + self.analysis.types = self.analyzer._types() + + def close(self, *, save: bool) -> dict[str, Any]: + saved: dict[str, Any] = {"saved": False} + if save and (self.dirty or self.unknown_changes or self._edited): + saved = self.save() + try: + self.analyzer.close() + except Exception: # noqa: BLE001 + pass + return saved + + _edited = False + + def save(self) -> dict[str, Any]: + """Write the database to disk (edits survive the session).""" + from . import edits + + return edits.save_database(self) + + # -- indexes and lookup --------------------------------------------------- + + def _build_indexes(self) -> None: + routines = self.analysis.routines + self._starts = sorted(routines) + self._by_name: dict[str, int] = {} + for address, routine in routines.items(): + self._by_name.setdefault(routine.name, address) + self._by_name.setdefault( + self.names.function_name(address, routine.name), address + ) + for symbol in self.analysis.symbols: + if symbol.name: + self._by_name.setdefault(symbol.name, symbol.vaddr) + for export in self.analysis.exports: + self._by_name.setdefault(export.name, export.address) + for address, item in self.analysis.imports.items(): + self._by_name.setdefault(item.name, address) + # Names given in this session to addresses without a symbol yet. + for name, address in self.assigned_names.items(): + self._by_name[name] = address + self._lower = {name.lower(): address for name, address in self._by_name.items()} + self._locations = _load_locations(self.target.export_dir) + + def resolve(self, value: Any, *, what: str = "target") -> int: + if isinstance(value, int): + return value + if not isinstance(value, str) or not value.strip(): + raise OpError(f"{what} must be an address (0x...) or a name") + text = value.strip() + if re.fullmatch(r"0x[0-9a-fA-F]+", text): + return int(text, 16) + if text.isdigit(): + return int(text) + if text in self._by_name: + return self._by_name[text] + if text.lower() in self._lower: + return self._lower[text.lower()] + # Auto-generated names carry their address (IDA, Ghidra, Binary Ninja). + match = _AUTO_NAME_RX.fullmatch(text) + if match: + return int(match.group(1), 16) + raise OpError(f"no function or symbol named {text!r}; try `lookup`") + + def routine_at(self, address: int) -> Routine | None: + routines = self.analysis.routines + if address in routines: + return routines[address] + index = bisect_right(self._starts, address) - 1 + if index >= 0: + routine = routines[self._starts[index]] + if routine.address <= address < routine.address + max(routine.size, 1): + return routine + return None + + def function_for(self, value: Any) -> Routine: + address = self.resolve(value) + routine = self.routine_at(address) + if routine is None: + raise OpError(f"no function at {hex(address)}") + return routine + + def name_of(self, address: int) -> str: + routine = self.analysis.routines.get(address) + if routine is not None: + return self.names.function_name(address, routine.name) + item = self.analysis.imports.get(address) + return item.name if item is not None else hex(address) + + # -- queries -------------------------------------------------------------- + + def _routine_row(self, routine: Routine) -> dict[str, Any]: + return { + "address": hex(routine.address), + "name": self.names.function_name(routine.address, routine.name), + "size": routine.size, + "signature": routine.signature, + "kind": "import" + if routine.imported + else "library" + if routine.library + else "thunk" + if routine.thunk + else routine.code_kind, + "callers": routine.indegree, + "callees": routine.outdegree, + } + + def op_functions(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + kind = args.get("kind") or "all" + rows = [] + for address in self._starts: + routine = self.analysis.routines[address] + row = self._routine_row(routine) + if kind != "all" and row["kind"] != kind: + continue + if pattern is not None and not pattern(row["name"]): + continue + rows.append(row) + return _page(rows, args) + + def op_function(self, args: dict[str, Any]) -> dict[str, Any]: + routine = self.function_for(args.get("target")) + address = routine.address + callees = self.analysis.callees.get(address, []) + callers = self.analysis.callers.get(address, []) + row = self._routine_row(routine) + row.update( + { + "return_type": routine.return_type, + "calltype": routine.calltype, + "noreturn": routine.noreturn, + "params": [{"name": n, "type": t} for n, t in routine.params], + "locals": [{"name": n, "type": t} for n, t in routine.local_vars], + "callers": [ + {"address": hex(item), "name": self.name_of(item)} + for item in callers + ], + "callees": [ + {"address": hex(item), "name": self.name_of(item)} + for item in callees + ], + "imports_called": self.analysis.import_calls.get(address, []), + "source": {"file": routine.source_file, "line": routine.source_line} + if routine.source_file + else None, + "export": self._locations.get(address), + } + ) + return row + + def op_decompile(self, args: dict[str, Any]) -> dict[str, Any]: + from ..exporter import render_one_with_limits + + routine = self.function_for(args.get("target")) + rendered = render_one_with_limits( + self.session, self.analysis, routine, self.names + ) + return { + "address": hex(routine.address), + "name": rendered.c_name, + "prototype": rendered.prototype, + "code": _clip(rendered.c_text), + "error": rendered.failure.message if rendered.failure else None, + "export": self._locations.get(routine.address), + } + + def op_disassemble(self, args: dict[str, Any]) -> dict[str, Any]: + routine = self.function_for(args.get("target")) + text = self.session.disasm(routine.address) + limit = int(args.get("max_lines") or 0) + if limit > 0: + lines = text.splitlines() + text = "\n".join(lines[:limit]) + ( + f"\n... {len(lines) - limit} more lines" if len(lines) > limit else "" + ) + return { + "address": hex(routine.address), + "name": self.name_of(routine.address), + "disassembly": _clip(text), + } + + def op_xrefs(self, args: dict[str, Any]) -> dict[str, Any]: + address = self.resolve(args.get("target")) + direction = args.get("direction") or "to" + routine = self.analysis.routines.get(address) + rows: list[dict[str, Any]] = [] + if direction == "from": + if routine is None: + raise OpError("xrefs from works on a function; use eval for data") + for item in self.analysis.callees.get(address, []): + rows.append( + {"to": hex(item), "name": self.name_of(item), "kind": "call"} + ) + for name in self.analysis.import_calls.get(address, []): + rows.append({"to": None, "name": name, "kind": "import call"}) + return {"target": hex(address), "direction": "from", "refs": rows} + if routine is not None: + for item in self.analysis.callers.get(address, []): + rows.append( + {"from": hex(item), "function": self.name_of(item), "kind": "call"} + ) + names = self._import_names_at(address) + if names: + # An import is called through its stub: name the real callers. + for caller, called in sorted(self.analysis.import_calls.items()): + if names & {_bare(name) for name in called}: + rows.append( + { + "from": hex(caller), + "function": self.name_of(caller), + "kind": "import call", + } + ) + seen = {row["from"] for row in rows} + refs = self.analysis.data_xrefs.get(address) + if refs is None: + try: + refs = self.session.data_xrefs([address]).get(address, []) + except Exception: # noqa: BLE001 + refs = [] + for source, write in refs or []: + owner = self.routine_at(source) + key = hex(source) + if key in seen: + continue + rows.append( + { + "from": key, + "function": self.name_of(owner.address) if owner else None, + "kind": "write" if write else "read", + } + ) + return {"target": hex(address), "direction": "to", "refs": rows} + + def _import_names_at(self, address: int) -> set[str]: + """Import names behind an address: the import itself or its stub.""" + known = {_bare(item.name) for item in self.analysis.imports.values()} + names: set[str] = set() + imported = self.analysis.imports.get(address) + if imported is not None: + names.add(_bare(imported.name)) + stub = self.analysis.routines.get(address) + if stub is not None and _bare(stub.name) in known: + names.add(_bare(stub.name)) + return names + + def op_strings(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + rows = [] + for item in self.analysis.strings: + if pattern is not None and not pattern(item.value): + continue + users = sorted( + { + self.name_of(owner.address) + for source, _ in self.analysis.data_xrefs.get(item.vaddr, []) + if (owner := self.routine_at(source)) is not None + } + ) + rows.append( + { + "address": hex(item.vaddr), + "value": item.value, + "kind": item.kind, + "used_by": users, + } + ) + return _page(rows, args) + + def op_imports(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + rows = [ + {"address": hex(address), "name": item.name, "library": item.dll} + for address, item in sorted(self.analysis.imports.items()) + if pattern is None or pattern(item.name) + ] + return _page(rows, args) + + def op_exports(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + rows = [ + {"address": hex(item.address), "name": item.name, "kind": item.kind} + for item in self.analysis.exports + if pattern is None or pattern(item.name) + ] + return _page(rows, args) + + def op_sections(self, _args: dict[str, Any]) -> dict[str, Any]: + return { + "sections": [ + { + "name": segment.name, + "address": hex(segment.vaddr), + "size": max(segment.vsize, segment.size), + "perms": segment.perms, + "kind": segment.kind, + } + for segment in self.analysis.segments + ] + } + + def op_types(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = _matcher(args.get("filter")) + kind = args.get("kind") + rows = [ + {"name": item.name, "kind": item.kind, "size": item.size} + for item in self.analysis.types + if (pattern is None or pattern(item.name)) + and (not kind or item.kind == kind) + ] + return _page(rows, args) + + def op_type(self, args: dict[str, Any]) -> dict[str, Any]: + name = str(args.get("name") or "").strip() + for item in self.analysis.types: + if item.name == name: + return { + "name": item.name, + "kind": item.kind, + "size": item.size, + "c_decl": item.c_decl, + "members": item.members, + } + raise OpError(f"no type named {name!r}; try `types` with a filter") + + def read(self, address: int, size: int) -> bytes: + reader = getattr(self.session, "read_bytes", None) + if callable(reader): + data = reader(address, size) + if data: + return bytes(data) + segment = self.analysis.segment_at(address) + if segment is None or segment.size <= 0: + return b"" + offset = address - segment.vaddr + if offset >= segment.size: + return b"\0" * size # uninitialized (bss-like) memory + path = self.analysis.binary.path + try: + with path.open("rb") as handle: + handle.seek(segment.paddr + offset) + return handle.read(min(size, segment.size - offset)) + except OSError: + return b"" + + def op_read_bytes(self, args: dict[str, Any]) -> dict[str, Any]: + address = self.resolve(args.get("address"), what="address") + size = max(1, min(int(args.get("size") or 64), MAX_READ)) + data = self.read(address, size) + return { + "address": hex(address), + "size": len(data), + "hex": data.hex(), + "ascii": "".join(chr(b) if 32 <= b < 127 else "." for b in data), + } + + def op_search(self, args: dict[str, Any]) -> dict[str, Any]: + pattern = str(args.get("pattern") or "") + kind = args.get("kind") or "text" + limit = max(1, min(int(args.get("limit") or 50), 1000)) + if not pattern: + raise OpError("search needs a pattern") + if kind == "hex": + try: + needle = re.compile(re.escape(bytes.fromhex(pattern.replace(" ", "")))) + except ValueError as exc: + raise OpError(f"bad hex pattern: {exc}") from exc + elif kind == "regex": + needle = re.compile(pattern.encode("utf-8")) + else: + needle = re.compile(re.escape(pattern.encode("utf-8"))) + hits: list[dict[str, Any]] = [] + budget = MAX_SEARCH_BYTES + for segment in self.analysis.segments: + size = min(segment.size, budget) + if size <= 0: + continue + budget -= size + blob = self.read(segment.vaddr, size) + for match in needle.finditer(blob): + address = segment.vaddr + match.start() + owner = self.routine_at(address) + hits.append( + { + "address": hex(address), + "section": segment.name, + "function": self.name_of(owner.address) if owner else None, + } + ) + if len(hits) >= limit: + return { + "pattern": pattern, + "kind": kind, + "hits": hits, + "truncated": True, + } + return {"pattern": pattern, "kind": kind, "hits": hits, "truncated": False} + + def op_lookup(self, args: dict[str, Any]) -> dict[str, Any]: + query = str(args.get("name") or "").strip().lower() + if not query: + raise OpError("lookup needs a name") + limit = max(1, min(int(args.get("limit") or 50), 500)) + rows = [] + for name, address in sorted(self._by_name.items()): + if query in name.lower(): + owner = self.routine_at(address) + rows.append( + { + "name": name, + "address": hex(address), + "function": self.name_of(owner.address) if owner else None, + } + ) + if len(rows) >= limit: + break + return {"query": query, "matches": rows} + + # -- eval ------------------------------------------------------------------- + + def touch(self, *addresses: Any) -> None: + for value in addresses: + self.dirty.add(self.resolve(value)) + + def _namespace(self) -> dict[str, Any]: + if self.eval_namespace is not None: + return self.eval_namespace + namespace: dict[str, Any] = { + "__name__": "__tocode_eval__", + "session": self.session, + "analysis": self.analysis, + "touch": self.touch, + "resolve": self.resolve, + "result": None, + } + backend = self.backend + if backend == "ida": + namespace["db"] = getattr(self.session, "_db", None) + for module in ( + "idaapi", + "idc", + "idautils", + "ida_funcs", + "ida_name", + "ida_typeinf", + "ida_hexrays", + "ida_bytes", + "ida_xref", + "ida_segment", + "ida_nalt", + "ida_lines", + "ida_ua", + "ida_loader", + ): + try: + namespace[module] = importlib.import_module(module) + except Exception: # noqa: BLE001 - optional modules + continue + elif backend == "binja": + namespace["bv"] = getattr(self.session, "_bv", None) + namespace["bn"] = getattr(self.session, "_bn", None) + elif backend == "r2": + namespace["r2"] = self.session + elif backend == "angr": + namespace["project"] = getattr(self.session, "project", None) + self.eval_namespace = namespace + return namespace + + def op_eval(self, args: dict[str, Any]) -> dict[str, Any]: + if not self.allow_eval: + raise OpError("eval is disabled for this server (--no-eval / --read-only)") + code = str(args.get("code") or "") + language = args.get("language") or "python" + if args.get("reset"): + self.eval_namespace = None + self.unknown_changes = True + self._edited = True + if language == "r2": + command = getattr(self.session, "cmd", None) + if not callable(command): + raise OpError( + f"r2 commands need the r2 backend (this is {self.backend})" + ) + return {"output": _clip(str(command(code)))} + if language != "python": + raise OpError("language must be python or r2") + before = set(self.dirty) + value, output, error = run_python(self._namespace(), code) + return { + "result": _jsonable(value), + "stdout": _clip(output), + "error": error, + "touched": sorted(hex(item) for item in self.dirty - before), + } + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _bare(name: str) -> str: + """`_printf`, `.printf`, `printf@@GLIBC_2.17` -> `printf`.""" + return name.split("@", 1)[0].lstrip("._") + + +def _load_locations(export_dir: Path | None) -> dict[int, dict[str, Any]]: + """Where each function lives in the export (from function-index.json).""" + if export_dir is None: + return {} + try: + payload = json.loads((export_dir / "function-index.json").read_text("utf-8")) + except (OSError, ValueError): + return {} + locations: dict[int, dict[str, Any]] = {} + for row in payload.get("functions", []): + try: + address = int(str(row["address"]), 16) + except (KeyError, ValueError): + continue + c = row.get("c") or {} + asm = row.get("asm") or {} + locations[address] = { + "c": f"{_relative(c.get('path'), export_dir)}:{c.get('line_start')}-{c.get('line_end')}", + "asm": f"{_relative(asm.get('path'), export_dir)}:{asm.get('line_start')}-{asm.get('line_end')}", + } + return locations + + +def _relative(path: Any, root: Path) -> str: + if not path: + return "" + try: + return Path(str(path)).resolve().relative_to(root.resolve()).as_posix() + except ValueError: + return str(path) + + +def _matcher(pattern: Any) -> Callable[[str], bool] | None: + if not pattern: + return None + text = str(pattern) + try: + compiled = re.compile(text, re.IGNORECASE) + except re.error: + lowered = text.lower() + return lambda value: lowered in value.lower() + return lambda value: compiled.search(value) is not None + + +def _page(rows: list[dict[str, Any]], args: dict[str, Any]) -> dict[str, Any]: + offset = max(0, int(args.get("offset") or 0)) + limit = max(1, min(int(args.get("limit") or DEFAULT_LIMIT), 2000)) + page = rows[offset : offset + limit] + result: dict[str, Any] = {"total": len(rows), "offset": offset, "items": page} + if offset + limit < len(rows): + result["next_offset"] = offset + limit + return result + + +def _clip(text: str, limit: int = MAX_TEXT) -> str: + if len(text) <= limit: + return text + return text[:limit] + f"\n... truncated ({len(text) - limit} more characters)" + + +def _jsonable(value: Any, depth: int = 0) -> Any: + if depth > 6: + return _clip(repr(value), 2000) + if value is None or isinstance(value, (bool, int, float, str)): + return _clip(value) if isinstance(value, str) else value + if isinstance(value, (bytes, bytearray)): + return bytes(value[:MAX_READ]).hex() + if isinstance(value, Path): + return str(value) + if isinstance(value, dict): + return { + str(key): _jsonable(item, depth + 1) + for key, item in list(value.items())[:2000] + } + if isinstance(value, (list, tuple, set, frozenset)): + return [_jsonable(item, depth + 1) for item in list(value)[:2000]] + if hasattr(value, "__dataclass_fields__"): + try: + return _jsonable(asdict(value), depth + 1) + except TypeError: + pass + return _clip(repr(value), 2000) diff --git a/src/tocode/skill/SKILL.md b/src/tocode/skill/SKILL.md new file mode 100644 index 0000000..4cb27db --- /dev/null +++ b/src/tocode/skill/SKILL.md @@ -0,0 +1,155 @@ +--- +name: tocode +description: Reverse engineer compiled programs with ToCode. It exports a binary, IDA/Binary Ninja database, APK, .NET assembly, or whole app folder into a source-like C project, and its MCP server (`tocode mcp`) keeps IDA, Binary Ninja, radare2, or angr open so you can decompile, cross-reference, rename, retype, declare structs and enums, run scripts in the database, and sync every improvement back into the exported source. Use for reverse engineering, binary or malware or firmware analysis, CTF reversing, recovering source code from stripped binaries, or whenever a folder holds a ToCode export (an AGENTS.md that mentions ToCode, functions.json, src/raw/). +--- + +# ToCode + +ToCode turns a binary into a project you can read like source code, and keeps +the decompiler session alive so you can improve it. You read files; the +database is the source of truth; `sync` regenerates the files from it. + +## 1. Get an export + +```bash +tocode ./program # -> program_decompiler/ (default backend: IDA, then Binary Ninja, r2, angr) +tocode ./program.i64 # an existing IDA or Binary Ninja (.bndb) database +tocode ./App/ # a whole folder: every binary, with the links between them +``` + +An export folder contains: + +| Path | What it is | +| --- | --- | +| `AGENTS.md` | Read first: layout and conventions for this binary | +| `src/raw//*.c` | Decompiled C, grouped by call graph (or by original source file with debug info) | +| `src/raw//*.asm` / `*.summary` | Matching assembly and per-function summaries | +| `functions.json` | Every function: address, name, prototype, params, locals, callers, callees, file and line range | +| `function-index.json` | Address -> file:lines, for jumping straight to a function | +| `triage.json` | Entry points, interesting functions, suspicious imports and strings: start here | +| `strings.json`, `imports.json`, `exports.json`, `types.json` | Strings with their users, imports, exports, recovered types | +| `include/*.h`, `include/*.types.h` | Prototypes and type definitions | +| `.i64` / `.bndb` | The database the live tools edit | + +Read with your normal file tools (grep, open at a line range). Use the MCP +tools when you need something the files cannot tell you, or to change the +analysis. + +## 2. Connect the live tools + +If the `tocode` MCP tools are not available, register the server once: + +```bash +claude mcp add --scope user tocode -- tocode mcp # Claude Code +tocode mcp --install codex # Codex (~/.codex/config.toml) +tocode mcp --config # snippets for other clients +``` + +Then `open` the export folder (or a folder export, a binary, a database). It +returns a `session`; pass it to every other tool. With one session open you +may omit it. For a folder export, address a binary as `session="bin/app"`. + +`open` uses the export's own database. Pass `database` to work on another +`.i64`/`.bndb` instead. + +## 3. The loop + +1. **Orient.** Read `AGENTS.md`, `triage.json`, and the C for entry points and + `main`. `strings`, `imports`, and `xrefs` tell you what code does what. +2. **Query live.** `decompile`, `disassemble`, `function`, `xrefs`, + `read_bytes`, `search`, `types`, `lookup` answer from the database now, + including your edits. +3. **Improve the database.** `rename`, `set_type`, `declare_types`, + `rename_local`, `set_local_type`, `set_comment`, `define_function`. + Each one names what it makes stale (`pending_sync`). +4. **Sync.** Call `sync` after a batch of edits. It re-renders what changed + (and callers of renamed or retyped functions, and every function using a + changed type or global), rewrites only files whose content changed, updates + the JSON and headers, and saves the database. The report lists what moved. +5. **Re-read** the changed files and continue. Better names and types make + the decompiler output better everywhere; iterate. + +Sync after every few edits rather than once at the end: you then read code +that already reflects what you learned. + +## 4. Tools + +| Group | Tools | +| --- | --- | +| Sessions | `open`, `sessions`, `close`, `save` | +| Query | `info`, `functions`, `function`, `decompile`, `disassemble`, `xrefs`, `strings`, `imports`, `exports`, `sections`, `types`, `type`, `read_bytes`, `search`, `lookup`, `status` | +| Edit | `rename`, `set_type`, `set_comment`, `declare_types`, `rename_local`, `set_local_type`, `define_function`, `undefine` | +| Script | `eval` (Python in the database; `language="r2"` for r2 commands) | +| Project | `sync`, `export` | + +Targets accept an address (`"0x401000"`) or a name; auto names such as +`sub_401000`, `qword_20020`, `DAT_00403000` work too. Full arguments: +[references/tools.md](references/tools.md). + +## 5. Recovering source code + +To reconstruct the original program (names, types, structs) work bottom-up: + +1. **Map the program.** List app functions (`functions` with `kind="app"`), + find `main` and entry points, and note the strings and imports each + function uses. Name obvious functions first: a function that only calls + `fopen`/`fread` on a path from argv is a file loader. +2. **Name leaves before callers.** A leaf's role is clearest from the + library calls and constants it uses; once named, callers read better. +3. **Recover structs from field accesses.** Repeated `*(ptr + 0x18)` patterns + on the same pointer are a struct. Collect offsets and sizes from every user + (`xrefs` to the allocator, `decompile` its callers), write the struct with + `declare_types` (pad unknown gaps: `char _pad_0x10[8];`), then `set_type` + the functions and globals that hold it and `set_local_type` the locals. + The decompiler then prints `pkt->payload_len` instead of offsets. +4. **Enums from constants** compared or switched on; `declare_types` an + `enum` and set it on the parameter or field. +5. **Globals**: `rename` and `set_type` them (`struct session *g_sessions`). +6. **Prototypes**: `set_type` with real parameter names and types; callers + then show meaningful arguments. +7. `sync`, re-read, repeat until the code reads like source. +8. If asked for source files, write them from the synced export, adding + includes and reordering into a buildable program. + +Details and a worked example: [references/workflows.md](references/workflows.md). + +## 6. .NET + +`tocode App.dll` (also apphosts, single-file bundles, `.nupkg`) exports C# +and IL, one `.cs`/`.il` file per type. The same tools work after `open`: + +- Targets: `Assembly!0x06000012` (a metadata token: 02 type, 04 field, + 06 method, 17 property, 14 event), `Namespace.Type`, or + `Namespace.Type::Member`. An ambiguous name lists the candidates. +- Edits: `rename` (types, methods, fields, properties, events), + `rename_many`, `rename_local` (parameters, and locals as the decompiled C# + names them), `set_comment`. Types come from metadata, so `set_type`, + `declare_types`, `set_local_type`, `define_function`, and `undefine` do not + apply. +- The database is an edit log in the export; ToCode also writes renamed + assemblies to `patched/` (open them in dnSpy/ILSpy, or run them). +- `sync` rewrites the types the edits touched; a renamed type's file moves. +- Obfuscated code: start at the entry point, `strings`, and `imports`; name + methods from the APIs they call, work outward, rename in batches with + `rename_many`, `sync`, and re-read. `eval` has the dnlib modules (`modules`) + and `decompile(target)`. + +## 7. Rules of thumb + +- Prefer the dedicated edit tools over `eval`: they report exactly what to + re-render. After changing things in `eval`, call `touch(addr, ...)` for the + functions you changed (or `sync` with `mode="full"`). +- `eval` runs inside the analysis process. Keep scripts bounded; a call that + runs past its timeout stops the session's worker, and unsaved edits since + the last `save`/`sync` are lost. Recipes per backend: + [references/eval.md](references/eval.md). +- Edits that the user made in the IDA or Binary Ninja UI are not tracked: + run `sync` with `mode="full"`, or `tocode sync --full` from a shell. +- Never edit files under `src/raw/` by hand to record findings; put them in + the database (names, types, comments) and `sync`. Hand edits are + overwritten when that file is regenerated. +- Backends: IDA and Binary Ninja support every edit; radare2 a subset; angr is + query-only. Some Binary Ninja builds cannot save databases headless; `save` + says so and edits then last for the session. +- A tool error explains what to do (wrong name, unsupported, timed out). Read + it before retrying. diff --git a/src/tocode/skill/references/eval.md b/src/tocode/skill/references/eval.md new file mode 100644 index 0000000..a6d3fb1 --- /dev/null +++ b/src/tocode/skill/references/eval.md @@ -0,0 +1,193 @@ +# `eval` recipes + +`eval` runs Python inside the session's analysis process. The last expression +(or a variable named `result`) is returned as JSON; `print` output comes back +in `stdout`. Names you define persist between calls (`reset=true` clears them). + +Always available: `session`, `analysis` (ToCode's inventory), `resolve(name)` +(name or `0x...` to an address), and `touch(addr, ...)` (mark functions your +script changed so the next `sync` re-renders them). + +Renames and prototype changes made in a script are found by `sync` on its +own (it diffs the database against the export). Changes inside a function +body (local variable names, comments) are not: `touch` those functions, or +run `sync` with `mode="full"`. + +Keep scripts bounded. A call that outlives its `timeout` (default 300 s) stops +the session's worker; unsaved edits since the last `save`/`sync` are lost. + +## IDA (`db` is the ida-domain database; `idaapi`, `idc`, `idautils`, `ida_*` are imported) + +Functions that call a function (for imports, use the `xrefs` tool: it follows +the PLT/thunk stub to the real callers): + +```python +import idautils, ida_funcs +target = resolve("parse_config") +sorted({hex(ida_funcs.get_func(x.frm).start_ea) for x in idautils.XrefsTo(target) + if ida_funcs.get_func(x.frm)}) +``` + +Find an immediate constant (crypto constants, magic values): + +```python +import idautils, idc +hits = [] +for f in idautils.Functions(): + for ea in idautils.FuncItems(f): + for n in range(3): + if idc.get_operand_type(ea, n) == idc.o_imm and idc.get_operand_value(ea, n) == 0x67452301: + hits.append((hex(f), hex(ea))) +hits +``` + +Pseudocode of a function as text, and its local variables: + +```python +import ida_hexrays +cfunc = ida_hexrays.decompile(resolve("parse_config")) +[(v.name, str(v.type())) for v in cfunc.get_lvars()], str(cfunc)[:2000] +``` + +Bulk renames (each also reported by the next `sync`): + +```python +import idc +for ea, name in [(0x401200, "read_header"), (0x401340, "check_crc")]: + assert idc.set_name(ea, name, idc.SN_CHECK), name +"ok" +``` + +Rename a local variable through Hex-Rays (then `touch` the function): + +```python +import ida_hexrays +f = resolve("parse_config") +ida_hexrays.rename_lvar(f, "v3", "cursor") +touch(f) +``` + +Read memory and strings: + +```python +import ida_bytes, idc +ida_bytes.get_bytes(0x403000, 32).hex(), idc.get_strlit_contents(0x403020, -1, idc.STRTYPE_C) +``` + +Segments and entry points: + +```python +import idautils, idc +[(idc.get_segm_name(s), hex(s), hex(idc.get_segm_end(s))) for s in idautils.Segments()] +``` + +## Binary Ninja (`bv` is the BinaryView, `bn` the module) + +High-level IL of a function: + +```python +f = bv.get_function_at(resolve("main")) +[str(line) for line in f.hlil.instructions][:40] +``` + +Code references to an address, and the functions they are in: + +```python +[(hex(r.address), r.function.name) for r in bv.get_code_refs(resolve("strcpy"))] +``` + +Define a type and apply it to a global: + +```python +t = bv.parse_type_string("struct packet { uint16_t len; uint8_t kind; uint8_t data[]; }")[0] +bv.define_user_type("packet", t) +bv.define_user_data_var(0x404010, bv.parse_type_string("struct packet *")[0]) +``` + +Rename a variable (then `touch` the function): + +```python +f = bv.get_function_at(resolve("handle_packet")) +var = next(v for v in f.vars if v.name == "var_18") +f.create_user_var(var, var.type, "pkt") +touch(f.start) +``` + +Strings referenced from a function: + +```python +f = bv.get_function_at(resolve("main")) +[s.value for s in bv.get_strings() if any(r.function == f for r in bv.get_code_refs(s.start))] +``` + +## radare2 (`language="r2"` runs a command; in Python, `r2.cmd(...)` / `r2.cmdj(...)`) + +| Command | Purpose | +| --- | --- | +| `afl~crypt` | functions whose name contains crypt | +| `pdf @ main` | disassemble a function | +| `pdg @ main` | decompile (r2ghidra) | +| `axt @ sym.imp.strcpy` | references to an import | +| `/x deadbeef` | search bytes | +| `izz~http` | strings containing http | +| `afvn new old @ fcn` | rename a local variable | + +```python +[f["name"] for f in r2.cmdj("aflj") if f["size"] > 1000] +``` + +## .NET (`modules` are dnlib `ModuleDefMD`s, `module` the first; `decompile(target)` returns C#) + +Types and methods whose names look obfuscated (short or non-ASCII): + +```python +import re +odd = re.compile(r"^[^A-Za-z_]|^.{1,2}$") +[(t.FullName, t.MDToken.Raw) for t in module.GetTypes() if odd.search(t.Name.String)][:50] +``` + +Methods that call an API (here `System.Net.WebClient::DownloadString`): + +```python +hits = [] +for t in module.GetTypes(): + for m in t.Methods: + if m.HasBody and any( + i.Operand is not None and "WebClient::DownloadString" in str(i.Operand) + for i in m.Body.Instructions + ): + hits.append((m.FullName, hex(m.MDToken.Raw))) +hits +``` + +A string-decryption helper: a static method taking an int and returning a +string, called from many places: + +```python +from collections import Counter +calls = Counter() +for t in module.GetTypes(): + for m in t.Methods: + if m.HasBody: + for i in m.Body.Instructions: + if i.OpCode.Name == "call" and i.Operand is not None and str(i.Operand).startswith("System.String ") and "(System.Int32)" in str(i.Operand): + calls[str(i.Operand)] += 1 +calls.most_common(5) +``` + +The IL of one method as text: + +```python +m = module.ResolveMethod(resolve("Billing.Invoice::Total").token & 0xFFFFFF) +[f"{i.Offset:04x} {i.OpCode.Name} {i.Operand}" for i in m.Body.Instructions] +``` + +Prefer the `rename` / `rename_many` tools for names: they record the edit in +the log and report what to re-render. + +## angr (`project` is the angr Project; query-only) + +```python +cfg = project.analyses.CFGFast(normalize=True) +[hex(a) for a in cfg.kb.functions.callgraph.predecessors(resolve("encrypt"))] +``` diff --git a/src/tocode/skill/references/tools.md b/src/tocode/skill/references/tools.md new file mode 100644 index 0000000..4aa0a94 --- /dev/null +++ b/src/tocode/skill/references/tools.md @@ -0,0 +1,230 @@ +# ToCode MCP tools + +Every tool except `open` and `sessions` takes `session`: the id `open` returned, a folder +binary's relative path (`bin/app`, or `:bin/app`), or nothing when one session is open. +Targets and addresses accept `0x...` addresses or names (auto names like `sub_401000` too). +List tools page with `offset`/`limit` and return `total` and `next_offset`. + +## Sessions + +### `open` + +Open a binary or database and keep it alive for the other tools. `path` is an export folder (made by `tocode `), a folder export, a binary, or a database (.i64/.idb/.bndb). By default an export's own database is used; pass `database` to edit another one. Returns the session id (a folder returns its binaries; address them as `session`=). + +- `path` (required): string. Export folder, folder export, binary, or database +- `database`: string. Database to open instead of the export's own copy +- `backend`: `auto` | `ida` | `binja` | `r2` | `angr`. Backend for a plain binary (databases pick their own) + +### `sessions` + +List open sessions and folders. + +### `close` + +Close a session (saving its database unless `save` is false). + +- `save`: boolean + +## Query + +### `info` + +Binary facts, backend, counts, database and export folder. + +### `functions` + +List functions (address, name, size, signature, kind, caller/callee counts). + +- `filter`: string. Case-insensitive regular expression (or substring) to match +- `kind`: `all` | `app` | `library` | `import` | `thunk` | `runtime` +- `offset`: integer. First row to return +- `limit`: integer. Rows per page + +### `function` + +One function: prototype, params, locals, callers, callees, imports it calls, and where it is in the export (file:lines). + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` + +### `decompile` + +Decompile a function now, from the live database (reflects edits). + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` + +### `disassemble` + +Disassemble a function from the live database. + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `max_lines`: integer + +### `xrefs` + +Cross-references: `to` (callers and data reads/writes of the target) or `from` (what a function calls). + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `direction`: `to` | `from` + +### `strings` + +Strings with the functions that use them. + +- `filter`: string. Case-insensitive regular expression (or substring) to match +- `offset`: integer. First row to return +- `limit`: integer. Rows per page + +### `imports` + +Imported symbols and their libraries. + +- `filter`: string. Case-insensitive regular expression (or substring) to match +- `offset`: integer. First row to return +- `limit`: integer. Rows per page + +### `exports` + +Exported symbols. + +- `filter`: string. Case-insensitive regular expression (or substring) to match +- `offset`: integer. First row to return +- `limit`: integer. Rows per page + +### `sections` + +Sections/segments with addresses, sizes and permissions. + +### `types` + +Types in the database (structs, unions, enums, typedefs). + +- `filter`: string. Case-insensitive regular expression (or substring) to match +- `kind`: `struct` | `union` | `enum` | `typedef` | `func` +- `offset`: integer. First row to return +- `limit`: integer. Rows per page + +### `type` + +One type's C declaration and members. + +- `name` (required): string + +### `read_bytes` + +Read bytes at an address (hex and ASCII, at most 64 KiB). + +- `address` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `size`: integer + +### `search` + +Search the binary's sections for text, hex bytes (`de ad be ef`), or a byte regex; returns addresses and the functions containing them. + +- `pattern` (required): string +- `kind`: `text` | `hex` | `regex` +- `limit`: integer + +### `lookup` + +Find functions and symbols whose name contains `name`. + +- `name` (required): string +- `limit`: integer + +### `status` + +What the next `sync` will update (functions changed in this session). + +## Edit + +### `save` + +Save the session's database so edits survive the session. + +### `rename` + +Rename a function or a global. Callers pick up the new name on the next `sync`. + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `name` (required): string. New name (a C identifier) + +### `rename_many` + +Rename many functions, types, members, or globals in one call (deobfuscation passes). Each entry is {target, name}; failures are reported per entry and do not stop the rest. + +- `renames` (required): array + +### `set_type` + +Set a function's prototype (`int parse(char *buf, size_t len)`) or a global's type (`struct config *g_config`). + +- `target` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `c_decl` (required): string + +### `set_comment` + +Comment an instruction (`line`) or a whole function (`function`). + +- `address` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `text` (required): string +- `kind`: `line` | `function` + +### `declare_types` + +Add or replace C types (structs, unions, enums, typedefs) in the database, e.g. `struct packet { uint16_t len; uint8_t type; uint8_t payload[]; };`. + +- `c_source` (required): string + +### `rename_local` + +Rename a local variable or parameter of a function (as the decompiler names it). + +- `function` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `old` (required): string +- `new` (required): string + +### `set_local_type` + +Set the type of a function's local variable or parameter. + +- `function` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` +- `name` (required): string +- `c_type` (required): string + +### `define_function` + +Create a function at an address the analysis missed. + +- `address` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` + +### `undefine` + +Delete the function (or undefine the data) at an address. + +- `address` (required): string. Function or symbol: an address (`0x401000`) or a name. For .NET: `Assembly!0x06000012`, `Namespace.Type`, or `Namespace.Type::Member` + +## Project + +### `sync` + +Update the export on disk from the database: new functions are added next to their callers, renamed/retyped functions and their callers and every function using a changed type are re-rendered, and only files whose content changed are rewritten (plus the JSON and headers). `mode`=`full` re-renders everything (after changes made outside this server, e.g. in the IDA UI); `addresses` forces some functions. Saves the database when it changed. + +- `mode`: `auto` | `full` +- `addresses`: array + +### `export` + +Write a ToCode export for a session opened on a plain binary (default folder: _decompiler next to it). + +- `out_dir`: string + +## Script + +### `eval` + +Run a script inside the database. Python by default: the last expression (or `result`) is returned and stdout is captured. Names: IDA `db` (ida-domain) plus idaapi/idc/ida_* modules; Binary Ninja `bv`, `bn`; r2 `r2` (`r2.cmd(...)`); angr `project`; everywhere `session`, `analysis`, `resolve(name)`, and `touch(addr, ...)` to mark functions your script changed for the next `sync`. Names you define persist between calls (`reset` clears them). `language`=`r2` runs an r2 command. + +- `code` (required): string +- `language`: `python` | `r2` +- `reset`: boolean +- `timeout`: number. Seconds (default 300) diff --git a/src/tocode/skill/references/workflows.md b/src/tocode/skill/references/workflows.md new file mode 100644 index 0000000..e1c0ae2 --- /dev/null +++ b/src/tocode/skill/references/workflows.md @@ -0,0 +1,113 @@ +# Workflows + +## Recovering source from a stripped binary (worked example) + +The binary is a small network daemon, stripped. After `tocode ./netd` and +`open`: + +1. **Find main and the pieces.** `xrefs` to `recv` shows one app caller, + `sub_4012a0`; `strings` shows `"bad checksum"`, `"PING"`, `"peer %s closed"` + used near it. `function sub_4012a0` lists its callees: `sub_401100` + (called with the receive buffer and its length), `sub_4011c0`, + `sub_401230`. +2. **Name by behaviour.** `decompile sub_401100`: it reads a 16-bit length at + offset 0, a type byte at offset 2, checks a 32-bit sum at offset 4, and + returns a pointer into the buffer or NULL: `rename sub_401100 parse_packet`. + `sub_4011c0` switches on the type byte and calls a different function per + value: `dispatch_packet`. `sub_401230` looks up a peer by address in a + global array: `find_session`. `sub_4012a0` is the receive loop, + `serve_peer`. +3. **Recover the struct.** Offsets used across those functions: 0x0 (uint16 + length, compared against the received size), 0x2 (uint8 compared against + 1, 2, 3), 0x3 (uint8 flags, masked with 0x80), 0x4 (uint32 checksum), + 0x8 (start of the payload). The type values are message kinds: + + ```c + enum msg_type { MSG_PING = 1, MSG_DATA = 2, MSG_CLOSE = 3 }; + struct packet { uint16_t payload_len; uint8_t type; uint8_t flags; uint32_t checksum; uint8_t payload[]; }; + ``` + + `declare_types` with that source. +4. **Apply types.** `set_type parse_packet "struct packet *parse_packet(uint8_t *buf, size_t len)"`, + `set_type dispatch_packet "int dispatch_packet(struct session *peer, struct packet *pkt)"`. + In `parse_packet`, `set_local_type` the result variable to `struct packet *` + and `rename_local` it to `pkt`. +5. **Globals.** `rename qword_404060 g_sessions`, + `set_type g_sessions "struct session *g_sessions"`, `rename dword_404058 g_session_count`. +6. **Sync and read.** After `sync`, `dispatch_packet` reads + `switch (pkt->type) { case MSG_PING: ...` and the receive loop reads + `pkt = parse_packet(buf, n);`. Repeat for `struct session` and the + handlers until every function reads like source. + +## Deobfuscating a .NET assembly + +1. `info` and `sections` show the assemblies and obfuscation hints; + `functions` with `kind="app"` lists the methods. +2. Start from the entry point (`Program::Main`, or the type `exports` and + `triage.json` point at) and from `strings`: each string's `used_by` names + the methods to read first. +3. Name from behaviour: a method calling `File.ReadAllBytes` then + `Aes.Create` is a decryptor; one building `HttpClient` requests is a + network client. `xrefs` from a method shows the APIs and fields it uses. +4. Rename in batches: `rename_many` with the types, methods, fields, and + properties you understood (`{"target": "a.b::c", "name": "LoadConfig"}`), + then `sync` and re-read: callers now show the new names. +5. String encryption: find the helper with the `eval` recipe, decrypt the + values in `eval`, and record them as `set_comment` on each caller. +6. `patched/.dll` holds the renamed assembly; reopen it in other + tools or run it. + +## Struct recovery checklist + +- Gather every access to the pointer: the allocator's callers, functions that + receive it, and globals that store it (`xrefs`). +- Record `offset -> size, how it is used` (compared, summed, passed to + `strcpy`, dereferenced as a pointer, called as a function pointer). +- The allocation size (`malloc(0x28)`, `calloc(n, 0x28)`, `new` size) is the + struct size, or an array element size. +- Fill gaps with padding (`char _pad_0x10[8];`) rather than guessing; refine + later. Arrays show up as `base + index * size`. +- Pointers to the same struct type are list or tree links; a pointer called + through is a function pointer (`int (*handler)(struct ctx *)`). +- After `declare_types`, set the type on the functions (`set_type`), locals + (`set_local_type`), and globals (`set_type` on the global), then `sync`. + +## Naming conventions + +- Functions: verb_object (`parse_header`, `send_packet`), keep library-like + names when the behaviour matches a known API (`memcpy`-like copy loops, + `crc32`). +- Globals: `g_` prefix; constants and enums in UPPER_CASE. +- Keep `sub_...` for functions you have not understood yet: a wrong name + misleads more than no name. + +## Malware and firmware triage + +1. `triage.json` and `imports`: networking, crypto, process injection, + persistence (registry, services, cron), anti-debugging. +2. `strings` with filters (`http|\.exe|cmd|password|key`), then `xrefs` to + each interesting string to find the code that uses it. +3. Constants (`search` for byte patterns, or the `eval` constant finder) to + spot crypto (AES S-box, SHA/MD5 initial values, RC4 key scheduling). +4. Decode obfuscated strings by reading the decoder with `decompile`, then + reproducing it in `eval` over `read_bytes` output; write the decoded + values back as comments (`set_comment`) and `sync`. + +## Folder exports (several binaries) + +- `open` the folder export; each binary is a session named by its relative + path (`bin/app`, `lib/libcrypto.so.3`), opened when first used. +- `mappings.json` at the top maps every imported symbol to the file and line + range that defines it in another binary's export; calls in the C carry + `/* tocode-link: name -> lib/x.so/src/raw/... */` comments. +- Work on the library side first when the executable calls into it: names + and types recovered there make the executable's calls readable. +- `sync` on one binary also refreshes the folder's links and comments. + +## When the export and the database disagree + +- Edits made in the IDA or Binary Ninja UI: `sync` with `mode="full"` (or + `tocode sync --full`). +- A function missing from the export (code the analysis did not find): + `define_function` at its address, then `sync`; it is placed next to its + callers. diff --git a/src/tocode/syncstate.py b/src/tocode/syncstate.py new file mode 100644 index 0000000..a1648ac --- /dev/null +++ b/src/tocode/syncstate.py @@ -0,0 +1,115 @@ +"""What an export keeps so ``sync`` can update it incrementally. + +``.tocode/sync/manifest.json`` records how functions were placed into files +(clusters with their members, labels, folders) and the export's options; +``.tocode/sync/renders/
.json`` keeps each function's rendered C, +assembly, and summary, so a cluster file can be rebuilt without decompiling +its untouched members again. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Iterable + +from .metadata import write_json, write_text_atomic +from .schema import Cluster, RenderedFunction + +SYNC_SCHEMA = 1 + + +def sync_dir(root: Path) -> Path: + return root / ".tocode" / "sync" + + +def renders_dir(root: Path) -> Path: + return sync_dir(root) / "renders" + + +def manifest_path(root: Path) -> Path: + return sync_dir(root) / "manifest.json" + + +def write_manifest( + root: Path, + *, + clusters: Iterable[Cluster], + backend: str, + decompiler: str, + tree: bool, + entropy: bool, + origin: str | None, +) -> None: + write_json( + manifest_path(root), + { + "schema": SYNC_SCHEMA, + "backend": backend, + "decompiler": decompiler, + "tree": tree, + "entropy": entropy, + "origin": origin, + "clusters": [cluster_to_json(cluster) for cluster in clusters], + }, + ) + + +def read_manifest(root: Path) -> dict[str, Any] | None: + try: + payload = json.loads(manifest_path(root).read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(payload, dict) or payload.get("schema") != SYNC_SCHEMA: + return None + return payload + + +def cluster_to_json(cluster: Cluster) -> dict[str, Any]: + return { + "root": f"0x{cluster.root:x}" if cluster.root >= 0 else cluster.root, + "label": cluster.label, + "summary": cluster.summary, + "folder": cluster.folder, + "file_base": cluster.file_base, + "members": [f"0x{address:x}" for address in cluster.members], + } + + +def cluster_from_json(payload: dict[str, Any]) -> Cluster: + root = payload.get("root") + return Cluster( + root=int(root, 16) if isinstance(root, str) else int(root or 0), + label=str(payload.get("label") or ""), + summary=str(payload.get("summary") or ""), + members=[int(str(item), 16) for item in payload.get("members", [])], + folder=str(payload.get("folder") or "app"), + file_base=payload.get("file_base"), + ) + + +def render_path(root: Path, address: int) -> Path: + return renders_dir(root) / f"{address:016x}.json" + + +def load_render(root: Path, address: int) -> RenderedFunction | None: + from .exporter import _rendered_from_json + + try: + payload = json.loads(render_path(root, address).read_text(encoding="utf-8")) + return _rendered_from_json(payload["rendered"]) + except (OSError, ValueError, KeyError, TypeError): + return None + + +def save_render(root: Path, rendered: RenderedFunction) -> None: + from .exporter import _rendered_to_json + + write_text_atomic( + render_path(root, rendered.address), + json.dumps({"rendered": _rendered_to_json(rendered)}, indent=1) + "\n", + ) + + +def drop_render(root: Path, address: int) -> None: + render_path(root, address).unlink(missing_ok=True) diff --git a/tests/conftest.py b/tests/conftest.py index a2501fc..fc343ec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,7 @@ from __future__ import annotations +from pathlib import Path + import pytest @@ -16,3 +18,22 @@ def _no_local_binary_ninja(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TOCODE_DISABLE_BINJA_HEADLESS", "1") # The CLI's --binja-license writes this variable; restore it after each test. monkeypatch.setenv("TOCODE_BINJA_LICENSE", "") + + +@pytest.fixture +def export(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """A ToCode export of the fake Binary Ninja program, with its own .bndb.""" + import importlib.util + import sys + + from live_client import FAKE_DIR, drop_binaryninja_modules, make_fake_export + + drop_binaryninja_modules() + if importlib.util.find_spec("binaryninja") is not None: + pytest.skip("a real Binary Ninja API is importable in this environment") + monkeypatch.setattr(sys, "path", list(sys.path)) + monkeypatch.delenv("TOCODE_DISABLE_BINJA_HEADLESS", raising=False) + monkeypatch.setenv("TOCODE_BINJA_DIR", str(FAKE_DIR)) + monkeypatch.setenv("TOCODE_BINJA_CACHE_DIR", str(tmp_path / "cache")) + monkeypatch.setenv("TOCODE_FUNCTION_TIMEOUT_SECONDS", "0") + return make_fake_export(tmp_path) diff --git a/tests/fake_binja/binaryninja/__init__.py b/tests/fake_binja/binaryninja/__init__.py index d171d90..4847598 100644 --- a/tests/fake_binja/binaryninja/__init__.py +++ b/tests/fake_binja/binaryninja/__init__.py @@ -20,6 +20,7 @@ import enum import os +import re from pathlib import Path import signal import threading @@ -108,6 +109,7 @@ def __str__(self) -> str: class TypeClass(enum.Enum): StructureTypeClass = 1 IntegerTypeClass = 2 + EnumerationTypeClass = 3 class DisassemblyOption: @@ -201,6 +203,72 @@ def get_llil_at(self, address: int): operation = self._llil.get(int(address)) return None if operation is None else SimpleNamespace(operation=operation) + # -- edits (the subset ToCode's live tools use) ---------------------------- + + comment = "" + + def set_user_type(self, type_obj) -> None: + """Takes a `_ParsedType` of `ret name(type a, ...)`.""" + match = re.match(r"\s*(.+?)\s*\b(\w+)\s*\((.*)\)\s*$", type_obj.decl) + if match is None: + raise ValueError(f"not a function type: {type_obj.decl}") + self.return_type = match.group(1).strip() + params = [] + for index, part in enumerate( + p for p in match.group(3).split(",") if p.strip() and p.strip() != "void" + ): + words = part.strip().rsplit(" ", 1) + name = words[1].lstrip("*") if len(words) == 2 else f"a{index + 1}" + ptype = ( + (words[0] + "*" * (len(words[1]) - len(words[1].lstrip("*")))) + if len(words) == 2 + else part.strip() + ) + params.append(_Var(name, ptype, 100 + index)) + locals_ = [var for var in self.vars if var not in self.parameter_vars] + self.parameter_vars = params + self.vars = params + locals_ + + def set_comment_at(self, address: int, text: str) -> None: + self._view._comments[int(address)] = text + + def create_user_var(self, var, type_, name) -> None: + var.name = name + var.type = str(getattr(type_, "decl", type_)) + + +class _ParsedType: + """What `parse_type_string` returns: keeps the declaration text.""" + + def __init__(self, decl: str) -> None: + self.decl = decl.strip().rstrip(";") + + def __str__(self) -> str: + return self.decl + + +class _UserType: + """A struct/union/enum declared through `define_user_type`.""" + + width = 8 + + def __init__(self, source: str) -> None: + self.source = source.strip() + kind = self.source.split(None, 1)[0] + self.type_class = ( + TypeClass.EnumerationTypeClass + if kind == "enum" + else TypeClass.StructureTypeClass + ) + self.type = ( + "StructureVariant.UnionStructureType" + if kind == "union" + else "StructureVariant.StructStructureType" + ) + + def __str__(self) -> str: + return self.source + class _StructType: type_class = TypeClass.StructureTypeClass @@ -228,6 +296,14 @@ def __init__(self, filename: str, original: str) -> None: def close(self) -> None: self.closed = True + def save_auto_snapshot(self, progress_func=None, settings=None) -> bool: + """Saves into the .bndb the view was opened from.""" + _note(f"save_auto_snapshot {Path(self.filename).name}") + Path(self.filename).write_text( + _MAGIC + self.original_filename, encoding="utf-8" + ) + return True + class BinaryView: def __init__(self, filename: str, original: str) -> None: @@ -356,16 +432,110 @@ def __init__(self, filename: str, original: str) -> None: } self.types = {"point": _StructType(), "score_t": _IntType()} self.type_container = SimpleNamespace(name="container") - self._linear = { - 0x401000: [ - "int32_t main(int32_t argc)", - "{", - " helper();", - ' return puts("hello");', - "}", - ], - 0x401080: ["void helper()", "{", " return;", "}"], - } + self._comments: dict[int, str] = {} + self._data_vars: dict[int, object] = {} + self._main, self._helper, self._puts = main, helper, puts + + # -- pseudo C, rendered from the current names, types and comments -------- + + def _prototype(self, func) -> str: + params = ", ".join(f"{var.type} {var.name}" for var in func.parameter_vars) + return f"{func.return_type} {func.name}({params})" + + def _pseudo_c(self, func) -> list[str]: + body: list[str] = [] + if func.comment: + body.append(f" // {func.comment}") + for address in sorted(self._comments): + if func.start <= address <= func.highest_address: + body.append(f" // {self._comments[address]}") + for var in func.vars: + if var not in func.parameter_vars and var.name != "total": + body.append(f" {var.type} {var.name};") + if func is self._main: + body += [ + f" {self._helper.name}();", + f' return {self._puts.name}("hello");', + ] + elif func is self._helper: + point = self.types.get("point") + fields = point.source.count(";") - 1 if isinstance(point, _UserType) else 2 + body.append(f" struct point *p; /* point has {fields} fields */") + body += [f" {callee.name}();" for callee in func.callees] + body.append(" return;") + else: + body.append(" return 0;") + return [self._prototype(func), "{", *body, "}"] + + # -- edits ------------------------------------------------------------------ + + def parse_type_string(self, decl: str): + match = re.search(r"(\w+)\s*(?:\(|$|\[)", decl.strip().rstrip(";")) + return _ParsedType(decl), (match.group(1) if match else "") + + def parse_types_from_string(self, source: str): + found = {} + for match in re.finditer( + r"\b(struct|union|enum)\s+(\w+)\s*\{[^}]*\}\s*;", source + ): + found[match.group(2)] = _UserType(match.group(0)) + if not found: + raise SyntaxError("no type declarations") + return SimpleNamespace(types=found, variables={}, functions={}) + + def define_user_type(self, name, type_obj) -> None: + _note(f"define_user_type {name}") + self.types[str(name)] = type_obj + + def define_user_symbol(self, symbol) -> None: + for existing in self._symbols: + if existing.address == symbol.address: + existing.name = existing.short_name = symbol.name + return + self._symbols.append( + SimpleNamespace( + address=symbol.address, + name=symbol.name, + short_name=symbol.name, + type=symbol.type, + binding=SymbolBinding.GlobalBinding, + namespace=SimpleNamespace(name=None), + ordinal=None, + ) + ) + + def define_user_data_var(self, address, type_obj) -> None: + self._data_vars[int(address)] = type_obj + + def undefine_user_data_var(self, address) -> None: + self._data_vars.pop(int(address), None) + + def set_comment_at(self, address, text) -> None: + self._comments[int(address)] = text + + def add_user_function(self, address) -> None: + address = int(address) + if address in self._by_start: + return + func = _Function( + self, + address, + f"sub_{address:x}", + size=0x10, + symbol_type=SymbolType.FunctionSymbol, + ) + # The fake's discovered function is called by helper. + self._helper.callees.append(func) + func.callers = [self._helper] + self.functions.append(func) + self._by_start[address] = func + _note(f"add_user_function 0x{address:x}") + + def remove_user_function(self, func) -> None: + self.functions = [item for item in self.functions if item is not func] + self._by_start.pop(func.start, None) + for item in self.functions: + item.callees = [callee for callee in item.callees if callee is not func] @property def segments(self): @@ -438,9 +608,17 @@ def get_next_linear_disassembly_lines(self, cursor): if cursor.done: return [] cursor.done = True - return [ - _Line(text, cursor.address) for text in self._linear.get(cursor.address, []) - ] + func = self._by_start.get(cursor.address) + if func is None or func.hlil is None: + return [] + return [_Line(text, cursor.address) for text in self._pseudo_c(func)] + + +class Symbol: + def __init__(self, sym_type, address, name) -> None: + self.type = sym_type + self.address = int(address) + self.name = name def load(source, update_analysis: bool = True, progress_func=None, options=None): diff --git a/tests/fake_binja/binaryninja/typeprinter.py b/tests/fake_binja/binaryninja/typeprinter.py index 5196b2d..57c048c 100644 --- a/tests/fake_binja/binaryninja/typeprinter.py +++ b/tests/fake_binja/binaryninja/typeprinter.py @@ -9,6 +9,9 @@ def get_type_lines( ): if escaping != 0: raise AssertionError("ToCode asks for unescaped C names") + source = getattr(type, "source", None) + if source is not None: # declared through define_user_type + return source.splitlines() if str(type) == "struct point": return [f"struct {name}", "{", " int32_t x;", " int32_t y;", "};"] return [f"typedef {type} {name};"] diff --git a/tests/live_client.py b/tests/live_client.py new file mode 100644 index 0000000..1ae81ed --- /dev/null +++ b/tests/live_client.py @@ -0,0 +1,162 @@ +"""A minimal MCP stdio client for driving `tocode mcp` in tests and by hand. + +python tests/live_client.py [tool json-args]... +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +import threading +from typing import Any + +SRC = Path(__file__).resolve().parents[1] / "src" + + +class LiveClient: + def __init__(self, *args: str, env: dict[str, str] | None = None) -> None: + environment = dict(os.environ) + environment["PYTHONPATH"] = os.pathsep.join( + [str(SRC), environment.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + environment.update(env or {}) + self.process = subprocess.Popen( + [sys.executable, "-m", "tocode", "mcp", *args], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=environment, + ) + self.stderr: list[str] = [] + self._next = 0 + threading.Thread(target=self._drain_stderr, daemon=True).start() + self.initialize = self.request( + "initialize", + { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "test", "version": "1"}, + }, + ) + self.notify("notifications/initialized", {}) + + def _drain_stderr(self) -> None: + assert self.process.stderr is not None + for raw in self.process.stderr: + self.stderr.append(raw.decode("utf-8", "replace").rstrip()) + + def notify(self, method: str, params: dict[str, Any]) -> None: + self._send({"jsonrpc": "2.0", "method": method, "params": params}) + + def _send(self, payload: dict[str, Any]) -> None: + assert self.process.stdin is not None + self.process.stdin.write(json.dumps(payload).encode() + b"\n") + self.process.stdin.flush() + + def request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: + self._next += 1 + request_id = self._next + self._send( + {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params} + ) + assert self.process.stdout is not None + while True: + line = self.process.stdout.readline() + if not line: + raise RuntimeError( + "server closed stdout:\n" + "\n".join(self.stderr[-30:]) + ) + message = json.loads(line) + if message.get("id") == request_id: + if "error" in message: + raise RuntimeError(message["error"]) + return message["result"] + + def tools(self) -> list[dict[str, Any]]: + return self.request("tools/list", {})["tools"] + + def call(self, tool: str, /, **arguments: Any) -> dict[str, Any]: + """structuredContent of a successful call; raises on a tool error.""" + result = self.request("tools/call", {"name": tool, "arguments": arguments}) + if result.get("isError"): + raise ToolError(result["content"][0]["text"]) + return result.get("structuredContent") or {} + + def close(self) -> None: + if self.process.stdin is not None: + self.process.stdin.close() + try: + self.process.wait(timeout=120) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait() + + def __enter__(self) -> "LiveClient": + return self + + def __exit__(self, *_exc: object) -> None: + self.close() + + +class ToolError(RuntimeError): + pass + + +FAKE_DIR = Path(__file__).resolve().parent / "fake_binja" + + +def fake_client(tmp_path: Path, *args: str) -> LiveClient: + """`tocode mcp` running against the on-disk fake Binary Ninja API.""" + return LiveClient( + *args, + env={ + "TOCODE_BINJA_DIR": str(FAKE_DIR), + "TOCODE_DISABLE_BINJA_HEADLESS": "", + "TOCODE_BINJA_CACHE_DIR": str(tmp_path / "cache"), + "FAKE_BINJA_LOG": str(tmp_path / "binja.log"), + "TOCODE_FUNCTION_TIMEOUT_SECONDS": "0", + }, + ) + + +def drop_binaryninja_modules() -> None: + for name in [ + m for m in sys.modules if m == "binaryninja" or m.startswith("binaryninja.") + ]: + sys.modules.pop(name, None) + + +def make_fake_export(tmp_path: Path) -> Path: + """A ToCode export of the fake program, with its own .bndb. + + The caller sets TOCODE_BINJA_DIR (and friends) for this process first. + """ + from tocode.analysis import create_analyzer + from tocode.exporter import export_binary + from tocode.progress import Progress + + drop_binaryninja_modules() + binary = tmp_path / "sample.elf" + binary.write_bytes(b"\x7fELF" + bytes(range(64))) + out = tmp_path / "sample_decompiler" + with create_analyzer( + binary, backend="binja", binja_dir=FAKE_DIR, progress=Progress(enabled=False) + ) as analyzer: + export_binary(analyzer, out_dir=out, jobs=1) + drop_binaryninja_modules() + return out + + +if __name__ == "__main__": + with LiveClient(sys.argv[1]) as client: + for spec in sys.argv[2:]: + name, _, raw = spec.partition("=") + print(f"== {name}") + print( + json.dumps( + client.call(name, **(json.loads(raw) if raw else {})), indent=1 + )[:4000] + ) diff --git a/tests/test_exporter.py b/tests/test_exporter.py index da05dd9..e9ce2b7 100644 --- a/tests/test_exporter.py +++ b/tests/test_exporter.py @@ -226,7 +226,14 @@ def test_export_binary_writes_source_tree_and_metadata(tmp_path: Path) -> None: assert "Export run started" in (summary.root_dir / "tocode.log").read_text( encoding="utf-8" ) - assert not (summary.root_dir / ".tocode").exists() + # The checkpoint is gone; what `sync` needs stays. + state = summary.root_dir / ".tocode" + assert not (state / "checkpoint.json").exists() + assert not (state / "rendered").exists() + assert (state / "sync" / "manifest.json").is_file() + assert ( + len(list((state / "sync" / "renders").glob("*.json"))) == summary.function_count + ) def test_relativize_sources_strips_common_root() -> None: @@ -391,7 +398,7 @@ def test_export_binary_resumes_from_checkpoint_after_interrupt(tmp_path: Path) - assert resumed.decompile_calls == 1 assert summary.function_count == 2 - assert not (export_root / ".tocode").exists() + assert not (export_root / ".tocode" / "checkpoint.json").exists() resumed_log = (export_root / "tocode.log").read_text(encoding="utf-8") assert "Checkpoint: resuming at 1/2 with 1 cached functions" in resumed_log assert "export main 0x1000 - 48 bytes cached" not in resumed_log diff --git a/tests/test_live_dotnet.py b/tests/test_live_dotnet.py new file mode 100644 index 0000000..7ae9016 --- /dev/null +++ b/tests/test_live_dotnet.py @@ -0,0 +1,294 @@ +"""`tocode mcp` on .NET: the edit log, text edits, targets, and a real session. + +The real-session test builds a small C# program and needs the .NET SDK plus +the decompiler libraries; it is skipped without them (CI has neither). +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from live_client import LiveClient, ToolError +from tocode.live import targets +from tocode.live.dotnet_edits import EditLog, apply_text_edits + + +def _fake_live(tmp_path: Path, entries: list[dict]) -> Any: + log = EditLog(path=tmp_path / "edits.json", entries=entries) + return SimpleNamespace( + edits=log, + session=SimpleNamespace(assemblies=[SimpleNamespace(name="app")]), + ) + + +def test_text_edits_rename_locals_in_their_method_only(tmp_path: Path) -> None: + text = "class A\n{\n void F() { int v1 = 0; }\n void G() { int v1 = 1; }\n}\n" + live = _fake_live( + tmp_path, + [ + { + "op": "rename_local", + "assembly": "app", + "token": "0x06000001", + "old": "v1", + "name": "count", + } + ], + ) + out, ranges = apply_text_edits( + live, 0, text, {0x06000001: (3, 3), 0x06000002: (4, 4)}, 0 + ) + assert "void F() { int count = 0; }" in out + assert "void G() { int v1 = 1; }" in out + assert ranges == {0x06000001: (3, 3), 0x06000002: (4, 4)} + + +def test_text_edits_insert_comments_and_shift_ranges(tmp_path: Path) -> None: + header = "// ToCode header\n" + body = "class A\n{\n void F() { }\n void G() { }\n}\n" + live = _fake_live( + tmp_path, + [ + { + "op": "comment", + "assembly": "app", + "token": "0x06000001", + "text": "first\nsecond", + }, + { + "op": "comment", + "assembly": "app", + "token": "0x06000002", + "text": "third", + }, + ], + ) + ranges = {0x02000002: (1, 5), 0x06000001: (3, 3), 0x06000002: (4, 4)} + out, moved = apply_text_edits(live, 0, header + body, ranges, 1) + lines = out.split("\n") + assert lines[3:8] == [ + " // first", + " // second", + " void F() { }", + " // third", + " void G() { }", + ] + # Ranges stay relative to the body (offset excluded); each member's range + # now starts at its comment. + assert moved[0x06000001] == (3, 5) + assert moved[0x06000002] == (6, 7) + assert moved[0x02000002] == (1, 8) + assert lines[moved[0x06000002][1]] == " void G() { }" # +offset, 0-based + + +def test_edit_log_round_trip_and_latest_comment_wins(tmp_path: Path) -> None: + log = EditLog.open(tmp_path / "missing.json") + assert log.entries == [] and log.synced == 0 + log.entries = [ + {"op": "comment", "assembly": "app", "token": "0x06000001", "text": "old"}, + {"op": "comment", "assembly": "app", "token": "0x06000001", "text": "new"}, + {"op": "comment", "assembly": "app", "token": "0x06000002", "text": "gone"}, + {"op": "comment", "assembly": "app", "token": "0x06000002", "text": ""}, + ] + log.synced = 3 + log.save() + again = EditLog.open(tmp_path / "missing.json") + assert again.synced == 3 and len(again.entries) == 4 + assert again.comments("app") == {0x06000001: "new"} + + +def test_a_dotnet_export_opens_with_its_edit_log(tmp_path: Path) -> None: + binary = tmp_path / "App.dll" + binary.write_bytes(b"MZ") + export = tmp_path / "App_decompiler" + export.mkdir() + (export / "functions.json").write_text('{"functions": []}') + (export / "project.json").write_text( + json.dumps({"binary": str(binary), "backend": "dotnet"}) + ) + target = targets.resolve(export) + assert isinstance(target, targets.Target) + assert target.backend == "dotnet" and target.open_path == binary + assert target.database == export / ".tocode" / "dotnet" / "edits.json" + + +# -- a real session ------------------------------------------------------------- + +PROGRAM = """ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace Demo +{ + internal sealed class Ledger + { + private readonly List _entries = new List(); + public int Limit { get; set; } = 100; + public void Record(int amount) => _entries.Add(amount); + public int Balance() => _entries.Sum(); + } + + internal static class Program + { + private static int Main(string[] args) + { + var ledger = new Ledger(); + ledger.Record(5); + ledger.Record(7); + Console.WriteLine($"balance={ledger.Balance()} limit={ledger.Limit}"); + return 0; + } + } +} +""" + + +def _toolchain() -> bool: + if shutil.which("dotnet") is None: + return False + try: + sdks = subprocess.run( + ["dotnet", "--list-sdks"], capture_output=True, text=True, timeout=60 + ).stdout + except (OSError, subprocess.SubprocessError): + return False + from tocode.backends import dotnet_libs + from tocode.backends.dotnet import probe_dotnet + + return bool(sdks.strip()) and probe_dotnet()[0] and dotnet_libs.is_installed() + + +@pytest.fixture(scope="module") +def dotnet_export(tmp_path_factory: pytest.TempPathFactory) -> Path: + if not _toolchain(): + pytest.skip( + "needs the .NET SDK, a .NET 9+ runtime, and the decompiler libraries" + ) + root = tmp_path_factory.mktemp("dotnet") + project = root / "demo" + project.mkdir() + (project / "demo.csproj").write_text( + 'Exe' + "net9.0Major" + "demodisable" + "" + ) + (project / "Program.cs").write_text(PROGRAM) + built = subprocess.run( + [ + "dotnet", + "build", + "-c", + "Release", + "-o", + str(root / "out"), + "-nologo", + "-v", + "q", + ], + cwd=project, + capture_output=True, + text=True, + timeout=600, + ) + if built.returncode != 0: + pytest.skip(f"dotnet build failed here: {built.stdout[-400:]}") + (root / "out" / "demo.pdb").unlink(missing_ok=True) + export = root / "demo_decompiler" + from tocode import cli + + assert cli.main([str(root / "out" / "demo.dll"), "-o", str(export), "-q"]) == 0 + return export + + +def test_dotnet_session_renames_and_syncs(dotnet_export: Path) -> None: + raw = dotnet_export / "src" / "raw" / "demo" / "Demo" + program_before = (raw / "Program.cs").read_bytes() + with LiveClient(str(dotnet_export)) as client: + info = client.call("info") + assert info["backend"] == "dotnet" + assert "Demo.Ledger::Balance" in { + row["name"] for row in client.call("functions", kind="app")["items"] + } + assert ( + "_entries.Sum()" + in client.call("decompile", target="Demo.Ledger::Balance")["code"] + ) + assert "Demo.Program::Main" in [ + r["function"] + for r in client.call("xrefs", target="Demo.Ledger::Record")["refs"] + ] + with pytest.raises(ToolError, match="not available for .NET"): + client.call("declare_types", c_source="struct x { int a; };") + + # One batch: the method is named by its type's old name. + batch = client.call( + "rename_many", + renames=[ + {"target": "Demo.Ledger", "name": "Account"}, + {"target": "Demo.Ledger::Record", "name": "Deposit"}, + ], + ) + assert batch["renamed"] == 2 and batch["errors"] == [] + client.call("rename", target="Demo.Account::Limit", name="Ceiling") + client.call( + "rename_local", function="Demo.Account::Deposit", old="amount", new="cents" + ) + many = client.call( + "rename_many", + renames=[ + {"target": "Demo.Account::Balance", "name": "Total"}, + {"target": "Nope::x", "name": "y"}, + ], + ) + assert many["renamed"] == 1 and len(many["errors"]) == 1 + client.call( + "set_comment", address="Demo.Account::Total", text="Sum of the deposits" + ) + report = client.call("sync") + assert report["moved"] == [ + { + "from": "src/raw/demo/Demo/Ledger.cs", + "to": "src/raw/demo/Demo/Account.cs", + } + ] + assert set(report["changed"]) == {"Demo.Account", "Demo.Program"} + + account = (raw / "Account.cs").read_text() + assert not (raw / "Ledger.cs").exists() + assert "public void Deposit(int cents)" in account + assert "// Sum of the deposits" in account and "public int Total()" in account + assert "public int Ceiling" in account + program = (raw / "Program.cs").read_text() + assert program != program_before.decode() + assert "Account" in program and ".Deposit(5)" in program and "Ceiling" in program + names = { + row["name"] + for row in json.loads((dotnet_export / "functions.json").read_text())[ + "functions" + ] + } + assert any("Demo.Account::Deposit" in name for name in names) + assert (dotnet_export / "patched" / "demo" / "demo.dll").is_file() + + # Reopening replays the log: nothing pending, names already in place. + with LiveClient(str(dotnet_export)) as client: + assert client.call("status") == {"pending_types": [], "pending_edits": 0} + assert ( + client.call("lookup", name="Deposit")["matches"][0]["name"] + == "Demo.Account::Deposit" + ) + assert client.call("sync")["files_rewritten"] == [] + + # `tocode sync` from the command line opens .NET exports too. + from tocode import cli + + assert cli.main(["sync", str(dotnet_export)]) == 0 diff --git a/tests/test_live_folder.py b/tests/test_live_folder.py new file mode 100644 index 0000000..5a719f5 --- /dev/null +++ b/tests/test_live_folder.py @@ -0,0 +1,88 @@ +"""`tocode mcp` over a folder export: one server, a session per binary.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +import sys + +import pytest + +from live_client import FAKE_DIR, ToolError, drop_binaryninja_modules, fake_client +from test_folder import _app +from tocode.apk_native import NativeOptions +from tocode.folder import FolderOptions, export_folder +from tocode.progress import Progress + + +@pytest.fixture +def folder_export(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + drop_binaryninja_modules() + if importlib.util.find_spec("binaryninja") is not None: + pytest.skip("a real Binary Ninja API is importable in this environment") + monkeypatch.setattr(sys, "path", list(sys.path)) + monkeypatch.delenv("TOCODE_DISABLE_BINJA_HEADLESS", raising=False) + monkeypatch.setenv("TOCODE_BINJA_DIR", str(FAKE_DIR)) + monkeypatch.setenv("TOCODE_BINJA_CACHE_DIR", str(tmp_path / "cache")) + monkeypatch.setenv("TOCODE_FUNCTION_TIMEOUT_SECONDS", "0") + root = _app(tmp_path / "app") + out = tmp_path / "app_decompiled" + summary = export_folder( + root, + options=FolderOptions( + out_dir=out, + parallel_files=1, + native_options=NativeOptions(backend="binja", binja_dir=FAKE_DIR), + ), + progress=Progress(enabled=False), + ) + assert summary.failed == 0, summary.errors + drop_binaryninja_modules() + return out + + +def test_folder_binaries_open_on_demand_and_sync_refreshes_links( + folder_export: Path, tmp_path: Path +) -> None: + with fake_client(tmp_path, "--max-sessions", "2", str(folder_export)) as client: + listed = client.call("sessions") + folder = listed["folders"][0] + relpaths = {row["relpath"] for row in folder["binaries"]} + assert {"bin/app", "lib/libfoo.so.1", "lib/libbar.so"} <= relpaths + assert listed["sessions"] == [] # nothing opened until asked + + app = client.call("info", session="bin/app") + assert app["relpath"] == "bin/app" and app["export_dir"].endswith("bin/app") + assert ( + client.call("info", session="lib/libfoo.so.1")["relpath"] + == "lib/libfoo.so.1" + ) + # Folder-qualified names work too. + assert ( + client.call("functions", session=f"{folder['session']}:bin/app")["total"] + == 3 + ) + + client.call("rename", session="bin/app", target="helper", name="parse_args") + report = client.call("sync", session="bin/app") + assert report["renamed"][0]["new"] == "parse_args" + assert report["folder_links"]["refreshed"] is True + source = "\n".join( + path.read_text() + for path in (folder_export / "bin" / "app" / "src" / "raw").rglob("*.c") + ) + assert "parse_args();" in source + assert json.loads((folder_export / "mappings.json").read_text())["links"] + + # A third binary with --max-sessions 2: the least recently used one + # closes and reopens on demand. + client.call("info", session="lib/libbar.so") + alive = { + row["relpath"]: row["alive"] for row in client.call("sessions")["sessions"] + } + assert sum(alive.values()) == 2 + assert client.call("info", session="bin/app")["relpath"] == "bin/app" + + with pytest.raises(ToolError, match="unknown session"): + client.call("info", session="lib/missing.so") diff --git a/tests/test_live_server.py b/tests/test_live_server.py new file mode 100644 index 0000000..a232501 --- /dev/null +++ b/tests/test_live_server.py @@ -0,0 +1,215 @@ +"""`tocode mcp`: protocol, targets, sessions, and the query tools. + +End-to-end tests run the real server in a subprocess against the on-disk fake +Binary Ninja API (`tests/fake_binja`), which spawned workers import like a +real install. +""" + +from __future__ import annotations + +import io +import json +from pathlib import Path +from typing import Any + +import pytest + +from live_client import ToolError, fake_client as _client +from tocode.live import targets +from tocode.live.jsonrpc import McpApp, StdioServer, ToolCall + +# -- transport ----------------------------------------------------------------- + + +def _serve( + *messages: dict[str, Any], app: McpApp | None = None +) -> list[dict[str, Any]]: + app = app or McpApp( + name="t", + version="1", + instructions="hi", + list_tools=lambda: [{"name": "echo"}], + call_tool=lambda call: { + "content": [{"type": "text", "text": json.dumps(call.arguments)}] + }, + ) + reader = io.BytesIO( + b"".join(json.dumps(m).encode() + b"\n" for m in messages) + b"{bad\n" + ) + writer = io.BytesIO() + StdioServer(app, reader, writer).serve_forever() + return [json.loads(line) for line in writer.getvalue().splitlines()] + + +def test_initialize_negotiates_the_protocol_version() -> None: + replies = _serve( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18"}, + }, + { + "jsonrpc": "2.0", + "id": 2, + "method": "initialize", + "params": {"protocolVersion": "1999-01-01"}, + }, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + {"jsonrpc": "2.0", "id": 3, "method": "tools/list"}, + {"jsonrpc": "2.0", "id": 4, "method": "nope"}, + { + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"a": 1}}, + }, + ) + by_id = {reply.get("id"): reply for reply in replies} + assert by_id[1]["result"]["protocolVersion"] == "2025-06-18" + assert by_id[1]["result"]["serverInfo"]["name"] == "t" + assert by_id[2]["result"]["protocolVersion"] == "2025-11-25" # our latest + assert by_id[3]["result"]["tools"] == [{"name": "echo"}] + assert by_id[4]["error"]["code"] == -32601 + assert json.loads(by_id[5]["result"]["content"][0]["text"]) == {"a": 1} + assert by_id[None]["error"]["code"] == -32700 # the bad line + assert len(replies) == 6 # the notification got no reply + + +def test_progress_notifications_carry_the_token() -> None: + def call_tool(call: ToolCall) -> dict[str, Any]: + call.progress(1, 2, "half") + return {"content": []} + + app = McpApp("t", "1", "", list_tools=list, call_tool=call_tool) + replies = _serve( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "x", "arguments": {}, "_meta": {"progressToken": "p"}}, + }, + app=app, + ) + note = next(r for r in replies if r.get("method") == "notifications/progress") + assert note["params"] == { + "progressToken": "p", + "progress": 1, + "total": 2, + "message": "half", + } + + +# -- targets ------------------------------------------------------------------- + + +def test_export_target_uses_the_exports_own_database( + export: Path, tmp_path: Path +) -> None: + target = targets.resolve(export) + assert isinstance(target, targets.Target) + assert target.database == (export / "sample.bndb").resolve() + assert target.open_path == target.database and target.backend == "binja" + assert target.export_dir == export.resolve() + + other = tmp_path / "elsewhere.bndb" + other.write_text((export / "sample.bndb").read_text()) + chosen = targets.resolve(export, database=other) + assert isinstance(chosen, targets.Target) + assert chosen.open_path == other.resolve() # the agent's choice wins + + # A binary with its default export next to it gets that export attached. + by_binary = targets.resolve(tmp_path / "sample.elf") + assert isinstance(by_binary, targets.Target) + assert by_binary.export_dir == export.resolve() + + +def test_unexported_folder_is_refused(tmp_path: Path) -> None: + (tmp_path / "plain").mkdir() + with pytest.raises(Exception, match="neither a ToCode export"): + targets.resolve(tmp_path / "plain") + + +# -- end to end ---------------------------------------------------------------- + + +def test_query_tools_over_a_live_session(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, str(export)) as client: + assert client.initialize["serverInfo"]["name"] == "tocode" + names = {tool["name"] for tool in client.tools()} + assert {"open", "decompile", "xrefs", "eval", "save"} <= names + + info = client.call("info") + assert info["backend"] == "binja" and info["session"] == "sample.elf" + functions = client.call("functions", filter="main|helper") + assert {row["name"] for row in functions["items"]} == {"main", "helper"} + main = client.call("function", target="main") + assert [callee["name"] for callee in main["callees"]][:1] == ["helper"] + assert main["export"]["c"].startswith("src/raw/") + code = client.call("decompile", target="0x401000") + assert code["prototype"] == "int32_t main(int32_t argc)" + assert "helper();" in code["code"] + assert "push" in client.call("disassemble", target="main")["disassembly"] + strings = client.call("strings", filter="hello") + assert strings["items"][0]["used_by"] == ["main"] + assert ( + client.call("read_bytes", address="0x401000", size=4)["hex"] == "00010203" + ) + hits = client.call("search", pattern="00 01 02", kind="hex")["hits"] + assert hits and hits[0]["address"] == "0x401000" + assert client.call("lookup", name="help")["matches"][0]["name"] == "helper" + point = client.call("type", name="point") + assert "int32_t x;" in point["c_decl"] + + with pytest.raises(ToolError, match="no function or symbol"): + client.call("decompile", target="missing_function") + + first = client.call("eval", code="count = len(bv.functions)\ncount") + assert first["result"] == 3 and first["error"] is None + assert client.call("eval", code="print('hi'); count + 1")["stdout"] == "hi\n" + assert "ZeroDivisionError" in client.call("eval", code="1/0")["error"] + assert client.call("save")["saved"] is True + + +def test_a_hung_call_is_stopped_and_the_session_reopens( + export: Path, tmp_path: Path +) -> None: + with _client(tmp_path, str(export)) as client: + with pytest.raises(ToolError, match="timed out"): + client.call("eval", code="import time; time.sleep(60)", timeout=2) + assert client.call("info")["note"].startswith("reopened after: call timed out") + with pytest.raises(ToolError, match="worker died"): + client.call("eval", code="import os; os._exit(3)") + assert client.call("functions")["total"] == 3 # reopened again + + +def test_session_cap_closes_the_least_recently_used( + export: Path, tmp_path: Path +) -> None: + copy = tmp_path / "copy.bndb" + copy.write_text((export / "sample.bndb").read_text()) + with _client(tmp_path, "--max-sessions", "1") as client: + first = client.call("open", path=str(export))["session"] + second = client.call("open", path=str(copy))["session"] + assert first != second + alive = { + row["session"]: row["alive"] for row in client.call("sessions")["sessions"] + } + assert alive == {first: False, second: True} + assert client.call("info", session=first)["session"] == first # reopens + + +def test_read_only_hides_edits_and_eval(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, "--read-only") as client: + names = {tool["name"] for tool in client.tools()} + assert "eval" not in names and "save" not in names + assert "decompile" in names + annotations = {tool["name"]: tool["annotations"] for tool in client.tools()} + assert annotations["decompile"]["readOnlyHint"] is True + + +def test_stdout_noise_never_reaches_the_protocol(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, str(export)) as client: + # A script printing to the real fd 1 must not corrupt the stream. + client.call("eval", code="import os; os.write(1, b'noise\\n')") + assert client.call("info")["session"] == "sample.elf" diff --git a/tests/test_live_setup.py b/tests/test_live_setup.py new file mode 100644 index 0000000..b7d235b --- /dev/null +++ b/tests/test_live_setup.py @@ -0,0 +1,157 @@ +"""The skill, MCP client registration, and `tocode sync` from the command line.""" + +from __future__ import annotations + +from importlib import resources +import json +from pathlib import Path +import re + +import pytest + +from tocode import cli +from tocode.live import setup +from tocode.live.tools import build_tools, render_reference + +SKILL = resources.files("tocode") / "skill" + + +def test_tools_reference_is_generated_from_the_tool_table() -> None: + # Regenerate: python -m tocode.live.tools > src/tocode/skill/references/tools.md + reference = (SKILL / "references" / "tools.md").read_text(encoding="utf-8") + assert reference == render_reference() + documented = set(re.findall(r"^### `(\w+)`", reference, re.M)) + assert documented == {tool.name for tool in build_tools()} + + +def test_skill_is_a_valid_agent_skill() -> None: + text = (SKILL / "SKILL.md").read_text(encoding="utf-8") + front = re.match(r"---\n(.*?)\n---\n", text, re.S) + assert front is not None + fields = dict(line.split(": ", 1) for line in front.group(1).splitlines()) + assert fields["name"] == "tocode" + assert 100 < len(fields["description"]) <= 1536 + assert len(text.splitlines()) < 500 + for link in re.findall(r"\]\((references/[\w.]+)\)", text): + assert (SKILL / link).is_file(), link + + +def test_skill_install_places_it_for_every_agent( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys +) -> None: + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path / "home")) + assert cli.main(["skill", "install"]) == 0 + for base in (".claude", ".agents"): + installed = tmp_path / "home" / base / "skills" / "tocode" + assert (installed / "SKILL.md").is_file() + assert (installed / "references" / "eval.md").is_file() + assert capsys.readouterr().out.count("Installed the ToCode skill") == 2 + + monkeypatch.chdir(tmp_path) + assert cli.main(["skill", "install", "--agent", "claude", "--project"]) == 0 + assert (tmp_path / ".claude" / "skills" / "tocode" / "SKILL.md").is_file() + assert not (tmp_path / ".agents").exists() + + assert cli.main(["skill", "install", "--path", str(tmp_path / "custom")]) == 0 + assert (tmp_path / "custom" / "tocode" / "SKILL.md").is_file() + + +def test_codex_config_is_added_then_updated_in_place(tmp_path: Path) -> None: + config = tmp_path / "config.toml" + config.write_text('model = "x"\n\n[mcp_servers.other]\ncommand = "o"\n') + setup._write_codex_config(config, ["/bin/tocode", "mcp"]) + first = config.read_text() + assert '[mcp_servers.tocode]\ncommand = "/bin/tocode"\nargs = ["mcp"]' in first + assert '[mcp_servers.other]\ncommand = "o"' in first + + setup._write_codex_config(config, ["/new/tocode", "mcp"]) + second = config.read_text() + assert second.count("[mcp_servers.tocode]") == 1 + assert 'command = "/new/tocode"' in second and "/bin/tocode" not in second + assert 'model = "x"' in second and "[mcp_servers.other]" in second + + +def test_mcp_config_prints_every_client(capsys) -> None: + assert cli.main(["mcp", "--config"]) == 0 + out = capsys.readouterr().out + assert "claude mcp add --scope user tocode --" in out + assert "[mcp_servers.tocode]" in out + assert '"mcpServers"' in out + + +def test_sync_command_updates_an_export(export: Path, capsys) -> None: + assert cli.main(["sync", str(export)]) == 0 + assert "0 re-rendered" in capsys.readouterr().out + assert cli.main(["sync", str(export), "--full", "-q"]) == 0 + assert cli.main(["sync", str(export.parent / "nope")]) == 1 + + +def test_subcommand_names_do_not_hide_the_flat_cli(tmp_path: Path, capsys) -> None: + # A file literally named like a subcommand is still exported with ./name. + with pytest.raises(SystemExit): + cli.main(["--help"]) + assert "--parallel-files" in capsys.readouterr().out + + +CORPUS = Path(__file__).resolve().parents[1] / "eval" / "corpus" +_COMMON = { + "main", + "argc", + "argv", + "line", + "value", + "count", + "result", + "error", + "buffer", + "length", + "index", + "next", + "name", + "type", + "data", + "size", + # Generic reverse-engineering vocabulary the docs need (a .NET "metadata + # token"); it names no program's design. + "token", +} + + +def _corpus_identifiers() -> set[str]: + """Names a blind agent must not be told: the benchmark's functions, + types, enum members, globals, and distinctive field names.""" + found: set[str] = set() + for path in list(CORPUS.rglob("*.c")) + list(CORPUS.rglob("*.h")): + text = path.read_text() + found.update( + re.findall(r"^\w[\w\s\*]*?\b([a-z_]\w*)\s*\([^;{)]*\)\s*\{", text, re.M) + ) + found.update(re.findall(r"\b(?:struct|union|enum)\s+(\w+)\s*\{", text)) + for body in re.findall(r"\benum\s+\w*\s*\{([^}]*)\}", text): + found.update(re.findall(r"\b([A-Z][A-Z0-9_]{2,})\b", body)) + found.update(re.findall(r"^(?:static\s+)?[\w\s\*]+?\b(g_\w+)", text, re.M)) + for body in re.findall(r"\bstruct\s+\w+\s*\{([^}]*)\}", text): + found.update( + n for n in re.findall(r"(\w+)\s*(?:\[[^\]]*\])?;", body) if "_" in n + ) + return {name for name in found if name not in _COMMON and len(name) > 3} + + +def test_the_skill_never_names_the_benchmarks_answers() -> None: + """The blind-recovery benchmark is only blind if the skill does not + contain its programs' identifiers (examples must use other programs).""" + identifiers = _corpus_identifiers() + assert {"item_new", "inventory_add", "kv_entry"} <= identifiers + from tocode.live.server import INSTRUCTIONS + + # Everything an agent is told: the skill, the tool descriptions, and the + # server's instructions. + text = "\n".join( + [path.read_text(encoding="utf-8") for path in Path(str(SKILL)).rglob("*.md")] + + [json.dumps(tool.schema()) for tool in build_tools()] + + [INSTRUCTIONS] + ) + leaked = sorted( + name for name in identifiers if re.search(rf"\b{re.escape(name)}\b", text) + ) + assert leaked == [] diff --git a/tests/test_live_sync.py b/tests/test_live_sync.py new file mode 100644 index 0000000..6b4df3e --- /dev/null +++ b/tests/test_live_sync.py @@ -0,0 +1,214 @@ +"""Edit tools and `sync`: the export follows the database, file by file.""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +import shutil + +import pytest + +from live_client import ToolError +from live_client import fake_client as _client + +SOURCE_SUFFIXES = (".c", ".asm", ".summary") + + +def _snapshot(root: Path) -> dict[str, str]: + files = {} + for path in root.rglob("*"): + if path.is_file() and ".tocode" not in path.parts: + files[path.relative_to(root).as_posix()] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + return files + + +def _changed(before: dict[str, str], after: dict[str, str]) -> set[str]: + return { + key for key in set(before) | set(after) if before.get(key) != after.get(key) + } + + +def _sources(changed: set[str]) -> set[str]: + return {path for path in changed if path.endswith(SOURCE_SUFFIXES)} + + +def _source_text(root: Path) -> str: + return "\n".join( + path.read_text() for path in sorted((root / "src" / "raw").rglob("*.c")) + ) + + +def _file_of(root: Path, name: str) -> str: + index = json.loads((root / "function-index.json").read_text()) + row = next(row for row in index["functions"] if row["name"] == name) + return Path(row["c"]["path"]).stem + + +def test_sync_follows_edits_and_touches_only_what_changed( + export: Path, tmp_path: Path +) -> None: + with _client(tmp_path, str(export)) as client: + before = _snapshot(export) + quiet = client.call("sync") + assert quiet["rerendered"] == 0 and quiet["files_rewritten"] == [] + assert _sources(_changed(before, _snapshot(export))) == set() + + # A rename reaches the caller's code. + before = _snapshot(export) + renamed = client.call("rename", target="helper", name="process_input") + assert set(renamed["pending_sync"]) == {"0x401000", "0x401080"} + assert client.call("status")["pending_functions"] == ["0x401000", "0x401080"] + report = client.call("sync") + assert report["renamed"] == [ + {"address": "0x401080", "old": "helper", "new": "process_input"} + ] + assert report["rerendered"] == 2 + text = _source_text(export) + assert "process_input();" in text and "helper" not in text + rows = { + row["address"]: row + for row in json.loads((export / "functions.json").read_text())["functions"] + } + assert rows["0x401080"]["name"] == "process_input" + changed = _sources(_changed(before, _snapshot(export))) + assert changed and all( + _file_of(export, "process_input") in path for path in changed + ) + assert client.call("status")["pending_functions"] == [] + + # A prototype change re-renders the function and its caller. + report = client.call( + "set_type", + target="process_input", + c_decl="int process_input(char *line, int len)", + ) + report = client.call("sync") + assert report["retyped"] == ["0x401080"] + assert "int process_input(char* line, int len)" in _source_text(export) or ( + "int process_input(char *line, int len)" in _source_text(export) + ) + + # A redefined struct reaches the function whose code uses it. + client.call("declare_types", c_source="struct point { int x; int y; };") + report = client.call("sync") + assert "point" in report["types_modified"] + assert "struct point *p; /* point has 2 fields */" in _source_text(export) + before = _snapshot(export) + client.call("declare_types", c_source="struct point { int x; int y; int z; };") + report = client.call("sync") + assert [item["name"] for item in report["changed"]] == ["process_input"] + assert "point has 3 fields" in _source_text(export) + assert "int z;" in (export / "include" / "sample.types.h").read_text() + assert "types.json" in _changed(before, _snapshot(export)) + + # Comments show up in the decompiled code. + client.call( + "set_comment", + address="process_input", + text="reads one line", + kind="function", + ) + client.call("sync") + assert "// reads one line" in _source_text(export) + + +def test_new_functions_land_next_to_their_caller(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, str(export)) as client: + home = _file_of(export, "helper") + client.call("define_function", address="0x4010c0") + report = client.call("sync") + assert report["added"] == [{"address": "0x4010c0", "name": "sub_4010c0"}] + assert _file_of(export, "sub_4010c0") == home # helper calls it + assert "sub_4010c0();" in _source_text(export) + assert client.call("function", target="sub_4010c0")["export"]["c"].endswith( + tuple(f"-{n}" for n in range(1, 500)) + ) + + client.call("undefine", address="0x4010c0") + report = client.call("sync") + assert report["removed"] == ["0x4010c0"] + assert "sub_4010c0" not in _source_text(export) + + +def test_data_renames_eval_touch_and_full_mode(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, str(export)) as client: + renamed = client.call("rename", target="0x404000", name="g_counter") + assert renamed["kind"] == "data" and renamed["pending_sync"] == ["0x401000"] + client.call("sync") + assert ( + client.call("lookup", name="g_counter")["matches"][0]["address"] + == "0x404000" + ) + + touched = client.call("eval", code="touch('main')") + assert touched["touched"] == ["0x401000"] + assert client.call("sync")["rerendered"] == 1 + + full = client.call("sync", mode="full") + assert ( + full["rerendered"] == 2 and full["files_rewritten"] == [] + ) # nothing differs + + forced = client.call("sync", addresses=["helper"]) + assert forced["rerendered"] == 1 + + +def test_edits_report_backend_errors(export: Path, tmp_path: Path) -> None: + with _client(tmp_path, str(export)) as client: + with pytest.raises(ToolError, match="not a valid name"): + client.call("rename", target="main", name="not valid") + with pytest.raises(ToolError, match="could not parse"): + client.call("declare_types", c_source="this is not C") + with pytest.raises(ToolError, match="no local variable"): + client.call("rename_local", function="main", old="nope", new="x") + client.call("rename_local", function="main", old="argc", new="count") + client.call("sync") + assert "int32_t main(int32_t count)" in _source_text(export) + + +def test_sync_updates_exports_made_before_sync_state( + export: Path, tmp_path: Path +) -> None: + shutil.rmtree(export / ".tocode" / "sync") # an export from an older ToCode + with _client(tmp_path, str(export)) as client: + before = _snapshot(export) + assert client.call("sync")["files_rewritten"] == [] + client.call("rename", target="helper", name="worker") + report = client.call("sync") + assert report["rerendered"] == 2 + assert "worker();" in _source_text(export) + changed = _sources(_changed(before, _snapshot(export))) + assert changed and all(_file_of(export, "worker") in path for path in changed) + assert (export / ".tocode" / "sync" / "manifest.json").is_file() + + +def test_sync_saves_the_database_after_edits(export: Path, tmp_path: Path) -> None: + database = export / "sample.bndb" + stamp = database.stat().st_mtime_ns + with _client(tmp_path, str(export)) as client: + client.call("rename", target="helper", name="saved_helper") + report = client.call("sync") + assert report["database"]["saved"] is True + assert database.stat().st_mtime_ns >= stamp + log = (tmp_path / "binja.log").read_text() + assert "save_auto_snapshot sample.bndb" in log + + +def test_export_tool_creates_a_project_for_a_plain_binary( + tmp_path: Path, export: Path +) -> None: + binary = tmp_path / "other.elf" + binary.write_bytes(b"\x7fELF" + bytes(range(32))) + with _client(tmp_path) as client: + session = client.call("open", path=str(binary), backend="binja")["session"] + with pytest.raises(ToolError, match="no export folder"): + client.call("sync", session=session) + made = client.call("export", session=session) + assert made["functions"] == 2 + assert Path(made["export_dir"]).name == "other_decompiler" + assert client.call("sync", session=session)["rerendered"] == 0 + assert os.path.isdir(tmp_path / "other_decompiler" / ".tocode" / "sync")