- Add .gitignore to prevent secret leakage - Implement YOLO guard with double-confirmation flow * Environment variable gate (YOLO_MODE=1) * Typed confirmation phrase + one-time code * Time-limited approval tokens (5 min TTL) * Single-use tokens with audit logging - Add rate limiting (10/min, 100/hour, 500/day) * Token bucket implementation * Per-session tracking * Automatic bucket reset - Integrate safeguards into existing code * Rate limiter in SecureBridge.send_message() * YOLO guard in YOLOMode.execute_command() * Dry-run mode by default - Add security test suite * .gitignore validation * YOLO guard functional tests * Rate limiter verification * Integration checks All security measures tested and verified. No secrets found in git history.
718 lines
26 KiB
Python
718 lines
26 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Secure Claude Code Multi-Agent Bridge
|
|
Production-lean MCP server with auth, redaction, and safety controls
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import hmac
|
|
import hashlib
|
|
import secrets
|
|
import re
|
|
import sqlite3
|
|
import sys
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
from contextlib import contextmanager
|
|
|
|
from mcp.server import Server
|
|
from mcp.types import Tool, TextContent
|
|
|
|
# Import YOLO mode (optional - only if yolo_mode.py is available)
|
|
try:
|
|
from yolo_mode import YOLOMode, create_yolo_config
|
|
YOLO_AVAILABLE = True
|
|
except ImportError:
|
|
YOLO_AVAILABLE = False
|
|
print("⚠️ YOLO mode not available (yolo_mode.py not found)", file=sys.stderr)
|
|
|
|
# Import rate limiter (critical security component)
|
|
try:
|
|
from rate_limiter import RateLimiter
|
|
RATE_LIMITER_AVAILABLE = True
|
|
except ImportError:
|
|
RATE_LIMITER_AVAILABLE = False
|
|
print("⚠️ Rate limiter not available (rate_limiter.py not found)", file=sys.stderr)
|
|
|
|
|
|
class SecretRedactor:
|
|
"""Redact sensitive data from messages"""
|
|
|
|
PATTERNS = [
|
|
(r'AKIA[0-9A-Z]{16}', 'AWS_KEY_REDACTED'),
|
|
(r'-----BEGIN[^-]+PRIVATE KEY-----.*?-----END[^-]+PRIVATE KEY-----', 'PRIVATE_KEY_REDACTED'),
|
|
(r'Bearer [A-Za-z0-9\-._~+/]+=*', 'BEARER_TOKEN_REDACTED'),
|
|
(r'(?i)password["\s:=]+[^\s"]+', 'PASSWORD_REDACTED'),
|
|
(r'(?i)api[_-]?key["\s:=]+[^\s"]+', 'API_KEY_REDACTED'),
|
|
(r'(?i)secret["\s:=]+[^\s"]+', 'SECRET_REDACTED'),
|
|
(r'ghp_[A-Za-z0-9]{36}', 'GITHUB_TOKEN_REDACTED'),
|
|
(r'sk-[A-Za-z0-9]{48}', 'OPENAI_KEY_REDACTED'),
|
|
]
|
|
|
|
@classmethod
|
|
def redact(cls, text: str) -> str:
|
|
"""Redact secrets from text"""
|
|
redacted = text
|
|
for pattern, replacement in cls.PATTERNS:
|
|
redacted = re.sub(pattern, replacement, redacted, flags=re.DOTALL)
|
|
return redacted
|
|
|
|
|
|
class SecureBridge:
|
|
"""Secure message bridge with HMAC authentication"""
|
|
|
|
def __init__(self, db_path: str):
|
|
self.db_path = db_path
|
|
self.master_secret = secrets.token_bytes(32) # Generate on startup
|
|
|
|
# Initialize rate limiter (10 req/min, 100 req/hour, 500 req/day)
|
|
if RATE_LIMITER_AVAILABLE:
|
|
self.rate_limiter = RateLimiter(
|
|
requests_per_minute=10,
|
|
requests_per_hour=100,
|
|
requests_per_day=500
|
|
)
|
|
else:
|
|
self.rate_limiter = None
|
|
|
|
self.init_db()
|
|
|
|
def init_db(self):
|
|
"""Initialize SQLite schema"""
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
|
|
# Conversations with session tokens
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
|
id TEXT PRIMARY KEY,
|
|
session_a_role TEXT NOT NULL,
|
|
session_b_role TEXT NOT NULL,
|
|
session_a_token TEXT NOT NULL,
|
|
session_b_token TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
expires_at TEXT NOT NULL
|
|
)
|
|
''')
|
|
|
|
# Messages with atomic read tracking
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
conversation_id TEXT NOT NULL,
|
|
from_session TEXT NOT NULL,
|
|
to_session TEXT NOT NULL,
|
|
message TEXT NOT NULL,
|
|
metadata TEXT,
|
|
timestamp TEXT NOT NULL,
|
|
read INTEGER DEFAULT 0,
|
|
FOREIGN KEY (conversation_id) REFERENCES conversations(id)
|
|
)
|
|
''')
|
|
|
|
# Session status
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS session_status (
|
|
conversation_id TEXT NOT NULL,
|
|
session_id TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
last_heartbeat TEXT NOT NULL,
|
|
PRIMARY KEY (conversation_id, session_id)
|
|
)
|
|
''')
|
|
|
|
# Audit log
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS audit_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
conversation_id TEXT,
|
|
session_id TEXT,
|
|
action TEXT NOT NULL,
|
|
details TEXT,
|
|
timestamp TEXT NOT NULL
|
|
)
|
|
''')
|
|
|
|
conn.commit()
|
|
|
|
@contextmanager
|
|
def _get_conn(self):
|
|
"""Thread-safe connection context manager"""
|
|
conn = sqlite3.connect(self.db_path, timeout=10.0)
|
|
conn.execute('PRAGMA journal_mode=WAL') # Better concurrency
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
def _generate_session_token(self, conv_id: str, session_id: str) -> str:
|
|
"""Generate HMAC token for session authentication"""
|
|
data = f"{conv_id}:{session_id}:{datetime.utcnow().isoformat()}"
|
|
return hmac.new(self.master_secret, data.encode(), hashlib.sha256).hexdigest()
|
|
|
|
def _verify_token(self, conv_id: str, session_id: str, token: str) -> bool:
|
|
"""Verify session token"""
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
SELECT session_a_token, session_b_token, expires_at
|
|
FROM conversations WHERE id = ?
|
|
''', (conv_id,))
|
|
row = c.fetchone()
|
|
|
|
if not row:
|
|
return False
|
|
|
|
# Check expiration
|
|
expires_at = datetime.fromisoformat(row[2])
|
|
if datetime.utcnow() > expires_at:
|
|
return False
|
|
|
|
# Verify token
|
|
expected_token = row[0] if session_id == 'a' else row[1]
|
|
return hmac.compare_digest(token, expected_token)
|
|
|
|
def _audit_log(self, conv_id: Optional[str], session_id: Optional[str],
|
|
action: str, details: dict):
|
|
"""Log action for audit trail"""
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
INSERT INTO audit_log (conversation_id, session_id, action, details, timestamp)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
''', (conv_id, session_id, action, json.dumps(details), datetime.utcnow().isoformat()))
|
|
conn.commit()
|
|
|
|
def create_conversation(self, session_a_role: str, session_b_role: str) -> dict:
|
|
"""Create new conversation with session tokens"""
|
|
conv_id = f"conv_{secrets.token_hex(8)}"
|
|
token_a = self._generate_session_token(conv_id, 'a')
|
|
token_b = self._generate_session_token(conv_id, 'b')
|
|
|
|
expires_at = datetime.utcnow() + timedelta(hours=3)
|
|
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
INSERT INTO conversations
|
|
(id, session_a_role, session_b_role, session_a_token, session_b_token, created_at, expires_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
''', (conv_id, session_a_role, session_b_role, token_a, token_b,
|
|
datetime.utcnow().isoformat(), expires_at.isoformat()))
|
|
conn.commit()
|
|
|
|
self._audit_log(conv_id, None, 'create_conversation', {
|
|
'roles': [session_a_role, session_b_role]
|
|
})
|
|
|
|
return {
|
|
'conversation_id': conv_id,
|
|
'session_a_token': token_a,
|
|
'session_b_token': token_b,
|
|
'expires_at': expires_at.isoformat()
|
|
}
|
|
|
|
def send_message(self, conv_id: str, session_id: str, token: str,
|
|
message: str, metadata: dict = None) -> dict:
|
|
"""Send message with authentication and redaction"""
|
|
|
|
# Check rate limit FIRST (before expensive operations)
|
|
if self.rate_limiter:
|
|
allowed, reason = self.rate_limiter.check_rate_limit(session_id)
|
|
if not allowed:
|
|
raise ValueError(f"Rate limit exceeded: {reason}")
|
|
|
|
# Verify authentication
|
|
if not self._verify_token(conv_id, session_id, token):
|
|
raise PermissionError("Invalid session token")
|
|
|
|
# Redact secrets
|
|
redacted_message = SecretRedactor.redact(message)
|
|
redacted_metadata = json.loads(SecretRedactor.redact(json.dumps(metadata or {})))
|
|
|
|
to_session = 'b' if session_id == 'a' else 'a'
|
|
|
|
# Atomic insert
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
INSERT INTO messages
|
|
(conversation_id, from_session, to_session, message, metadata, timestamp)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
''', (conv_id, session_id, to_session, redacted_message,
|
|
json.dumps(redacted_metadata), datetime.utcnow().isoformat()))
|
|
conn.commit()
|
|
|
|
self._audit_log(conv_id, session_id, 'send_message', {
|
|
'to': to_session,
|
|
'message_length': len(redacted_message),
|
|
'redacted': message != redacted_message
|
|
})
|
|
|
|
return {'status': 'sent', 'redacted': message != redacted_message}
|
|
|
|
def get_unread_messages(self, conv_id: str, session_id: str, token: str) -> list:
|
|
"""Get and mark messages as read atomically"""
|
|
|
|
if not self._verify_token(conv_id, session_id, token):
|
|
raise PermissionError("Invalid session token")
|
|
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
|
|
# Atomic read + mark
|
|
c.execute('BEGIN IMMEDIATE')
|
|
|
|
c.execute('''
|
|
SELECT id, from_session, message, metadata, timestamp
|
|
FROM messages
|
|
WHERE conversation_id = ? AND to_session = ? AND read = 0
|
|
ORDER BY timestamp ASC
|
|
''', (conv_id, session_id))
|
|
|
|
messages = []
|
|
message_ids = []
|
|
|
|
for row in c.fetchall():
|
|
messages.append({
|
|
'id': row[0],
|
|
'from': row[1],
|
|
'message': row[2],
|
|
'metadata': json.loads(row[3]) if row[3] else {},
|
|
'timestamp': row[4]
|
|
})
|
|
message_ids.append(row[0])
|
|
|
|
# Mark as read
|
|
if message_ids:
|
|
placeholders = ','.join('?' * len(message_ids))
|
|
c.execute(f'UPDATE messages SET read = 1 WHERE id IN ({placeholders})', message_ids)
|
|
|
|
conn.commit()
|
|
|
|
self._audit_log(conv_id, session_id, 'get_messages', {
|
|
'count': len(messages)
|
|
})
|
|
|
|
return messages
|
|
|
|
def update_status(self, conv_id: str, session_id: str, token: str, status: str):
|
|
"""Update session status with heartbeat"""
|
|
|
|
if not self._verify_token(conv_id, session_id, token):
|
|
raise PermissionError("Invalid session token")
|
|
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
INSERT OR REPLACE INTO session_status
|
|
(conversation_id, session_id, status, last_heartbeat)
|
|
VALUES (?, ?, ?, ?)
|
|
''', (conv_id, session_id, status, datetime.utcnow().isoformat()))
|
|
conn.commit()
|
|
|
|
def get_partner_status(self, conv_id: str, session_id: str, token: str) -> dict:
|
|
"""Get partner session status"""
|
|
|
|
if not self._verify_token(conv_id, session_id, token):
|
|
raise PermissionError("Invalid session token")
|
|
|
|
partner = 'b' if session_id == 'a' else 'a'
|
|
|
|
with self._get_conn() as conn:
|
|
c = conn.cursor()
|
|
c.execute('''
|
|
SELECT status, last_heartbeat
|
|
FROM session_status
|
|
WHERE conversation_id = ? AND session_id = ?
|
|
''', (conv_id, partner))
|
|
|
|
row = c.fetchone()
|
|
|
|
if row:
|
|
heartbeat = datetime.fromisoformat(row[1])
|
|
age = (datetime.utcnow() - heartbeat).total_seconds()
|
|
return {
|
|
'status': row[0],
|
|
'last_heartbeat': row[1],
|
|
'age_seconds': int(age),
|
|
'alive': age < 120 # Consider alive if heartbeat within 2 min
|
|
}
|
|
|
|
return {'status': 'unknown', 'alive': False}
|
|
|
|
|
|
# MCP Server Setup
|
|
app = Server("claude-code-bridge-secure")
|
|
bridge = None # Will be initialized with db_path
|
|
yolo = None # Will be initialized if YOLO mode enabled
|
|
|
|
|
|
@app.list_tools()
|
|
async def list_tools() -> list[Tool]:
|
|
"""MCP tool definitions with strict schemas"""
|
|
tools = [
|
|
Tool(
|
|
name="create_conversation",
|
|
description="Initialize a new secure conversation. Returns tokens for both sessions.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"my_role": {
|
|
"type": "string",
|
|
"description": "Your role (e.g., 'backend_developer')",
|
|
"minLength": 3,
|
|
"maxLength": 100
|
|
},
|
|
"partner_role": {
|
|
"type": "string",
|
|
"description": "Partner's role (e.g., 'frontend_developer')",
|
|
"minLength": 3,
|
|
"maxLength": 100
|
|
}
|
|
},
|
|
"required": ["my_role", "partner_role"]
|
|
}
|
|
),
|
|
Tool(
|
|
name="send_to_partner",
|
|
description="Send a message to partner session (authenticated, redacted)",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64},
|
|
"message": {"type": "string", "maxLength": 50000},
|
|
"action_type": {
|
|
"type": "string",
|
|
"enum": ["question", "info", "proposal", "blocked", "complete"]
|
|
},
|
|
"files_involved": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"maxItems": 20
|
|
}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token", "message"]
|
|
}
|
|
),
|
|
Tool(
|
|
name="check_messages",
|
|
description="Check for new messages (atomic read + mark)",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token"]
|
|
}
|
|
),
|
|
Tool(
|
|
name="update_my_status",
|
|
description="Update status with heartbeat",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64},
|
|
"status": {
|
|
"type": "string",
|
|
"enum": ["working", "waiting", "blocked", "complete"]
|
|
}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token", "status"]
|
|
}
|
|
),
|
|
Tool(
|
|
name="check_partner_status",
|
|
description="Get partner session status and liveness",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token"]
|
|
}
|
|
)
|
|
]
|
|
|
|
# Add YOLO mode tools if available
|
|
if YOLO_AVAILABLE:
|
|
tools.extend([
|
|
Tool(
|
|
name="enable_yolo_mode",
|
|
description="⚠️ DANGEROUS: Enable command execution for this conversation. Use with extreme caution!",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64},
|
|
"mode": {
|
|
"type": "string",
|
|
"enum": ["safe", "restricted", "yolo"],
|
|
"description": "safe=read-only, restricted=git/npm/pip, yolo=most commands"
|
|
},
|
|
"workspace": {
|
|
"type": "string",
|
|
"description": "Working directory for command execution"
|
|
},
|
|
"timeout": {
|
|
"type": "integer",
|
|
"description": "Command timeout in seconds (default: 30)",
|
|
"default": 30
|
|
},
|
|
"sandbox": {
|
|
"type": "boolean",
|
|
"description": "Run in Docker sandbox (requires Docker)",
|
|
"default": False
|
|
}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token", "mode"]
|
|
}
|
|
),
|
|
Tool(
|
|
name="execute_command",
|
|
description="Execute a command (requires YOLO mode enabled). Both agents will see the result.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_id": {"type": "string", "pattern": "^conv_[a-f0-9]{16}$"},
|
|
"session_id": {"type": "string", "enum": ["a", "b"]},
|
|
"token": {"type": "string", "minLength": 64, "maxLength": 64},
|
|
"command": {
|
|
"type": "string",
|
|
"description": "Shell command to execute",
|
|
"maxLength": 1000
|
|
}
|
|
},
|
|
"required": ["conversation_id", "session_id", "token", "command"]
|
|
}
|
|
)
|
|
])
|
|
|
|
return tools
|
|
|
|
|
|
@app.call_tool()
|
|
async def call_tool(name: str, arguments: Any) -> list[TextContent]:
|
|
"""Handle tool calls with validation and error handling"""
|
|
|
|
try:
|
|
if name == "create_conversation":
|
|
result = bridge.create_conversation(
|
|
arguments["my_role"],
|
|
arguments["partner_role"]
|
|
)
|
|
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""✅ Secure conversation created!
|
|
|
|
Conversation ID: {result['conversation_id']}
|
|
|
|
Your token (keep secret): {result['session_a_token']}
|
|
Partner token (share securely): {result['session_b_token']}
|
|
|
|
Expires: {result['expires_at']}
|
|
|
|
IMPORTANT: Tokens are required for all operations. Store your token securely.
|
|
Share the conversation ID and partner token with your partner session via a secure channel."""
|
|
)]
|
|
|
|
elif name == "send_to_partner":
|
|
result = bridge.send_message(
|
|
arguments["conversation_id"],
|
|
arguments["session_id"],
|
|
arguments["token"],
|
|
arguments["message"],
|
|
{
|
|
"action_type": arguments.get("action_type", "info"),
|
|
"files_involved": arguments.get("files_involved", [])
|
|
}
|
|
)
|
|
|
|
redacted_notice = "\n⚠️ Secrets were redacted from your message" if result['redacted'] else ""
|
|
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"📤 Message sent to partner session{redacted_notice}"
|
|
)]
|
|
|
|
elif name == "check_messages":
|
|
messages = bridge.get_unread_messages(
|
|
arguments["conversation_id"],
|
|
arguments["session_id"],
|
|
arguments["token"]
|
|
)
|
|
|
|
if not messages:
|
|
return [TextContent(type="text", text="📭 No new messages")]
|
|
|
|
response = f"📬 {len(messages)} new message(s):\n\n"
|
|
for msg in messages:
|
|
response += f"From: Session {msg['from']}\n"
|
|
response += f"Time: {msg['timestamp']}\n"
|
|
if msg['metadata'].get('action_type'):
|
|
response += f"Type: {msg['metadata']['action_type']}\n"
|
|
response += f"\nMessage:\n{msg['message']}\n"
|
|
if msg['metadata'].get('files_involved'):
|
|
response += f"\nFiles: {', '.join(msg['metadata']['files_involved'])}\n"
|
|
response += "\n" + "="*60 + "\n\n"
|
|
|
|
return [TextContent(type="text", text=response)]
|
|
|
|
elif name == "update_my_status":
|
|
bridge.update_status(
|
|
arguments["conversation_id"],
|
|
arguments["session_id"],
|
|
arguments["token"],
|
|
arguments["status"]
|
|
)
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"✅ Status updated: {arguments['status']}"
|
|
)]
|
|
|
|
elif name == "check_partner_status":
|
|
status = bridge.get_partner_status(
|
|
arguments["conversation_id"],
|
|
arguments["session_id"],
|
|
arguments["token"]
|
|
)
|
|
|
|
partner_id = "B" if arguments["session_id"] == "a" else "A"
|
|
alive_indicator = "🟢" if status['alive'] else "🔴"
|
|
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""{alive_indicator} Partner Session {partner_id}
|
|
|
|
Status: {status['status']}
|
|
Last heartbeat: {status.get('last_heartbeat', 'Never')}
|
|
Age: {status.get('age_seconds', 'N/A')}s
|
|
Alive: {status['alive']}"""
|
|
)]
|
|
|
|
# YOLO mode tools
|
|
elif name == "enable_yolo_mode" and YOLO_AVAILABLE:
|
|
global yolo
|
|
if yolo is None:
|
|
yolo = YOLOMode(bridge)
|
|
|
|
config = yolo.set_mode(
|
|
arguments["conversation_id"],
|
|
arguments["mode"],
|
|
workspace=arguments.get("workspace"),
|
|
timeout=arguments.get("timeout", 30),
|
|
sandbox=arguments.get("sandbox", False)
|
|
)
|
|
|
|
mode_warnings = {
|
|
"safe": "✅ Safe mode: Read-only commands only",
|
|
"restricted": "⚠️ Restricted mode: git, npm, pip commands allowed with validation",
|
|
"yolo": "🔥 YOLO MODE: Most commands allowed! Use with extreme caution!"
|
|
}
|
|
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""{mode_warnings[config['mode']]}
|
|
|
|
Workspace: {config['workspace']}
|
|
Timeout: {config['timeout']}s
|
|
Sandbox: {'Enabled (Docker)' if config['sandbox'] else 'Disabled'}
|
|
|
|
Both agents can now execute commands using execute_command tool.
|
|
Results will be visible to both sessions."""
|
|
)]
|
|
|
|
elif name == "execute_command" and YOLO_AVAILABLE:
|
|
if yolo is None:
|
|
return [TextContent(
|
|
type="text",
|
|
text="❌ YOLO mode not initialized. Use enable_yolo_mode first."
|
|
)]
|
|
|
|
result = yolo.execute_command(
|
|
arguments["conversation_id"],
|
|
arguments["session_id"],
|
|
arguments["token"],
|
|
arguments["command"]
|
|
)
|
|
|
|
if result.get('blocked'):
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""🚫 Command blocked
|
|
|
|
Command: {result['command']}
|
|
Reason: {result['reason']}"""
|
|
)]
|
|
|
|
if not result.get('success', False):
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""❌ Command failed
|
|
|
|
{result.get('error', 'Unknown error')}"""
|
|
)]
|
|
|
|
snapshot_info = f"\n📸 Git snapshot: {result['snapshot']}" if result.get('snapshot') else ""
|
|
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"""✅ Command executed
|
|
|
|
Command: {result['command']}
|
|
Exit code: {result['exit_code']}
|
|
Duration: {result['duration']:.2f}s{snapshot_info}
|
|
|
|
STDOUT:
|
|
```
|
|
{result['stdout'][:2000]}{'...' if len(result['stdout']) > 2000 else ''}
|
|
```
|
|
|
|
STDERR:
|
|
```
|
|
{result['stderr'][:1000]}{'...' if len(result['stderr']) > 1000 else ''}
|
|
```
|
|
|
|
Note: Your partner can see this result via check_messages"""
|
|
)]
|
|
|
|
return [TextContent(type="text", text="❌ Unknown tool")]
|
|
|
|
except PermissionError as e:
|
|
return [TextContent(type="text", text=f"🔒 Authentication failed: {str(e)}")]
|
|
except Exception as e:
|
|
return [TextContent(type="text", text=f"❌ Error: {str(e)}")]
|
|
|
|
|
|
async def main(db_path: str = "/tmp/claude_bridge_secure.db"):
|
|
"""Run the secure MCP server"""
|
|
global bridge
|
|
bridge = SecureBridge(db_path)
|
|
|
|
from mcp.server.stdio import stdio_server
|
|
|
|
async with stdio_server() as (read_stream, write_stream):
|
|
await app.run(
|
|
read_stream,
|
|
write_stream,
|
|
app.create_initialization_options()
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
db_path = sys.argv[1] if len(sys.argv) > 1 else "/tmp/claude_bridge_secure.db"
|
|
print(f"Starting secure bridge with database: {db_path}", file=sys.stderr)
|
|
asyncio.run(main(db_path))
|