#!/usr/bin/env python3
"""
Memory Forensics MCP Server
Provides memory analysis capabilities via Volatility 3 through the MCP protocol.
Works with any MCP client (Claude Code, Claude Desktop, custom clients, etc.)
"""
import asyncio
import logging
from pathlib import Path
from typing import Any, Optional
import zipfile
import tempfile
import shutil
from datetime import datetime
from mcp.server.models import InitializationOptions
from mcp.server import NotificationOptions, Server
from mcp.server.stdio import stdio_server
from mcp.types import (
Resource,
Tool,
TextContent,
ImageContent,
EmbeddedResource,
LoggingLevel
)
from database import ForensicsDatabase
from volatility_handler import VolatilityHandler
from provenance import ProvenanceTracker
from hashing import get_or_calculate_hashes, format_hashes
from exporters import DataExporter
from timeline import TimelineGenerator
from anomaly_detector import AnomalyDetector
from extractors import create_extractor
from validation import DataValidator
from cleanup import (
ManagedExtraction, cleanup_old_extractions, cleanup_all_extractions,
get_disk_usage, list_extractions
)
from config import (
DB_PATH, DUMPS_DIR, EXPORT_DIR, EXTRACTED_FILES_DIR, EXTRACTION_DIR,
LLM_PROFILE, CURRENT_PROFILE_SETTINGS,
EXTRACTION_RETENTION_HOURS, AUTO_CLEANUP_ON_STARTUP, DATA_DIR
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("memory-forensics-mcp")
logger.info(f"LLM Profile: {LLM_PROFILE}")
logger.info(f"Output format: {CURRENT_PROFILE_SETTINGS.get('format')}")
# Global database instance
db = ForensicsDatabase(DB_PATH)
# Global provenance tracker
provenance_tracker = ProvenanceTracker(db)
# Cache for VolatilityHandler instances
vol_handlers = {}
def get_dump_id(file_path: str) -> str:
"""Generate dump ID from filename"""
return Path(file_path).stem
def extract_dump_if_needed(dump_path: Path) -> Path:
"""Extract memory dump from zip if necessary"""
if dump_path.suffix.lower() == '.zip':
logger.info(f"Extracting {dump_path.name}...")
# Use persistent extraction directory instead of tmpfs
import time
dump_id = get_dump_id(str(dump_path))
timestamp = int(time.time())
temp_dir = EXTRACTION_DIR / f"memdump_{dump_id}_{timestamp}"
temp_dir.mkdir(parents=True, exist_ok=True)
with zipfile.ZipFile(dump_path, 'r') as zip_ref:
# Extract all files
zip_ref.extractall(temp_dir)
# Find the actual memory dump file (usually .raw, .mem, .dmp, .vmem)
for ext in ['.raw', '.mem', '.dmp', '.vmem', '.bin']:
extracted_files = list(temp_dir.rglob(f'*{ext}'))
if extracted_files:
logger.info(f"Extracted to: {extracted_files[0]}")
return extracted_files[0]
# If no specific extension, return first file
files = list(temp_dir.iterdir())
if files:
logger.info(f"Extracted to: {files[0]}")
return files[0]
raise ValueError(f"No memory dump found in {dump_path}")
return dump_path
def get_volatility_handler(dump_id: str) -> Optional[VolatilityHandler]:
"""Get or create VolatilityHandler for a dump"""
if dump_id in vol_handlers:
return vol_handlers[dump_id]
# Find the dump file
for dump_file in DUMPS_DIR.iterdir():
if get_dump_id(str(dump_file)) == dump_id:
# Extract if needed
actual_dump = extract_dump_if_needed(dump_file)
# Create handler with dump_id and provenance tracking
handler = VolatilityHandler(
dump_path=actual_dump,
dump_id=dump_id,
provenance_tracker=provenance_tracker
)
vol_handlers[dump_id] = handler
return handler
return None
def adapt_description(description: str) -> str:
"""Adapt tool description based on LLM profile"""
max_length = CURRENT_PROFILE_SETTINGS.get('max_description_length', 500)
if len(description) <= max_length:
return description
# Truncate and add ellipsis
return description[:max_length-3] + "..."
# Create the MCP server
server = Server("memory-forensics")
@server.list_resources()
async def handle_list_resources() -> list[Resource]:
"""List available memory dumps as resources"""
dumps = await db.list_dumps()
resources = []
for dump in dumps:
resources.append(Resource(
uri=f"memdump://{dump['dump_id']}",
name=f"Memory Dump: {dump['dump_id']}",
description=f"Memory dump from {dump['file_path']} ({dump.get('os_type', 'Unknown OS')})",
mimeType="application/x-memory-dump"
))
return resources
@server.list_tools()
async def handle_list_tools() -> list[Tool]:
"""List available memory forensics tools"""
return [
Tool(
name="list_dumps",
description=adapt_description("List all available memory dumps for analysis"),
inputSchema={
"type": "object",
"properties": {},
"required": []
}
),
Tool(
name="process_dump",
description=adapt_description("Process a memory dump with Volatility 3 and extract artifacts (processes, network, etc.)"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump to process (from list_dumps)"
}
},
"required": ["dump_id"]
}
),
Tool(
name="list_processes",
description=adapt_description("List all processes from a memory dump"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"suspicious_only": {
"type": "boolean",
"description": "Only show suspicious processes",
"default": False
}
},
"required": ["dump_id"]
}
),
Tool(
name="analyze_process",
description=adapt_description("Get detailed information about a specific process (command line, DLLs, network connections)"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"pid": {
"type": "integer",
"description": "Process ID to analyze"
}
},
"required": ["dump_id", "pid"]
}
),
Tool(
name="detect_code_injection",
description=adapt_description("Detect potential code injection using malfind (unbacked executable memory regions)"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"pid": {
"type": "integer",
"description": "Optional: only check specific process",
"default": None
}
},
"required": ["dump_id"]
}
),
Tool(
name="network_analysis",
description=adapt_description("Analyze network connections and correlate with processes"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"remote_ip": {
"type": "string",
"description": "Optional: filter by remote IP address"
}
},
"required": ["dump_id"]
}
),
Tool(
name="detect_hidden_processes",
description=adapt_description("Find hidden processes by comparing psscan and pslist"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
}
},
"required": ["dump_id"]
}
),
Tool(
name="get_process_tree",
description=adapt_description("Get hierarchical process tree showing parent-child relationships"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
}
},
"required": ["dump_id"]
}
),
Tool(
name="get_dump_metadata",
description=adapt_description("Get detailed metadata about a memory dump including hashes, OS info, and processing statistics"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
}
},
"required": ["dump_id"]
}
),
Tool(
name="export_data",
description=adapt_description("Export forensic data in JSON, CSV, or HTML format"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"format": {
"type": "string",
"description": "Export format: 'json', 'csv', or 'html'",
"enum": ["json", "csv", "html"]
},
"data_type": {
"type": "string",
"description": "Type of data to export (for csv): 'processes', 'network', 'memory_regions', or 'all' (for json/html)",
"default": "all"
},
"output_filename": {
"type": "string",
"description": "Optional output filename (will be placed in exports directory)",
"default": None
}
},
"required": ["dump_id", "format"]
}
),
Tool(
name="get_command_history",
description=adapt_description("View all Volatility commands executed for a dump (provenance/audit trail)"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"limit": {
"type": "integer",
"description": "Maximum number of commands to return",
"default": 50
}
},
"required": ["dump_id"]
}
),
Tool(
name="generate_timeline",
description=adapt_description("Generate chronological timeline of events from memory dump"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"format": {
"type": "string",
"description": "Output format: 'summary' (display), 'json', 'csv', or 'text'",
"enum": ["summary", "json", "csv", "text"],
"default": "summary"
},
"suspicious_only": {
"type": "boolean",
"description": "Only include suspicious events",
"default": False
},
"output_filename": {
"type": "string",
"description": "Optional output filename for file exports",
"default": None
}
},
"required": ["dump_id"]
}
),
Tool(
name="detect_anomalies",
description=adapt_description("Detect suspicious process behavior (wrong parents, typosquatting, unusual paths)"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
}
},
"required": ["dump_id"]
}
),
Tool(
name="health_check",
description=adapt_description("Check data integrity and consistency for a processed memory dump"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump to validate"
}
},
"required": ["dump_id"]
}
),
Tool(
name="extract_process",
description=adapt_description("Extract detailed process information to JSON file"),
inputSchema={
"type": "object",
"properties": {
"dump_id": {
"type": "string",
"description": "ID of the memory dump"
},
"pid": {
"type": "integer",
"description": "Process ID to extract"
}
},
"required": ["dump_id", "pid"]
}
),
Tool(
name="cleanup_extractions",
description=adapt_description("Clean up old memory dump extractions to free disk space"),
inputSchema={
"type": "object",
"properties": {
"mode": {
"type": "string",
"description": "Cleanup mode: 'old' (remove extractions older than retention period), 'all' (remove all), 'list' (show extractions)",
"enum": ["old", "all", "list"],
"default": "old"
},
"dry_run": {
"type": "boolean",
"description": "If true, show what would be removed without actually deleting",
"default": False
}
},
"required": []
}
),
Tool(
name="get_disk_usage",
description=adapt_description("Get disk space usage statistics for the MCP server (database, exports, extractions)"),
inputSchema={
"type": "object",
"properties": {},
"required": []
}
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict) -> list[TextContent]:
"""Handle tool execution"""
try:
if name == "list_dumps":
# Check for available dumps in the dumps directory
available_dumps = []
for dump_file in DUMPS_DIR.iterdir():
if dump_file.is_file() and dump_file.suffix in ['.zip', '.raw', '.mem', '.dmp', '.vmem']:
dump_id = get_dump_id(str(dump_file))
file_size = dump_file.stat().st_size
# Check if already in database
existing = await db.get_dump(dump_id)
if not existing:
# Register new dump
await db.add_dump(dump_id, str(dump_file), file_size)
available_dumps.append({
'dump_id': dump_id,
'file_name': dump_file.name,
'file_size_mb': round(file_size / (1024*1024), 2),
'status': existing.get('status', 'new') if existing else 'new'
})
result = f"Found {len(available_dumps)} memory dumps:\n\n"
for dump in available_dumps:
result += f"- **{dump['dump_id']}**\n"
result += f" File: {dump['file_name']}\n"
result += f" Size: {dump['file_size_mb']} MB\n"
result += f" Status: {dump['status']}\n\n"
return [TextContent(type="text", text=result)]
elif name == "process_dump":
dump_id = arguments["dump_id"]
logger.info(f"Processing dump: {dump_id}")
# Get Volatility handler
vol = get_volatility_handler(dump_id)
if not vol:
return [TextContent(type="text", text=f"Error: Dump '{dump_id}' not found")]
result = f"Processing memory dump: {dump_id}\n\n"
# Extract processes
result += "Extracting processes...\n"
processes = await vol.list_processes()
await db.add_processes(dump_id, processes)
result += f"[OK] Found {len(processes)} processes\n\n"
# Extract network connections
result += "Extracting network connections...\n"
connections = await vol.get_network_connections()
# Clear old network data before inserting to prevent duplicates on reprocessing
await db.clear_network_connections(dump_id)
await db.add_network_connections(dump_id, connections)
result += f"[OK] Found {len(connections)} network connections\n\n"
# Detect code injection
result += "Scanning for code injection...\n"
suspicious_regions = await vol.detect_malfind()
# Clear old memory region data before inserting to prevent duplicates on reprocessing
await db.clear_memory_regions(dump_id)
await db.add_memory_regions(dump_id, suspicious_regions)
result += f"[OK] Found {len(suspicious_regions)} suspicious memory regions\n\n"
# Detect hidden processes
result += "Detecting hidden processes...\n"
hidden_pids = await vol.detect_hidden_processes()
# Mark hidden processes in database
if hidden_pids:
for pid in hidden_pids:
proc = await db.get_process_by_pid(dump_id, pid)
if proc:
proc['is_hidden'] = True
proc['is_suspicious'] = True
await db.add_processes(dump_id, [proc])
result += f"[OK] Found {len(hidden_pids)} hidden processes: {hidden_pids}\n\n"
else:
result += "[OK] No hidden processes detected\n\n"
# Update dump status
dump = await db.get_dump(dump_id)
if dump:
await db.add_dump(dump_id, dump['file_path'], dump['file_size'], dump.get('os_type'))
# Validate data integrity
result += "Validating data integrity...\n"
validator = DataValidator(db)
# Compare Volatility results with database
db_processes = await db.get_processes(dump_id)
db_connections = await db.get_network_connections(dump_id)
validation_warnings = await validator.compare_volatility_to_database(
dump_id,
{
'processes': len(processes),
'network_connections': len(connections),
'suspicious_memory_regions': len(suspicious_regions)
},
{
'processes': len(db_processes),
'network_connections': len(db_connections),
'suspicious_memory_regions': len(suspicious_regions)
}
)
if validation_warnings:
result += "\n**Data Integrity Warnings:**\n"
for warning in validation_warnings:
result += f"- WARNING: {warning}\n"
logger.warning(f"[{dump_id}] {warning}")
else:
result += "[OK] All data validated successfully\n"
result += "\n**Processing complete!** You can now use other tools to analyze the extracted data."
return [TextContent(type="text", text=result)]
elif name == "list_processes":
dump_id = arguments["dump_id"]
suspicious_only = arguments.get("suspicious_only", False)
processes = await db.get_processes(dump_id, suspicious_only)
if not processes:
return [TextContent(type="text", text=f"No processes found. Run 'process_dump' first for dump: {dump_id}")]
result = f"**Processes in {dump_id}**"
if suspicious_only:
result += " (suspicious only)"
result += f"\n\nTotal: {len(processes)} processes\n\n"
# Format as table
result += "| PID | PPID | Name | Created | Flags |\n"
result += "|-----|------|------|---------|-------|\n"
for proc in processes:
flags = []
if proc.get('is_hidden'):
flags.append('HIDDEN')
if proc.get('is_suspicious'):
flags.append('SUSPICIOUS')
flag_str = ', '.join(flags) if flags else '-'
result += f"| {proc['pid']} | {proc.get('ppid', '-')} | {proc.get('name', 'N/A')} | {proc.get('create_time', 'N/A')[:19]} | {flag_str} |\n"
return [TextContent(type="text", text=result)]
elif name == "analyze_process":
dump_id = arguments["dump_id"]
pid = arguments["pid"]
# Get process info
process = await db.get_process_by_pid(dump_id, pid)
if not process:
return [TextContent(type="text", text=f"Process {pid} not found in dump {dump_id}")]
result = f"**Process Analysis: PID {pid}**\n\n"
result += f"**Basic Information:**\n"
result += f"- Name: {process.get('name', 'N/A')}\n"
result += f"- Path: {process.get('path', 'N/A')}\n"
result += f"- Parent PID: {process.get('ppid', 'N/A')}\n"
result += f"- Created: {process.get('create_time', 'N/A')}\n"
if process.get('is_hidden'):
result += f"- [WARNING] **HIDDEN PROCESS** (not in standard process list)\n"
if process.get('is_suspicious'):
result += f"- [WARNING] **MARKED AS SUSPICIOUS**\n"
result += "\n"
# Get command line
vol = get_volatility_handler(dump_id)
if vol:
cmdlines = await vol.get_cmdline(pid)
if cmdlines:
result += f"**Command Line:**\n```\n{cmdlines[0].get('cmdline', 'N/A')}\n```\n\n"
# Get DLLs
result += "**Loaded DLLs:**\n"
dlls = await vol.get_dlls(pid)
if dlls:
for dll in dlls[:10]: # Show first 10
result += f"- {dll.get('name', 'N/A')} @ {dll.get('base_address', 'N/A')}\n"
if len(dlls) > 10:
result += f"- ... and {len(dlls) - 10} more\n"
else:
result += "- No DLLs found\n"
result += "\n"
# Get network connections
connections = await db.get_network_connections(dump_id, pid)
if connections:
result += f"**Network Connections ({len(connections)}):**\n"
for conn in connections:
result += f"- {conn.get('protocol', 'TCP')} {conn.get('local_addr')}:{conn.get('local_port')} → "
result += f"{conn.get('remote_addr')}:{conn.get('remote_port')} [{conn.get('state', 'N/A')}]\n"
result += "\n"
# Get suspicious memory regions
regions = await db.get_suspicious_memory_regions(dump_id, pid)
if regions:
result += f"**[WARNING] Suspicious Memory Regions ({len(regions)}):**\n"
for region in regions:
result += f"- Base: {region.get('base_address')} | Protection: {region.get('protection')} | "
result += f"Unbacked: {not region.get('is_file_backed')}\n"
result += "\n**This may indicate code injection!**\n"
return [TextContent(type="text", text=result)]
elif name == "detect_code_injection":
dump_id = arguments["dump_id"]
target_pid = arguments.get("pid")
# Get suspicious memory regions
regions = await db.get_suspicious_memory_regions(dump_id, target_pid)
if not regions:
return [TextContent(type="text", text="No code injection detected. All processes appear clean.")]
# Group by PID
by_pid = {}
for region in regions:
pid = region['pid']
if pid not in by_pid:
by_pid[pid] = []
by_pid[pid].append(region)
result = f"**Code Injection Detection Results**\n\n"
result += f"Found suspicious memory regions in {len(by_pid)} process(es):\n\n"
for pid, pid_regions in by_pid.items():
# Get process name
proc = await db.get_process_by_pid(dump_id, pid)
proc_name = proc.get('name', 'Unknown') if proc else 'Unknown'
result += f"**PID {pid} ({proc_name})** - {len(pid_regions)} suspicious region(s)\n"
for region in pid_regions:
result += f" - Base: {region.get('base_address')} | Protection: {region.get('protection')}\n"
result += "\n"
result += "**Recommendation:** Investigate these processes further. Unbacked executable memory often indicates:\n"
result += "- Process hollowing\n"
result += "- Reflective DLL injection\n"
result += "- Shellcode injection\n"
return [TextContent(type="text", text=result)]
elif name == "network_analysis":
dump_id = arguments["dump_id"]
remote_ip = arguments.get("remote_ip")
connections = await db.get_network_connections(dump_id)
if not connections:
return [TextContent(type="text", text="No network connections found.")]
# Filter by IP if specified
if remote_ip:
connections = [c for c in connections if c.get('remote_addr') == remote_ip]
if not connections:
return [TextContent(type="text", text=f"No connections found to {remote_ip}")]
result = f"**Network Analysis**\n\n"
result += f"Total connections: {len(connections)}\n\n"
# Group by process
by_pid = {}
for conn in connections:
pid = conn.get('pid')
if pid not in by_pid:
by_pid[pid] = []
by_pid[pid].append(conn)
for pid, pid_conns in sorted(by_pid.items(), key=lambda x: (x[0] is None, x[0] or 0)):
if pid:
proc = await db.get_process_by_pid(dump_id, pid)
proc_name = proc.get('name', 'Unknown') if proc else 'Unknown'
result += f"**PID {pid} ({proc_name})** - {len(pid_conns)} connection(s)\n"
else:
result += f"**Unknown PID** - {len(pid_conns)} connection(s)\n"
for conn in pid_conns:
result += f" - {conn.get('protocol', 'TCP')} {conn.get('local_addr')}:{conn.get('local_port')} → "
result += f"{conn.get('remote_addr')}:{conn.get('remote_port')} [{conn.get('state', 'N/A')}]\n"
result += "\n"
return [TextContent(type="text", text=result)]
elif name == "detect_hidden_processes":
dump_id = arguments["dump_id"]
vol = get_volatility_handler(dump_id)
if not vol:
return [TextContent(type="text", text=f"Error: Dump '{dump_id}' not found")]
hidden_pids = await vol.detect_hidden_processes()
if not hidden_pids:
return [TextContent(type="text", text="No hidden processes detected. All processes appear in the standard process list.")]
result = f"**[WARNING] Hidden Process Detection**\n\n"
result += f"Found {len(hidden_pids)} hidden process(es):\n\n"
for pid in hidden_pids:
proc = await db.get_process_by_pid(dump_id, pid)
if proc:
result += f"- PID {pid}: {proc.get('name', 'Unknown')}\n"
else:
result += f"- PID {pid}: (process details not available)\n"
result += "\n**Explanation:** These processes appear in memory scans but not in the standard process list, "
result += "which may indicate rootkit activity or process hiding techniques.\n"
return [TextContent(type="text", text=result)]
elif name == "get_process_tree":
dump_id = arguments["dump_id"]
processes = await db.get_processes(dump_id)
if not processes:
return [TextContent(type="text", text=f"No processes found. Run 'process_dump' first.")]
# Build process tree
proc_map = {p['pid']: p for p in processes}
children_map = {}
roots = []
for proc in processes:
ppid = proc.get('ppid')
if ppid and ppid in proc_map:
if ppid not in children_map:
children_map[ppid] = []
children_map[ppid].append(proc['pid'])
else:
roots.append(proc['pid'])
def format_tree(pid, indent=0):
proc = proc_map.get(pid)
if not proc:
return ""
prefix = " " * indent
flags = []
if proc.get('is_hidden'):
flags.append('HIDDEN')
if proc.get('is_suspicious'):
flags.append('SUSP')
flag_str = f" [{', '.join(flags)}]" if flags else ""
line = f"{prefix}├─ {pid}: {proc.get('name', 'N/A')}{flag_str}\n"
# Add children
if pid in children_map:
for child_pid in children_map[pid]:
line += format_tree(child_pid, indent + 1)
return line
result = f"**Process Tree for {dump_id}**\n\n```\n"
for root_pid in sorted(roots):
result += format_tree(root_pid)
result += "```\n"
return [TextContent(type="text", text=result)]
elif name == "get_dump_metadata":
dump_id = arguments["dump_id"]
# Get dump info
dump = await db.get_dump(dump_id)
if not dump:
return [TextContent(type="text", text=f"Dump '{dump_id}' not found")]
result = f"**Memory Dump Metadata**\n\n"
# File information
result += "**File Information:**\n"
result += f"- Dump ID: {dump_id}\n"
result += f"- File Path: {dump.get('file_path', 'N/A')}\n"
result += f"- File Size: {dump.get('file_size', 0) / (1024*1024):.2f} MB\n"
# Get or calculate hashes
dump_path = Path(dump['file_path'])
if dump_path.suffix == '.zip':
# For zip files, find the actual dump
dump_path = extract_dump_if_needed(dump_path)
hashes = await get_or_calculate_hashes(db, dump_id, dump_path)
result += "\n" + format_hashes(hashes)
# OS information
if dump.get('os_type'):
result += f"\n**Operating System:**\n"
result += f"- OS Type: {dump.get('os_type', 'Unknown')}\n"
# Processing status
result += f"\n**Processing Status:**\n"
result += f"- Status: {dump.get('status', 'new')}\n"
result += f"- Last Processed: {dump.get('last_processed', 'Never')}\n"
# Get command statistics
stats = await db.get_command_stats(dump_id)
if stats and stats.get('total_commands', 0) > 0:
result += f"- Commands Executed: {stats.get('total_commands', 0)}\n"
result += f"- Average Execution Time: {int(stats.get('avg_execution_time', 0))} ms\n"
# Count suspicious findings
processes = await db.get_processes(dump_id)
suspicious_procs = sum(1 for p in processes if p.get('is_suspicious'))
regions = await db.get_suspicious_memory_regions(dump_id)
if suspicious_procs > 0 or len(regions) > 0:
result += f"- Suspicious Findings: {suspicious_procs} processes, {len(regions)} memory regions\n"
return [TextContent(type="text", text=result)]
elif name == "export_data":
dump_id = arguments["dump_id"]
export_format = arguments["format"]
data_type = arguments.get("data_type", "all")
output_filename = arguments.get("output_filename")
# Generate filename if not provided
if not output_filename:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
if export_format == "csv":
output_filename = f"{dump_id}_{data_type}_{timestamp}.csv"
elif export_format == "json":
output_filename = f"{dump_id}_{timestamp}.json"
else: # html
output_filename = f"{dump_id}_report_{timestamp}.html"
output_path = EXPORT_DIR / output_filename
# Create exporter
exporter = DataExporter(db)
# Export data
try:
if export_format == "json":
data_types = [data_type] if data_type != "all" else ["processes", "network", "memory_regions"]
stats = await exporter.export_json(dump_id, output_path, data_types=data_types)
elif export_format == "csv":
if data_type == "all":
return [TextContent(type="text", text="Error: CSV format requires a specific data_type (processes, network, or memory_regions)")]
stats = await exporter.export_csv(dump_id, data_type, output_path)
elif export_format == "html":
stats = await exporter.export_html(dump_id, output_path)
else:
return [TextContent(type="text", text=f"Unknown format: {export_format}")]
result = f"**Data Export Complete**\n\n"
result += f"Format: {stats['format']}\n"
result += f"Output: {stats['output_path']}\n"
result += f"Size: {stats.get('file_size', 0) / 1024:.2f} KB\n"
if 'total_records' in stats:
result += f"Records: {stats['total_records']}\n"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Export failed: {e}", exc_info=True)
return [TextContent(type="text", text=f"Export failed: {str(e)}")]
elif name == "get_command_history":
dump_id = arguments["dump_id"]
limit = arguments.get("limit", 50)
# Get provenance summary
summary = await provenance_tracker.get_provenance_summary(dump_id)
return [TextContent(type="text", text=summary)]
elif name == "generate_timeline":
dump_id = arguments["dump_id"]
format_type = arguments.get("format", "summary")
suspicious_only = arguments.get("suspicious_only", False)
output_filename = arguments.get("output_filename")
# Create timeline generator
timeline_gen = TimelineGenerator(db)
if format_type == "summary":
# Display summary
summary = await timeline_gen.get_timeline_summary(dump_id)
return [TextContent(type="text", text=summary)]
else:
# Export to file
if not output_filename:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_filename = f"{dump_id}_timeline_{timestamp}.{format_type}"
output_path = EXPORT_DIR / output_filename
try:
if format_type == "json":
stats = await timeline_gen.export_timeline_json(
dump_id, output_path, suspicious_only=suspicious_only
)
elif format_type == "csv":
stats = await timeline_gen.export_timeline_csv(
dump_id, output_path, suspicious_only=suspicious_only
)
elif format_type == "text":
stats = await timeline_gen.export_timeline_text(
dump_id, output_path, suspicious_only=suspicious_only
)
else:
return [TextContent(type="text", text=f"Unknown format: {format_type}")]
result = f"**Timeline Export Complete**\n\n"
result += f"Format: {stats['format']}\n"
result += f"Output: {stats['output_path']}\n"
result += f"Events: {stats['event_count']}\n"
result += f"Size: {stats.get('file_size', 0) / 1024:.2f} KB\n"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Timeline export failed: {e}", exc_info=True)
return [TextContent(type="text", text=f"Timeline export failed: {str(e)}")]
elif name == "detect_anomalies":
dump_id = arguments["dump_id"]
# Create anomaly detector
detector = AnomalyDetector(db)
# Detect anomalies (returns list of anomaly dicts)
anomalies = await detector.detect_anomalies(dump_id)
# Extract PIDs of suspicious processes and persist to database
suspicious_pids = set()
for anomaly in anomalies:
if 'pid' in anomaly:
suspicious_pids.add(anomaly['pid'])
# For duplicate instances, add all PIDs
if 'pids' in anomaly:
suspicious_pids.update(anomaly['pids'])
# Mark processes as suspicious in the database
if suspicious_pids:
await db.mark_processes_suspicious(dump_id, list(suspicious_pids))
logger.info(f"Marked {len(suspicious_pids)} processes as suspicious in database")
# Get anomaly report
report = await detector.get_anomaly_report(dump_id)
return [TextContent(type="text", text=report)]
elif name == "health_check":
dump_id = arguments["dump_id"]
# Create validator
validator = DataValidator(db)
# Run integrity check
validation_result = await validator.validate_dump_integrity(dump_id)
# Format result
result = f"**Data Integrity Check - {dump_id}**\n\n"
if validation_result['valid']:
result += "Status: HEALTHY\n\n"
else:
result += "Status: ISSUES FOUND\n\n"
# Show statistics
stats = validation_result.get('stats', {})
result += "**Statistics:**\n"
result += f"- Total Commands: {stats.get('total_commands', 0)}\n"
result += f"- Failed Commands: {stats.get('failed_commands', 0)}\n"
result += f"- Process Count: {stats.get('process_count', 0)}\n"
result += f"- Network Connection Count: {stats.get('network_count', 0)}\n\n"
# Show issues
if validation_result.get('issues'):
result += "**Critical Issues:**\n"
for issue in validation_result['issues']:
result += f"- ERROR: {issue}\n"
result += "\n"
# Show warnings
if validation_result.get('warnings'):
result += "**Warnings:**\n"
for warning in validation_result['warnings']:
result += f"- WARNING: {warning}\n"
result += "\n"
if not validation_result.get('issues') and not validation_result.get('warnings'):
result += "No issues or warnings detected. Data appears consistent.\n"
return [TextContent(type="text", text=result)]
elif name == "extract_process":
dump_id = arguments["dump_id"]
pid = arguments["pid"]
# Get volatility handler
vol = get_volatility_handler(dump_id)
if not vol:
return [TextContent(type="text", text=f"Error: Dump '{dump_id}' not found")]
# Create extractor
extractor = await create_extractor(dump_id, vol, db)
try:
# Extract process info
info = await extractor.extract_process_info(pid, EXTRACTED_FILES_DIR)
result = f"**Process Extraction Complete**\n\n"
result += f"Process ID: {info['pid']}\n"
result += f"Output: {info['output_path']}\n"
result += f"Size: {info['file_size'] / 1024:.2f} KB\n"
result += f"SHA256: {info['sha256']}\n\n"
result += f"**Extracted Information:**\n"
result += f"- Loaded DLLs: {info['dll_count']}\n"
result += f"- Network Connections: {info['connection_count']}\n"
result += f"- Suspicious Memory Regions: {info['suspicious_region_count']}\n\n"
result += f"File contains comprehensive JSON with all process details.\n"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Process extraction failed: {e}", exc_info=True)
return [TextContent(type="text", text=f"Process extraction failed: {str(e)}")]
elif name == "cleanup_extractions":
mode = arguments.get("mode", "old")
dry_run = arguments.get("dry_run", False)
try:
if mode == "list":
# List all extractions
extractions = list_extractions(EXTRACTION_DIR)
if not extractions:
return [TextContent(type="text", text="No extractions found.")]
result = f"**Current Extractions**\n\n"
result += f"Found {len(extractions)} extraction(s):\n\n"
for extraction in extractions:
result += f"**{extraction['name']}**\n"
result += f"- Size: {extraction['size_gb']:.2f} GB\n"
result += f"- Created: {extraction['created']}\n"
result += f"- Age: {extraction['age_hours']:.1f} hours\n\n"
return [TextContent(type="text", text=result)]
elif mode == "all":
# Remove all extractions
stats = cleanup_all_extractions(EXTRACTION_DIR)
result = f"**Cleanup Complete**\n\n"
result += f"Removed: {stats['removed_count']} extraction(s)\n"
result += f"Freed: {stats['freed_gb']:.2f} GB\n"
if stats.get('errors'):
result += f"\n**Errors:**\n"
for error in stats['errors']:
result += f"- {error}\n"
return [TextContent(type="text", text=result)]
else: # mode == "old"
# Remove old extractions
stats = cleanup_old_extractions(
EXTRACTION_DIR,
retention_hours=EXTRACTION_RETENTION_HOURS,
dry_run=dry_run
)
if dry_run:
result = f"**Cleanup Preview (Dry Run)**\n\n"
else:
result = f"**Cleanup Complete**\n\n"
result += f"Retention period: {EXTRACTION_RETENTION_HOURS} hours\n"
result += f"Removed: {stats['removed_count']} extraction(s)\n"
result += f"Freed: {stats['freed_gb']:.2f} GB\n"
if dry_run and stats['removed_count'] > 0:
result += f"\nRun with dry_run=false to actually remove these files.\n"
if stats.get('errors'):
result += f"\n**Errors:**\n"
for error in stats['errors']:
result += f"- {error}\n"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Cleanup failed: {e}", exc_info=True)
return [TextContent(type="text", text=f"Cleanup failed: {str(e)}")]
elif name == "get_disk_usage":
try:
stats = get_disk_usage(DATA_DIR, EXTRACTION_DIR)
result = f"**Disk Usage Statistics**\n\n"
result += f"**Database:**\n"
result += f"- Size: {stats['database_size_mb']:.2f} MB\n"
result += f"- Location: {stats['data_dir']}/artifacts.db\n\n"
result += f"**Exports:**\n"
result += f"- Size: {stats['exports_size_mb']:.2f} MB\n"
result += f"- Location: {stats['data_dir']}/exports\n\n"
result += f"**Extractions:**\n"
result += f"- Size: {stats['extractions_size_gb']:.2f} GB\n"
result += f"- Count: {stats['extractions_count']}\n"
result += f"- Location: {stats['extraction_dir']}\n\n"
result += f"**Total:**\n"
result += f"- Size: {stats['total_size_gb']:.2f} GB\n"
# Add recommendation if extractions are large
if stats['extractions_size_gb'] > 5.0:
result += f"\n**Recommendation:** Extractions are using {stats['extractions_size_gb']:.2f} GB. "
result += f"Consider running cleanup_extractions to free space.\n"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Failed to get disk usage: {e}", exc_info=True)
return [TextContent(type="text", text=f"Failed to get disk usage: {str(e)}")]
else:
return [TextContent(type="text", text=f"Unknown tool: {name}")]
except Exception as e:
logger.error(f"Error executing tool {name}: {e}", exc_info=True)
return [TextContent(type="text", text=f"Error: {str(e)}")]
async def main():
"""Main entry point for the MCP server"""
# Initialize database
await db.initialize()
logger.info("Memory Forensics MCP Server starting...")
logger.info(f"Database: {DB_PATH}")
logger.info(f"Dumps directory: {DUMPS_DIR}")
# Clean up old extractions on startup
if AUTO_CLEANUP_ON_STARTUP:
try:
stats = cleanup_old_extractions(
EXTRACTION_DIR,
retention_hours=EXTRACTION_RETENTION_HOURS,
dry_run=False
)
if stats['removed_count'] > 0:
logger.info(
f"Startup cleanup: Removed {stats['removed_count']} old extraction(s), "
f"freed {stats['freed_gb']:.2f} GB"
)
except Exception as e:
logger.warning(f"Startup cleanup failed: {e}")
# Run the server
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
InitializationOptions(
server_name="memory-forensics",
server_version="0.1.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={}
)
)
)
if __name__ == "__main__":
asyncio.run(main())
