mirror of
https://github.com/rizinorg/rizin-mcp.git
synced 2026-11-07 06:46:00 -10:00
Add export_annotations and save_project to keep the work
* Everything an agent renamed, commented or patched was lost with `close_file`. * Record every state-changing command the tools run (the annotation commands and verified writes, at resolved addresses) in a journal kept until the next open, also past `close_file`. `export_annotations` returns it as a Rizin script to replay with `. <script>`. * `save_project` saves the session as a Rizin project with `Ps`; it writes a file, so it is offered only with `--allow-write`, and it checks that the file was written. * A test replays an exported script on a fresh analysis and reopens a saved project with `rizin -p`.
This commit is contained in:
1 parent
aa53d2dbec
commit
bb18b601ba
5 files changed
+183
-4
No files matched your search
@@ -129,6 +129,14 @@ produces the full listing, which the server then slices.
|
||||
understood so later disassembly and decompilation show it. Each returns
|
||||
the state Rizin stored, and fails when Rizin only logs an error.
|
||||
|
||||
**Keeping the work**
|
||||
- `export_annotations` — the renames, comments, flags, signatures and byte
|
||||
writes made through the tools since the file was opened, as a Rizin script
|
||||
(also after `close_file`); replay it with `. <script>` on the same binary.
|
||||
- `save_project` — save the session as a Rizin project (`.rzdb`, reopened with
|
||||
`rizin -p`); writes a file, so it needs `--allow-write`. Bytes written to a
|
||||
file opened with `write: true` go to the file directly.
|
||||
|
||||
**Writing** (disabled by default — see *Security*)
|
||||
- `write_hex`, `write_string`, `write_assembly` (instructions separated by `;`
|
||||
or newlines, assembled for the destination address). Every write is read
|
||||
|
||||
@@ -248,6 +248,13 @@ pub struct SetSignatureParams {
|
||||
pub signature: String,
|
||||
}
|
||||
|
||||
/// Parameters for `save_project`.
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
pub struct SaveProjectParams {
|
||||
/// Where to write the Rizin project (`.rzdb`), e.g. `/tmp/sample.rzdb`.
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
/// Parameters for `write_hex` (write mode only).
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
pub struct WriteHexParams {
|
||||
|
||||
+60
-1
@@ -75,6 +75,7 @@ pub const WRITE_TOOLS: &[&str] = &[
|
||||
"write_string",
|
||||
"write_assembly",
|
||||
"flirt_create",
|
||||
"save_project",
|
||||
];
|
||||
|
||||
/// The tools that change the analysis (names, comments, flags, types),
|
||||
@@ -255,7 +256,8 @@ impl RizinServer {
|
||||
}
|
||||
|
||||
/// Run a command that changes state, failing also when Rizin logs an
|
||||
/// error but reports success, as `f`, `wx` and `Fs` do.
|
||||
/// error but reports success, as `f`, `wx` and `Fs` do. A successful
|
||||
/// command is recorded for `export_annotations`.
|
||||
async fn cmd_strict(&self, cmd: &str) -> Result<String, ErrorData> {
|
||||
let out = self.manager.exec(cmd).await.mcp()?;
|
||||
let out = crate::backend::check(cmd, out).mcp()?;
|
||||
@@ -273,6 +275,7 @@ impl RizinServer {
|
||||
}
|
||||
.to_error_data());
|
||||
}
|
||||
self.manager.record(cmd.to_string()).await;
|
||||
Ok(out.output)
|
||||
}
|
||||
|
||||
@@ -390,6 +393,10 @@ impl RizinServer {
|
||||
to_hex(bytes)
|
||||
)));
|
||||
}
|
||||
let addr = self.resolve(loc).await?;
|
||||
self.manager
|
||||
.record(format!("wx {} @ {addr}", to_hex(bytes)))
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1217,6 +1224,58 @@ impl RizinServer {
|
||||
Ok(self.json_result(self.manager.cmdj(&format!("afsj @ {addr}")).await.mcp()?))
|
||||
}
|
||||
|
||||
// --- Persistence -----------------------------------------------------
|
||||
|
||||
/// The recorded changes as a Rizin script.
|
||||
#[tool(
|
||||
description = "Export the renames, comments, flags, signatures and byte writes made \
|
||||
through the tools since the file was opened, as a Rizin script (also after \
|
||||
close_file). Replay it with `. <script>` in Rizin (or the command tool) on the same \
|
||||
binary to restore that work."
|
||||
)]
|
||||
async fn export_annotations(&self) -> Result<CallToolResult, ErrorData> {
|
||||
let (file, commands) = self.manager.journal().await;
|
||||
let mut script = format!(
|
||||
"# rizin-mcp changes for {}\n",
|
||||
file.as_deref().unwrap_or("(no file)")
|
||||
);
|
||||
for cmd in &commands {
|
||||
script.push_str(cmd);
|
||||
script.push('\n');
|
||||
}
|
||||
Ok(self.json_result(json!({
|
||||
"file": file,
|
||||
"commands": commands.len(),
|
||||
"script": script,
|
||||
})))
|
||||
}
|
||||
|
||||
/// Save a Rizin project (`Ps`).
|
||||
#[tool(
|
||||
description = "Save the session (analysis, names, comments, flags, types) as a Rizin \
|
||||
project file (.rzdb), to reopen later with `rizin -p <file>`. Requires --allow-write \
|
||||
(it writes a file to disk) (Ps)."
|
||||
)]
|
||||
async fn save_project(
|
||||
&self,
|
||||
Parameters(p): Parameters<SaveProjectParams>,
|
||||
) -> Result<CallToolResult, ErrorData> {
|
||||
if !self.manager.config().allow_write {
|
||||
return Err(RizinError::NotPermitted(
|
||||
"save_project writes a file to disk; start the server with --allow-write".into(),
|
||||
)
|
||||
.to_error_data());
|
||||
}
|
||||
let path = quote_arg(&p.path).mcp()?;
|
||||
let out = self.manager.exec(&format!("Ps {path}")).await.mcp()?;
|
||||
crate::backend::check(&format!("Ps {path}"), out).mcp()?;
|
||||
let bytes = std::fs::metadata(&p.path).map(|m| m.len()).map_err(|e| {
|
||||
RizinError::Io(format!("the project was not written to {}: {e}", p.path))
|
||||
.to_error_data()
|
||||
})?;
|
||||
Ok(self.json_result(json!({ "path": p.path, "bytes": bytes })))
|
||||
}
|
||||
|
||||
// --- Writing (guarded) ------------------------------------------------
|
||||
|
||||
/// Write hex bytes (`wx`).
|
||||
|
||||
+23
-3
@@ -72,6 +72,11 @@ pub struct SessionManager {
|
||||
struct Slot {
|
||||
session: Option<Session>,
|
||||
ended: Option<String>,
|
||||
/// The file the journal belongs to.
|
||||
journal_file: Option<String>,
|
||||
/// The state-changing commands run through the tools since the file was
|
||||
/// opened. Kept past a close, so the work can still be exported.
|
||||
journal: Vec<String>,
|
||||
}
|
||||
|
||||
impl Slot {
|
||||
@@ -193,6 +198,8 @@ impl SessionManager {
|
||||
*self.inner.lock().await = Slot {
|
||||
session: Some(session),
|
||||
ended: None,
|
||||
journal_file: Some(path.to_string()),
|
||||
journal: Vec::new(),
|
||||
};
|
||||
|
||||
Ok(OpenInfo {
|
||||
@@ -245,9 +252,22 @@ impl SessionManager {
|
||||
|
||||
/// Close the current session, if any. Returns whether one was open.
|
||||
pub async fn close(&self) -> bool {
|
||||
std::mem::take(&mut *self.inner.lock().await)
|
||||
.session
|
||||
.is_some()
|
||||
let mut slot = self.inner.lock().await;
|
||||
slot.ended = None;
|
||||
slot.session.take().is_some()
|
||||
}
|
||||
|
||||
/// Record a command that changed the analysis or the bytes, for
|
||||
/// [`Self::journal`].
|
||||
pub async fn record(&self, cmd: String) {
|
||||
self.inner.lock().await.journal.push(cmd);
|
||||
}
|
||||
|
||||
/// The file of the last opened session and the state-changing commands
|
||||
/// recorded for it, in order.
|
||||
pub async fn journal(&self) -> (Option<String>, Vec<String>) {
|
||||
let slot = self.inner.lock().await;
|
||||
(slot.journal_file.clone(), slot.journal.clone())
|
||||
}
|
||||
|
||||
/// A snapshot of the current session state. A session that died (e.g.
|
||||
|
||||
@@ -781,3 +781,88 @@ async fn annotations_change_the_analysis() {
|
||||
assert_eq!(sig["args"][1]["name"], "args", "{t:?}: {sig}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn work_can_be_exported_and_saved() {
|
||||
let tmp = std::env::temp_dir();
|
||||
let id = std::process::id();
|
||||
for t in transports() {
|
||||
let script = tmp.join(format!("rizin-mcp-{id}-{t:?}.rz"));
|
||||
let project = tmp.join(format!("rizin-mcp-{id}-{t:?}.rzdb"));
|
||||
let s = Session::start(t, |c| c.allow_write = true).await;
|
||||
s.ok(
|
||||
"open_file",
|
||||
json!({ "path": target(), "analysis": "basic" }),
|
||||
)
|
||||
.await;
|
||||
s.ok(
|
||||
"rename_function",
|
||||
json!({ "at": "main", "name": "entry_logic" }),
|
||||
)
|
||||
.await;
|
||||
s.ok(
|
||||
"set_comment",
|
||||
json!({ "address": "entry_logic", "text": "parses; argv @ start" }),
|
||||
)
|
||||
.await;
|
||||
s.ok(
|
||||
"set_flag",
|
||||
json!({ "name": "mcp.mark", "address": "entry0" }),
|
||||
)
|
||||
.await;
|
||||
let saved = s
|
||||
.ok("save_project", json!({ "path": project.to_string_lossy() }))
|
||||
.await;
|
||||
assert!(saved["bytes"].as_u64().unwrap() > 0, "{t:?}");
|
||||
|
||||
// The journal survives closing the file.
|
||||
s.ok("close_file", json!({})).await;
|
||||
let export = s.ok("export_annotations", json!({})).await;
|
||||
assert_eq!(export["commands"], 4, "{t:?}: {export}"); // afn, CC-, CC, f+
|
||||
std::fs::write(&script, export["script"].as_str().unwrap()).unwrap();
|
||||
|
||||
// Replaying the script on a fresh analysis restores the work.
|
||||
s.ok(
|
||||
"open_file",
|
||||
json!({ "path": target(), "analysis": "basic" }),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(s.ok("export_annotations", json!({})).await["commands"], 0);
|
||||
let r = s
|
||||
.call(
|
||||
"command",
|
||||
json!({ "command": format!(". {}", script.to_string_lossy()) }),
|
||||
)
|
||||
.await
|
||||
.expect("replay");
|
||||
assert_ne!(r.is_error, Some(true), "{t:?}: {r:?}");
|
||||
let info = s
|
||||
.items("function_info", json!({ "at": "entry_logic" }))
|
||||
.await;
|
||||
assert_eq!(info[0]["name"], "entry_logic", "{t:?}");
|
||||
let c = s
|
||||
.call("command", json!({ "command": "CC. @ entry_logic" }))
|
||||
.await
|
||||
.expect("CC.");
|
||||
assert_eq!(
|
||||
c.content[0].as_text().unwrap().text.trim(),
|
||||
"parses; argv @ start",
|
||||
"{t:?}"
|
||||
);
|
||||
s.ok("seek", json!({ "address": "mcp.mark" })).await;
|
||||
|
||||
// The project reopens in rizin with the renamed function.
|
||||
if pipe_available() {
|
||||
let out = std::process::Command::new(rizin_binary())
|
||||
.args(["-q", "-p"])
|
||||
.arg(&project)
|
||||
.args(["-c", "afl~entry_logic"])
|
||||
.output()
|
||||
.expect("rizin -p");
|
||||
let stdout = String::from_utf8_lossy(&out.stdout);
|
||||
assert!(stdout.contains("entry_logic"), "{t:?}: {stdout:?}");
|
||||
}
|
||||
let _ = std::fs::remove_file(&script);
|
||||
let _ = std::fs::remove_file(&project);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user