from typing import Any, List, Optional
import json
import subprocess
from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from mcp.server.sse import SseServerTransport
from starlette.requests import Request
from starlette.routing import Mount, Route
from mcp.server import Server
import uvicorn
import os
mcp = FastMCP("die")
DIE_EXECUTABLE = "diec.exe"
def run_die_command(target_path: str, options: List[str]) -> str:
cmd = [DIE_EXECUTABLE] + options + [target_path]
try:
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
return result.stdout
except subprocess.CalledProcessError as e:
return f"Error executing DIE: {e.stderr if e.stderr else str(e)}"
except FileNotFoundError:
return f"Error: DIE executable not found at '{DIE_EXECUTABLE}'. Please set the correct path."
@mcp.tool()
async def analyze_file(file_path: str, json_output: bool = True, deep_scan: bool = False,
entropy: bool = False, verbose: bool = False) -> str:
if not os.path.isfile(file_path):
return f"Error: File not found at path '{file_path}'"
options = []
if json_output:
options.append("--json")
if deep_scan:
options.append("--deepscan")
if entropy:
options.append("--entropy")
if verbose:
options.append("--verbose")
result = run_die_command(file_path, options)
return result
@mcp.tool()
async def show_special_info(file_path: str, method: str) -> str:
if not os.path.isfile(file_path):
return f"Error: File not found at path '{file_path}'"
options = ["--special", method, "--json"]
result = run_die_command(file_path, options)
return result
@mcp.tool()
async def list_available_methods(file_path: str) -> str:
if not os.path.isfile(file_path):
return f"Error: File not found at path '{file_path}'"
options = ["--showmethods", "--json"]
result = run_die_command(file_path, options)
return result
def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
sse = SseServerTransport("/messages/")
async def handle_sse(request: Request) -> None:
async with sse.connect_sse(
request.scope,
request.receive,
request._send,
) as (read_stream, write_stream):
await mcp_server.run(
read_stream,
write_stream,
mcp_server.create_initialization_options(),
)
return Starlette(
debug=debug,
routes=[
Route("/sse", endpoint=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
],
)
if __name__ == "__main__":
mcp_server = mcp._mcp_server
import argparse
parser = argparse.ArgumentParser(description='Run DIE MCP server with configurable transport')
parser.add_argument('--transport', choices=['stdio', 'sse'], default='stdio',
help='Transport mode (stdio or sse)')
parser.add_argument('--host', default='0.0.0.0',
help='Host to bind to (for SSE mode)')
parser.add_argument('--port', type=int, default=8080,
help='Port to listen on (for SSE mode)')
parser.add_argument('--die-path',
help='Path to the DIE command-line executable (diec.exe)')
args = parser.parse_args()
if args.die_path:
DIE_EXECUTABLE = args.die_path
if args.transport == 'stdio':
mcp.run(transport='stdio')
else:
starlette_app = create_starlette_app(mcp_server, debug=True)
uvicorn.run(starlette_app, host=args.host, port=args.port)
