Menu
Public repository / f0de6c29b6eff8151dc0f8eebcffcf409026c988
AkurAI-PassVault / crates/passvault-mcp/src/tools.rs
Latest change 11e144cbad3ab1009e59df3196dd0848fb7e016b - AKURAI-PASSVAULT-5 report missing deletes accurately by Ólafur Búi Ólafsson
use akurai_json::Value;
use akurai_storage::BTree;
use passvault_core::crypto::Crypto;
use passvault_core::entities::{normalize_entry_name, Folder, VaultEntry};
use passvault_core::types::Timestamp;
use std::sync::{Arc, Mutex};
fn upper_bound(prefix: &[u8]) -> Vec<u8> {
let mut bound = prefix.to_vec();
bound.push(0xff);
bound
}
fn now() -> Timestamp {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64
}
pub struct ToolRegistry {
db: Arc<Mutex<BTree>>,
crypto: Arc<Crypto>,
}
impl ToolRegistry {
pub fn new(db: Arc<Mutex<BTree>>, crypto: Arc<Crypto>) -> Self {
ToolRegistry { db, crypto }
}
fn json_rpc_result(&self, id: Option<Value>, result: Value) -> String {
let mut resp = Vec::new();
resp.push(("jsonrpc".into(), Value::Str("2.0".into())));
resp.push(("result".into(), result));
if let Some(id_val) = id {
resp.push(("id".into(), id_val));
}
Value::Object(resp).to_json()
}
fn json_rpc_error(&self, id: Option<Value>, code: i64, message: &str) -> String {
let mut resp = Vec::new();
resp.push(("jsonrpc".into(), Value::Str("2.0".into())));
resp.push((
"error".into(),
Value::Object(vec![
("code".into(), Value::Int(code)),
("message".into(), Value::Str(message.into())),
]),
));
if let Some(id_val) = id {
resp.push(("id".into(), id_val));
}
Value::Object(resp).to_json()
}
pub fn handle_request(&self, body: &str) -> String {
let req = match akurai_json::parse(body) {
Ok(v) => v,
Err(e) => return self.json_rpc_error(None, -32700, &format!("Parse error: {e}")),
};
let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");
let id = req.get("id").cloned();
// JSON-RPC notifications carry no `id` and MUST NOT receive a response body.
// Return an empty string; the HTTP layer turns this into `202 Accepted` with no
// body (per the MCP Streamable-HTTP spec). Strict clients (codex/rmcp) error out
// if a notification gets a JSON-RPC response.
if id.is_none() || method.starts_with("notifications/") {
return String::new();
}
match method {
"initialize" => self.json_rpc_result(
id,
Value::Object(vec![
("protocolVersion".into(), Value::Str("2024-11-05".into())),
(
"capabilities".into(),
Value::Object(vec![("tools".into(), Value::Object(vec![]))]),
),
(
"serverInfo".into(),
Value::Object(vec![
("name".into(), Value::Str("akurai-passvault".into())),
("version".into(), Value::Str("0.0.1".into())),
]),
),
]),
),
"tools/list" => {
let tools = self.list_tools();
self.json_rpc_result(
id,
Value::Object(vec![("tools".into(), Value::Array(tools))]),
)
}
"tools/get" => {
let params = match req.get("params") {
Some(Value::Object(pairs)) => pairs.clone(),
_ => vec![],
};
let tool_name = params
.iter()
.find(|(k, _)| k == "name")
.and_then(|(_, v)| v.as_str())
.unwrap_or("");
let tools = self.list_tools();
let tool = tools
.into_iter()
.find(|t| t.get("name").and_then(|v| v.as_str()) == Some(tool_name));
match tool {
Some(t) => self.json_rpc_result(id, Value::Object(vec![("tool".into(), t)])),
None => self.json_rpc_error(id, -32602, &format!("Unknown tool: {tool_name}")),
}
}
"tools/call" => {
let params = match req.get("params") {
Some(Value::Object(pairs)) => pairs.clone(),
_ => vec![],
};
let tool_name = params
.iter()
.find(|(k, _)| k == "name")
.and_then(|(_, v)| v.as_str())
.unwrap_or("");
let arguments: Vec<(String, Value)> = params
.iter()
.find(|(k, _)| k == "arguments")
.and_then(|(_, v)| match v {
Value::Object(pairs) => Some(pairs.clone()),
_ => None,
})
.unwrap_or_default();
match self.execute(tool_name, &arguments) {
Ok(result) => self.json_rpc_result(
id,
Value::Object(vec![(
"content".into(),
Value::Array(vec![Value::Object(vec![
("type".into(), Value::Str("text".into())),
("text".into(), Value::Str(result.to_json())),
])]),
)]),
),
Err(e) => self.json_rpc_error(id, -32603, &format!("Tool error: {e}")),
}
}
_ => self.json_rpc_error(id, -32601, &format!("Method not found: {method}")),
}
}
pub fn list_tools(&self) -> Vec<Value> {
vec![
self.tool_def(
"list_entries",
"List all vault entries (name, uri, username)",
vec![],
),
self.tool_def(
"get_entry",
"Get a single entry by ID with all fields",
vec![("id", "number", "Entry ID", true)],
),
self.tool_def(
"create_entry",
"Create a new vault entry",
vec![
("name", "string", "Entry name", true),
("username", "string", "Username", false),
("password", "string", "Password", false),
("uri", "string", "URI/URL", false),
("notes", "string", "Notes", false),
("folderId", "number", "Folder ID", false),
],
),
self.tool_def(
"update_entry",
"Update an existing entry",
vec![
("id", "number", "Entry ID", true),
("name", "string", "Entry name", false),
("username", "string", "Username", false),
("password", "string", "Password", false),
("uri", "string", "URI/URL", false),
("notes", "string", "Notes", false),
("folderId", "number", "Folder ID", false),
],
),
self.tool_def(
"delete_entry",
"Delete a vault entry",
vec![("id", "number", "Entry ID", true)],
),
self.tool_def(
"generate_password",
"Generate a random password",
vec![
("length", "number", "Password length (default 24)", false),
(
"symbols",
"boolean",
"Include symbols (default true)",
false,
),
],
),
self.tool_def(
"search_entries",
"Search entries by name, URI, or username",
vec![("query", "string", "Search query", true)],
),
self.tool_def("list_folders", "List all folders", vec![]),
self.tool_def(
"create_folder",
"Create a new folder",
vec![("name", "string", "Folder name", true)],
),
self.tool_def(
"delete_folder",
"Delete a folder",
vec![("id", "number", "Folder ID", true)],
),
self.tool_def(
"system_status",
"Report PassVault storage health and record counts",
vec![],
),
]
}
fn tool_def(
&self,
name: &str,
description: &str,
params: Vec<(&str, &str, &str, bool)>,
) -> Value {
let mut properties = Vec::new();
let mut required = Vec::new();
for (pname, ptype, pdesc, is_required) in params {
properties.push((
pname.to_string(),
Value::Object(vec![
("type".into(), Value::Str(ptype.into())),
("description".into(), Value::Str(pdesc.into())),
]),
));
if is_required {
required.push(Value::Str(pname.to_string()));
}
}
Value::Object(vec![
("name".into(), Value::Str(name.into())),
("description".into(), Value::Str(description.into())),
(
"inputSchema".into(),
Value::Object(vec![
("type".into(), Value::Str("object".into())),
("properties".into(), Value::Object(properties)),
("required".into(), Value::Array(required)),
]),
),
])
}
pub fn execute(&self, tool_name: &str, args: &[(String, Value)]) -> Result<Value, String> {
match tool_name {
"list_entries" => self.list_entries(),
"get_entry" => self.get_entry(args),
"create_entry" => self.create_entry(args),
"update_entry" => self.update_entry(args),
"delete_entry" => self.delete_entry(args),
"generate_password" => self.generate_password(args),
"search_entries" => self.search_entries(args),
"list_folders" => self.list_folders(),
"create_folder" => self.create_folder(args),
"delete_folder" => self.delete_folder(args),
"system_status" => self.system_status(),
_ => Err(format!("unknown tool: {tool_name}")),
}
}
fn system_status(&self) -> Result<Value, String> {
let entries = self.scan_entries()?.len() as i64;
let folders = match self.list_folders()? {
Value::Array(values) => values.len() as i64,
_ => 0,
};
Ok(Value::Object(vec![
("status".into(), Value::Str("ok".into())),
("entries".into(), Value::Int(entries)),
("folders".into(), Value::Int(folders)),
]))
}
fn scan_entries(&self) -> Result<Vec<(i64, VaultEntry)>, String> {
let prefix = b"vault:entry:";
let start = prefix.to_vec();
let end = upper_bound(prefix);
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let entries = db.range(&start, &end).map_err(|e| format!("range: {e}"))?;
let mut results = Vec::new();
for (key, val) in entries {
let key_str = String::from_utf8_lossy(&key).to_string();
if key_str.contains("counter") {
continue;
}
let id = key_str
.strip_prefix("vault:entry:")
.and_then(|s| s.parse::<i64>().ok())
.unwrap_or(0);
let decrypted = self
.crypto
.decrypt(&val)
.map_err(|e| format!("decrypt: {e}"))?;
if let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&decrypted)) {
if let Ok(entry) = VaultEntry::from_value(&v) {
results.push((id, entry));
}
}
}
Ok(results)
}
fn get_entry_raw(&self, id: i64) -> Result<VaultEntry, String> {
let key = format!("vault:entry:{id}");
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let bytes = db
.get(key.as_bytes())
.map_err(|e| format!("storage: {e}"))?
.ok_or_else(|| "entry not found".to_string())?;
let decrypted = self
.crypto
.decrypt(&bytes)
.map_err(|e| format!("decrypt: {e}"))?;
let v = akurai_json::parse(&String::from_utf8_lossy(&decrypted))
.map_err(|e| format!("json: {e}"))?;
VaultEntry::from_value(&v).map_err(|e| format!("parse: {e}"))
}
fn list_entries(&self) -> Result<Value, String> {
let entries = self.scan_entries()?;
Ok(Value::Array(
entries
.into_iter()
.map(|(_, e)| e.summary_value())
.collect(),
))
}
fn get_entry(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = get_i64(args, "id").ok_or("missing id")?;
let entry = self.get_entry_raw(id)?;
Ok(entry.to_value())
}
fn create_entry(&self, args: &[(String, Value)]) -> Result<Value, String> {
let name = normalize_entry_name(&get_str(args, "name").ok_or("name is required")?)
.map_err(|e| e.to_string())?;
let username = get_str(args, "username").unwrap_or_default();
let password = get_str(args, "password").unwrap_or_default();
let uri = get_str(args, "uri").unwrap_or_default();
let notes = get_str(args, "notes").unwrap_or_default();
let folder_id = get_i64(args, "folderId");
let ts = now();
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let next_id = next_counter(&mut db, "vault:entry:counter");
let entry = VaultEntry {
id: next_id,
name,
username,
password,
uri,
notes,
folder_id,
created_at: ts,
updated_at: ts,
};
let encrypted = self
.crypto
.encrypt(entry.to_value().to_json().as_bytes())
.map_err(|e| format!("encrypt: {e}"))?;
db.insert(format!("vault:entry:{next_id}").as_bytes(), &encrypted)
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(entry.to_value())
}
fn update_entry(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = get_i64(args, "id").ok_or("missing id")?;
let mut entry = self.get_entry_raw(id)?;
if let Some(v) = get_str(args, "name") {
entry.name = normalize_entry_name(&v).map_err(|e| e.to_string())?;
}
if let Some(v) = get_str(args, "username") {
entry.username = v;
}
if let Some(v) = get_str(args, "password") {
entry.password = v;
}
if let Some(v) = get_str(args, "uri") {
entry.uri = v;
}
if let Some(v) = get_str(args, "notes") {
entry.notes = v;
}
if let Some(value) = get_arg(args, "folderId") {
entry.folder_id = value.as_i64();
}
entry.updated_at = now();
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let encrypted = self
.crypto
.encrypt(entry.to_value().to_json().as_bytes())
.map_err(|e| format!("encrypt: {e}"))?;
db.insert(format!("vault:entry:{id}").as_bytes(), &encrypted)
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(entry.to_value())
}
fn delete_entry(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = get_i64(args, "id").ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let removed = db
.delete(format!("vault:entry:{id}").as_bytes())
.map_err(|e| format!("delete: {e}"))?;
if !removed {
return Err(format!("entry {id} not found"));
}
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("deleted".into(), Value::Bool(true))]))
}
fn generate_password(&self, args: &[(String, Value)]) -> Result<Value, String> {
let length = get_i64(args, "length").unwrap_or(24).clamp(8, 128) as usize;
let symbols = get_bool(args, "symbols").unwrap_or(true);
let mut charset = Vec::new();
charset.extend(b'a'..=b'z');
charset.extend(b'A'..=b'Z');
charset.extend(b'0'..=b'9');
if symbols {
charset.extend(b"!@#$%^&*()_+-=[]{}|;:,.<>?");
}
use rand::Rng;
let password: String = (0..length)
.map(|_| charset[rand::thread_rng().gen_range(0..charset.len())] as char)
.collect();
Ok(Value::Object(vec![(
"password".into(),
Value::Str(password),
)]))
}
fn search_entries(&self, args: &[(String, Value)]) -> Result<Value, String> {
let query = get_str(args, "query")
.ok_or("missing query")?
.to_lowercase();
let entries = self.scan_entries()?;
let results: Vec<Value> = entries
.into_iter()
.filter(|(_, e)| {
e.name.to_lowercase().contains(&query)
|| e.uri.to_lowercase().contains(&query)
|| e.username.to_lowercase().contains(&query)
|| e.notes.to_lowercase().contains(&query)
})
.map(|(_, e)| e.summary_value())
.collect();
Ok(Value::Array(results))
}
fn list_folders(&self) -> Result<Value, String> {
let prefix = b"vault:folder:";
let start = prefix.to_vec();
let end = upper_bound(prefix);
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let entries = db.range(&start, &end).map_err(|e| format!("range: {e}"))?;
let folders: Vec<Value> = entries
.into_iter()
.filter_map(|(_, val)| {
let v = akurai_json::parse(&String::from_utf8_lossy(&val)).ok()?;
Folder::from_value(&v).ok().map(|f| f.to_value())
})
.collect();
Ok(Value::Array(folders))
}
fn create_folder(&self, args: &[(String, Value)]) -> Result<Value, String> {
let name = get_str(args, "name").ok_or("name is required")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let next_id = next_counter(&mut db, "vault:folder:counter");
let folder = Folder {
id: next_id,
name,
created_at: now(),
};
db.insert(
format!("vault:folder:{next_id}").as_bytes(),
folder.to_value().to_json().as_bytes(),
)
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(folder.to_value())
}
fn delete_folder(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = get_i64(args, "id").ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let removed = db
.delete(format!("vault:folder:{id}").as_bytes())
.map_err(|e| format!("delete: {e}"))?;
if !removed {
return Err(format!("folder {id} not found"));
}
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("deleted".into(), Value::Bool(true))]))
}
}
fn get_arg<'a>(args: &'a [(String, Value)], name: &str) -> Option<&'a Value> {
args.iter().find(|(k, _)| k == name).map(|(_, v)| v)
}
fn get_str(args: &[(String, Value)], name: &str) -> Option<String> {
get_arg(args, name).and_then(|v| v.as_str().map(String::from))
}
fn get_i64(args: &[(String, Value)], name: &str) -> Option<i64> {
get_arg(args, name).and_then(|v| v.as_i64())
}
fn get_bool(args: &[(String, Value)], name: &str) -> Option<bool> {
get_arg(args, name).and_then(|v| v.as_bool())
}
fn next_counter(db: &mut BTree, key: &str) -> i64 {
let next = db
.get(key.as_bytes())
.ok()
.flatten()
.and_then(|v| std::str::from_utf8(&v).ok()?.parse::<i64>().ok())
.unwrap_or(0)
+ 1;
let _ = db.insert(key.as_bytes(), next.to_string().as_bytes());
next
}