439 lines
18 KiB
Python
439 lines
18 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Any, TextIO
|
|
|
|
from .config import BrokerConfig, ConfigError
|
|
from .executor import ExecutionError, RemoteExecutor
|
|
from .keychain import (
|
|
CredentialError,
|
|
CredentialNotFound,
|
|
CredentialVault,
|
|
MacKeychain,
|
|
UnsupportedPlatformError,
|
|
validate_profile,
|
|
)
|
|
from .popup import PromptCancelled, SecretPrompt, macos_secret_prompt
|
|
|
|
_DEFAULT_CONFIG = Path("~/.config/umbra-trans/remote-broker.json").expanduser()
|
|
_DEFAULT_STATE = Path("~/.config/umbra-trans/remote-broker-state.json").expanduser()
|
|
_MAX_RESULT_CHARS = 1_000_000
|
|
|
|
|
|
class BrokerState:
|
|
"""Non-secret metadata about profiles that have been configured locally."""
|
|
|
|
def __init__(self, path: Path) -> None:
|
|
self.path = path
|
|
self._lock = threading.RLock()
|
|
|
|
def _read(self) -> dict[str, Any]:
|
|
try:
|
|
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
|
except FileNotFoundError:
|
|
return {"profiles": {}}
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError("Unable to read broker state") from exc
|
|
if not isinstance(raw, dict) or not isinstance(raw.get("profiles", {}), dict):
|
|
raise RuntimeError("Broker state is malformed")
|
|
return raw
|
|
|
|
def _write(self, value: dict[str, Any]) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temporary_name = tempfile.mkstemp(prefix=".remote-broker-state.", dir=self.path.parent)
|
|
try:
|
|
os.fchmod(fd, 0o600)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|
json.dump(value, handle, indent=2, sort_keys=True)
|
|
handle.write("\n")
|
|
os.replace(temporary_name, self.path)
|
|
except BaseException:
|
|
try:
|
|
os.unlink(temporary_name)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
|
|
def record(self, profile: str, purpose: str, ttl_seconds: int | None = None) -> None:
|
|
with self._lock:
|
|
value = self._read()
|
|
metadata: dict[str, Any] = {
|
|
"purpose": purpose,
|
|
"updated_at": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
if ttl_seconds is not None:
|
|
metadata["expires_at"] = (
|
|
datetime.now(timezone.utc) + timedelta(seconds=int(ttl_seconds))
|
|
).isoformat()
|
|
value["profiles"][profile] = metadata
|
|
self._write(value)
|
|
|
|
def remove(self, profile: str) -> None:
|
|
with self._lock:
|
|
value = self._read()
|
|
value["profiles"].pop(profile, None)
|
|
self._write(value)
|
|
|
|
def profiles(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
return dict(self._read().get("profiles", {}))
|
|
|
|
|
|
class RemoteCredentialBroker:
|
|
"""MCP-facing broker; secret values never cross the tool boundary."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: CredentialVault,
|
|
config: BrokerConfig,
|
|
state: BrokerState,
|
|
prompt: SecretPrompt,
|
|
executor: RemoteExecutor,
|
|
) -> None:
|
|
self.store = store
|
|
self.config = config
|
|
self.state = state
|
|
self.prompt = prompt
|
|
self.executor = executor
|
|
|
|
@staticmethod
|
|
def _text_result(value: Any, *, is_error: bool = False) -> dict[str, Any]:
|
|
return {
|
|
"content": [{"type": "text", "text": json.dumps(value, ensure_ascii=False)}],
|
|
"isError": is_error,
|
|
}
|
|
|
|
def tool_definitions(self) -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"name": "credential_set",
|
|
"description": (
|
|
"Open a native macOS hidden-password dialog and store the entered secret "
|
|
"in the login Keychain. The secret is never returned to the model."
|
|
),
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"required": ["profile", "purpose"],
|
|
"properties": {
|
|
"profile": {
|
|
"type": "string",
|
|
"pattern": "^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$",
|
|
"description": "Non-secret profile name, for example oldbox-umbra-ssh.",
|
|
},
|
|
"purpose": {
|
|
"type": "string",
|
|
"maxLength": 160,
|
|
"description": "Short label shown in the popup; never include a password.",
|
|
},
|
|
"ttl_seconds": {
|
|
"type": "integer",
|
|
"minimum": 1,
|
|
"maximum": 3600,
|
|
"description": "Optional in-memory lifetime; omit for persistent Keychain storage.",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"name": "credential_status",
|
|
"description": "Report whether configured credential profiles exist; never reveal secrets.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {"profile": {"type": "string"}},
|
|
},
|
|
},
|
|
{
|
|
"name": "credential_delete",
|
|
"description": "Delete a stored credential profile from the local Keychain.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"required": ["profile"],
|
|
"properties": {"profile": {"type": "string"}},
|
|
},
|
|
},
|
|
{
|
|
"name": "remote_exec",
|
|
"description": (
|
|
"Run one bounded command as a configured SSH account. The command is a list, "
|
|
"not a shell string; output is redacted and size-limited."
|
|
),
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"required": ["profile", "command"],
|
|
"properties": {
|
|
"profile": {"type": "string"},
|
|
"command": {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"items": {"type": "string"},
|
|
},
|
|
"timeout_seconds": {
|
|
"type": "integer",
|
|
"minimum": 5,
|
|
"maximum": 3600,
|
|
"default": 120,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"name": "remote_sudo_exec",
|
|
"description": (
|
|
"Run one exact, configured allowlist command through remote sudo. "
|
|
"The password is read internally from Keychain and is never returned."
|
|
),
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"required": ["profile", "command"],
|
|
"properties": {
|
|
"profile": {"type": "string"},
|
|
"command": {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"items": {"type": "string"},
|
|
},
|
|
"timeout_seconds": {
|
|
"type": "integer",
|
|
"minimum": 5,
|
|
"maximum": 3600,
|
|
"default": 120,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"name": "remote_status",
|
|
"description": "Run a small fixed status probe against a configured SSH account.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"required": ["profile"],
|
|
"properties": {"profile": {"type": "string"}},
|
|
},
|
|
},
|
|
]
|
|
|
|
def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
|
|
try:
|
|
if name == "credential_set":
|
|
profile = validate_profile(str(arguments["profile"]))
|
|
purpose = str(arguments["purpose"]).strip()
|
|
if not purpose or len(purpose) > 160:
|
|
raise ExecutionError("purpose must be 1-160 characters")
|
|
ttl_value = arguments.get("ttl_seconds")
|
|
ttl_seconds = None if ttl_value is None else int(ttl_value)
|
|
if ttl_seconds is not None and not 1 <= ttl_seconds <= 3600:
|
|
raise ExecutionError("ttl_seconds must be between 1 and 3600")
|
|
secret = self.prompt(profile, purpose)
|
|
if secret is None:
|
|
raise PromptCancelled("Credential entry was cancelled")
|
|
self.store.set(profile, secret, ttl_seconds=ttl_seconds)
|
|
self.state.record(profile, purpose, ttl_seconds=ttl_seconds)
|
|
return self._text_result(
|
|
{
|
|
"profile": profile,
|
|
"stored": True,
|
|
"temporary": ttl_seconds is not None,
|
|
"ttl_seconds": ttl_seconds,
|
|
}
|
|
)
|
|
if name == "credential_status":
|
|
profile = arguments.get("profile")
|
|
if profile is not None:
|
|
profile = validate_profile(str(profile))
|
|
return self._text_result(
|
|
{"profile": profile, "configured": self.store.exists(profile)}
|
|
)
|
|
statuses = []
|
|
for name_, metadata in sorted(self.state.profiles().items()):
|
|
statuses.append(
|
|
{
|
|
"profile": name_,
|
|
"configured": self.store.exists(name_),
|
|
"purpose": metadata.get("purpose", ""),
|
|
}
|
|
)
|
|
return self._text_result({"profiles": statuses})
|
|
if name == "credential_delete":
|
|
profile = validate_profile(str(arguments["profile"]))
|
|
self.store.delete(profile)
|
|
self.state.remove(profile)
|
|
return self._text_result({"profile": profile, "deleted": True})
|
|
if name in {"remote_exec", "remote_sudo_exec"}:
|
|
profile = self.config.get(str(arguments["profile"]))
|
|
command = arguments.get("command")
|
|
if not isinstance(command, list):
|
|
raise ExecutionError("command must be an array of strings")
|
|
if not all(isinstance(part, str) and part for part in command):
|
|
raise ExecutionError("command must contain only non-empty strings")
|
|
command_tuple = tuple(command)
|
|
if name == "remote_sudo_exec" and not profile.sudo_allowed(command_tuple):
|
|
raise ExecutionError(
|
|
"The requested sudo command is not an exact allowlist entry for this profile"
|
|
)
|
|
timeout = int(arguments.get("timeout_seconds", 120))
|
|
result = self.executor.run(
|
|
profile,
|
|
command,
|
|
sudo=name == "remote_sudo_exec",
|
|
timeout_seconds=timeout,
|
|
)
|
|
output = {
|
|
"exit_code": result.exit_code,
|
|
"stdout": result.stdout[:_MAX_RESULT_CHARS],
|
|
"stderr": result.stderr[:_MAX_RESULT_CHARS],
|
|
"duration_ms": result.duration_ms,
|
|
}
|
|
return self._text_result(output, is_error=result.exit_code != 0)
|
|
if name == "remote_status":
|
|
profile = self.config.get(str(arguments["profile"]))
|
|
results = []
|
|
for command in (["/usr/bin/id", "-un"], ["/usr/bin/hostname"]):
|
|
result = self.executor.run(profile, command, timeout_seconds=30)
|
|
results.append(
|
|
{
|
|
"command": command,
|
|
"exit_code": result.exit_code,
|
|
"stdout": result.stdout.strip(),
|
|
"stderr": result.stderr.strip(),
|
|
}
|
|
)
|
|
return self._text_result({"profile": profile.name, "checks": results})
|
|
raise ExecutionError(f"Unknown tool: {name}")
|
|
except (CredentialError, ConfigError, ExecutionError, PromptCancelled, KeyError, ValueError) as exc:
|
|
return self._text_result({"error": str(exc)}, is_error=True)
|
|
|
|
def handle_json_rpc(self, message: dict[str, Any]) -> dict[str, Any] | None:
|
|
request_id = message.get("id")
|
|
method = message.get("method")
|
|
if method == "initialize":
|
|
params = message.get("params") or {}
|
|
requested_protocol = params.get("protocolVersion")
|
|
supported_protocols = {
|
|
"2024-11-05",
|
|
"2025-03-26",
|
|
"2025-06-18",
|
|
"2025-11-25",
|
|
}
|
|
protocol = (
|
|
requested_protocol
|
|
if isinstance(requested_protocol, str) and requested_protocol in supported_protocols
|
|
else "2024-11-05"
|
|
)
|
|
result: dict[str, Any] = {
|
|
"protocolVersion": protocol,
|
|
"capabilities": {"tools": {"listChanged": False}},
|
|
"serverInfo": {"name": "umbratrans-remote-broker", "version": "0.1.0"},
|
|
}
|
|
return self._response(request_id, result=result)
|
|
if method in {"notifications/initialized", "notifications/cancelled"}:
|
|
return None
|
|
if method == "ping":
|
|
return self._response(request_id, result={})
|
|
if method == "tools/list":
|
|
return self._response(request_id, result={"tools": self.tool_definitions()})
|
|
if method == "resources/list":
|
|
return self._response(request_id, result={"resources": []})
|
|
if method == "prompts/list":
|
|
return self._response(request_id, result={"prompts": []})
|
|
if method == "tools/call":
|
|
params = message.get("params") or {}
|
|
name = params.get("name")
|
|
arguments = params.get("arguments") or {}
|
|
if not isinstance(name, str) or not isinstance(arguments, dict):
|
|
return self._error(request_id, -32602, "Invalid tools/call parameters")
|
|
return self._response(request_id, result=self.call_tool(name, arguments))
|
|
if request_id is not None:
|
|
return self._error(request_id, -32601, f"Method not found: {method}")
|
|
return None
|
|
|
|
@staticmethod
|
|
def _response(request_id: Any, *, result: Any) -> dict[str, Any]:
|
|
return {"jsonrpc": "2.0", "id": request_id, "result": result}
|
|
|
|
@staticmethod
|
|
def _error(request_id: Any, code: int, message: str) -> dict[str, Any]:
|
|
return {"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}}
|
|
|
|
def serve(self, input_stream: TextIO = sys.stdin, output_stream: TextIO = sys.stdout) -> None:
|
|
for line in input_stream:
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
message = json.loads(line)
|
|
if not isinstance(message, dict):
|
|
raise ValueError("message must be a JSON object")
|
|
response = self.handle_json_rpc(message)
|
|
except json.JSONDecodeError:
|
|
response = self._error(None, -32700, "Parse error")
|
|
except Exception:
|
|
response = self._error(None, -32603, "Internal broker error")
|
|
if response is not None:
|
|
output_stream.write(json.dumps(response, ensure_ascii=False) + "\n")
|
|
output_stream.flush()
|
|
|
|
|
|
def _load_runtime() -> tuple[CredentialVault, BrokerConfig, BrokerState, RemoteCredentialBroker]:
|
|
config_path = Path(os.environ.get("UMBRA_BROKER_CONFIG", str(_DEFAULT_CONFIG))).expanduser()
|
|
state_path = Path(os.environ.get("UMBRA_BROKER_STATE", str(_DEFAULT_STATE))).expanduser()
|
|
store = CredentialVault()
|
|
config = BrokerConfig.load(config_path)
|
|
state = BrokerState(state_path)
|
|
audit_path = Path(
|
|
os.environ.get("UMBRA_BROKER_AUDIT_LOG", "~/.local/state/umbratrans/remote-broker.log")
|
|
).expanduser()
|
|
executor = RemoteExecutor(store, audit_path=audit_path)
|
|
broker = RemoteCredentialBroker(store, config, state, macos_secret_prompt, executor)
|
|
return store, config, state, broker
|
|
|
|
|
|
def _cli_credential_set(profile: str, purpose: str) -> int:
|
|
store, _, state, _ = _load_runtime()
|
|
try:
|
|
secret = macos_secret_prompt(validate_profile(profile), purpose)
|
|
if secret is None:
|
|
return 2
|
|
store.set(profile, secret)
|
|
state.record(profile, purpose)
|
|
except (CredentialError, UnsupportedPlatformError, PromptCancelled, RuntimeError) as exc:
|
|
print(f"credential setup failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
print(f"stored credential profile {profile!r}; secret was not printed")
|
|
return 0
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description="UmbraTrans local remote credential MCP broker")
|
|
parser.add_argument("--credential-profile", help="open the popup for one profile and exit")
|
|
parser.add_argument("--purpose", default="remote access", help="label shown in the credential popup")
|
|
args = parser.parse_args(argv)
|
|
if args.credential_profile:
|
|
return _cli_credential_set(args.credential_profile, args.purpose)
|
|
try:
|
|
_, _, _, broker = _load_runtime()
|
|
broker.serve()
|
|
except (ConfigError, CredentialError, UnsupportedPlatformError, RuntimeError) as exc:
|
|
print(f"remote broker startup failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|