Menu
Public repository / 653879bcf622c47638f5aab413ebab46937365ab
AkurAI-PassVault / crates/passvault-mcp/src/tools.rs
Latest change 653879bcf622c47638f5aab413ebab46937365ab - Initial: AkurAI-PassVault AES-256-GCM encrypted vault with MCP API by Claude Code
use std::sync::{Arc, Mutex};
use akurai_json::Value;
use akurai_storage::BTree;
use passvault_core::crypto::Crypto;
use passvault_core::entities::{Folder, VaultEntry};
use passvault_core::types::Timestamp;
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();
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())),
])),
]),
),
"notifications/initialized" => self.json_rpc_result(id, Value::Null),
"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}")),
}
}
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),
]),
]
}
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)),
])),
])
}
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),
_ => Err(format!("unknown tool: {tool_name}")),
}
}
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 = get_str(args, "name").ok_or("name is required")?;
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 = v; }
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(v) = get_i64(args, "folderId") { entry.folder_id = Some(v); }
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}"))?;
db.delete(format!("vault:entry:{id}").as_bytes())
.map_err(|e| format!("delete: {e}"))?;
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}"))?;
db.delete(format!("vault:folder:{id}").as_bytes())
.map_err(|e| format!("delete: {e}"))?;
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
}