ffmpeg_server.py

from mcp.server.fastmcp import FastMCP
import subprocess
import os
import glob

# Check if FastMCP is properly imported
print("FastMCP imported from:", FastMCP.__module__)

# Initialize the MCP server
try:
    mcp = FastMCP(
        name="FFmpegServer",
        description="A server for multimedia processing with FFmpeg",
        version="0.1.0"
    )
except Exception as e:
    print(f"Error initializing FastMCP: {e}")
    exit(1)

# Define working directory for security
WORKING_DIR = os.path.expanduser("~/ffmpeg_mcp_files")
os.makedirs(WORKING_DIR, exist_ok=True)

def is_safe_path(base_path: str, path: str) -> bool:
    """Ensure path stays within working directory"""
    abs_base = os.path.abspath(base_path)
    abs_path = os.path.abspath(path)
    return abs_path.startswith(abs_base)

def run_ffmpeg_command(command: list) -> dict:
    """Run FFmpeg command and return structured response"""
    try:
        # Get ffmpeg path from which command
        ffmpeg_path = subprocess.check_output(['which', 'ffmpeg']).decode().strip()
        # Replace the first element (ffmpeg) with the full path
        command[0] = ffmpeg_path
        
        result = subprocess.run(
            command,
            check=True,
            capture_output=True,
            text=True
        )
        return {
            "success": True,
            "output": result.stdout.strip(),
            "error": None
        }
    except subprocess.CalledProcessError as e:
        return {
            "success": False,
            "output": None,
            "error": e.stderr.strip()
        }
    except Exception as e:
        return {
            "success": False,
            "output": None,
            "error": str(e)
        }

@mcp.tool()
def trim_video(input_path: str, start_time: float, duration: float, output_path: str) -> dict:
    """
    Trim a video to a specified start time and duration.
    """
    full_input_path = os.path.join(WORKING_DIR, input_path)
    full_output_path = os.path.join(WORKING_DIR, output_path)

    if not is_safe_path(WORKING_DIR, full_input_path) or not is_safe_path(WORKING_DIR, full_output_path):
        return {"success": False, "output": None, "error": "Path outside working directory not allowed"}
    
    if not os.path.exists(full_input_path):
        return {"success": False, "output": None, "error": f"Input file {input_path} does not exist"}
    
    if start_time < 0 or duration <= 0:
        return {"success": False, "output": None, "error": "Invalid time parameters"}

    command = [
        "ffmpeg",
        "-i", full_input_path,
        "-ss", str(start_time),
        "-t", str(duration),
        "-c:v", "copy",
        "-c:a", "copy",
        "-y",
        full_output_path
    ]
    result = run_ffmpeg_command(command)
    if result["success"]:
        result["output"] = f"Video trimmed to {full_output_path}"
    return result

@mcp.tool()
def convert_video_format(input_path: str, output_format: str, output_path: str) -> dict:
    """
    Convert a video to a specified format.
    """
    full_input_path = os.path.join(WORKING_DIR, input_path)
    full_output_path = os.path.join(WORKING_DIR, output_path)
    output_format = output_format.lower().strip()

    if not is_safe_path(WORKING_DIR, full_input_path) or not is_safe_path(WORKING_DIR, full_output_path):
        return {"success": False, "output": None, "error": "Path outside working directory not allowed"}
    
    if not os.path.exists(full_input_path):
        return {"success": False, "output": None, "error": f"Input file {input_path} does not exist"}
    
    if not full_output_path.endswith(f".{output_format}"):
        full_output_path = f"{full_output_path}.{output_format}"

    command = [
        "ffmpeg",
        "-i", full_input_path,
        "-f", output_format,
        "-y",
        full_output_path
    ]
    result = run_ffmpeg_command(command)
    if result["success"]:
        result["output"] = f"Video converted to {full_output_path}"
    return result

@mcp.tool()
def extract_audio(input_path: str, audio_format: str, output_path: str) -> dict:
    """
    Extract audio from a video file.
    """
    full_input_path = os.path.join(WORKING_DIR, input_path)
    full_output_path = os.path.join(WORKING_DIR, output_path)
    audio_format = audio_format.lower().strip()
    
    if not is_safe_path(WORKING_DIR, full_input_path) or not is_safe_path(WORKING_DIR, full_output_path):
        return {"success": False, "output": None, "error": "Path outside working directory not allowed"}
    
    if not os.path.exists(full_input_path):
        return {"success": False, "output": None, "error": f"Input file {input_path} does not exist"}
    
    if not full_output_path.endswith(f".{audio_format}"):
        full_output_path = f"{full_output_path}.{audio_format}"
    
    command = [
        "ffmpeg",
        "-i", full_input_path,
        "-vn",  # No video
        "-acodec", "copy" if audio_format in ["aac", "mp3", "ogg"] else audio_format,
        "-y",
        full_output_path
    ]
    
    result = run_ffmpeg_command(command)
    if result["success"]:
        result["output"] = f"Audio extracted to {full_output_path}"
    return result

@mcp.tool()
def add_watermark(input_path: str, watermark_path: str, position: str, output_path: str) -> dict:
    """
    Add a watermark image to a video.
    Position can be: "topleft", "topright", "bottomleft", "bottomright", "center"
    """
    full_input_path = os.path.join(WORKING_DIR, input_path)
    full_watermark_path = os.path.join(WORKING_DIR, watermark_path)
    full_output_path = os.path.join(WORKING_DIR, output_path)
    
    if not is_safe_path(WORKING_DIR, full_input_path) or not is_safe_path(WORKING_DIR, full_output_path) or not is_safe_path(WORKING_DIR, full_watermark_path):
        return {"success": False, "output": None, "error": "Path outside working directory not allowed"}
    
    if not os.path.exists(full_input_path):
        return {"success": False, "output": None, "error": f"Input file {input_path} does not exist"}
    
    if not os.path.exists(full_watermark_path):
        return {"success": False, "output": None, "error": f"Watermark file {watermark_path} does not exist"}
    
    # Define overlay position
    position_map = {
        "topleft": "10:10",
        "topright": "(main_w-overlay_w-10):10",
        "bottomleft": "10:(main_h-overlay_h-10)",
        "bottomright": "(main_w-overlay_w-10):(main_h-overlay_h-10)",
        "center": "(main_w-overlay_w)/2:(main_h-overlay_h)/2"
    }
    
    if position not in position_map:
        position = "bottomright"  # Default position
    
    overlay_position = position_map[position]
    
    command = [
        "ffmpeg",
        "-i", full_input_path,
        "-i", full_watermark_path,
        "-filter_complex", f"[0:v][1:v] overlay={overlay_position}:enable='between(t,0,999999)'",
        "-codec:a", "copy",
        "-y",
        full_output_path
    ]
    
    result = run_ffmpeg_command(command)
    if result["success"]:
        result["output"] = f"Video with watermark saved to {full_output_path}"
    return result

@mcp.tool()
def adjust_video_quality(input_path: str, quality: str, output_path: str) -> dict:
    """
    Adjust video quality. Quality can be: "low", "medium", "high", "veryhigh"
    """
    full_input_path = os.path.join(WORKING_DIR, input_path)
    full_output_path = os.path.join(WORKING_DIR, output_path)
    
    if not is_safe_path(WORKING_DIR, full_input_path) or not is_safe_path(WORKING_DIR, full_output_path):
        return {"success": False, "output": None, "error": "Path outside working directory not allowed"}
    
    if not os.path.exists(full_input_path):
        return {"success": False, "output": None, "error": f"Input file {input_path} does not exist"}
    
    # Define quality presets
    quality_presets = {
        "low": ["-crf", "28", "-preset", "veryfast"],
        "medium": ["-crf", "23", "-preset", "medium"],
        "high": ["-crf", "18", "-preset", "slow"],
        "veryhigh": ["-crf", "15", "-preset", "veryslow"]
    }
    
    if quality not in quality_presets:
        quality = "medium"  # Default quality
    
    command = [
        "ffmpeg",
        "-i", full_input_path,
        *quality_presets[quality],
        "-c:a", "copy",
        "-y",
        full_output_path
    ]
    
    result = run_ffmpeg_command(command)
    if result["success"]:
        result["output"] = f"Video with adjusted quality saved to {full_output_path}"
    return result

@mcp.tool()
def batch_process(input_pattern: str, operation: str, params: dict) -> dict:
    """
    Process multiple files matching a pattern.
    Operation can be: "trim", "convert", "extract_audio", "watermark", "quality"
    """
    full_pattern = os.path.join(WORKING_DIR, input_pattern)
    
    # Ensure pattern is within working directory
    if not is_safe_path(WORKING_DIR, os.path.dirname(full_pattern)):
        return {"success": False, "output": None, "error": "Pattern outside working directory not allowed"}
    
    # Find matching files
    matching_files = glob.glob(full_pattern)
    if not matching_files:
        return {"success": False, "output": None, "error": f"No files match pattern {input_pattern}"}
    
    results = []
    operations = {
        "trim": trim_video,
        "convert": convert_video_format,
        "extract_audio": extract_audio,
        "watermark": add_watermark,
        "quality": adjust_video_quality
    }
    
    if operation not in operations:
        return {"success": False, "output": None, "error": f"Unknown operation: {operation}"}
    
    for file_path in matching_files:
        # Get relative path for function calls
        rel_path = os.path.relpath(file_path, WORKING_DIR)
        
        # Prepare output path
        filename = os.path.basename(file_path)
        name, ext = os.path.splitext(filename)
        output_filename = f"{name}_processed{ext}"
        
        # Update params with input and output paths
        current_params = params.copy()
        current_params["input_path"] = rel_path
        
        if "output_path" not in current_params:
            current_params["output_path"] = output_filename
        
        # Call the appropriate function
        result = operations[operation](**current_params)
        results.append({
            "file": rel_path,
            "result": result
        })
    
    return {
        "success": True,
        "output": f"Processed {len(results)} files",
        "details": results
    }

@mcp.tool()
def list_media_files() -> dict:
    """
    List all media files in the working directory.
    """
    try:
        # Common media file extensions
        media_extensions = [".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm", 
                           ".mp3", ".wav", ".ogg", ".aac", ".flac", ".jpg", ".png", ".gif"]
        
        all_files = []
        for ext in media_extensions:
            pattern = os.path.join(WORKING_DIR, f"*{ext}")
            all_files.extend(glob.glob(pattern))
        
        # Convert to relative paths
        relative_files = [os.path.relpath(f, WORKING_DIR) for f in all_files]
        
        return {
            "success": True,
            "output": f"Found {len(relative_files)} media files",
            "files": relative_files
        }
    except Exception as e:
        return {
            "success": False,
            "output": None,
            "error": str(e)
        }

if __name__ == "__main__":
    print(f"Starting FFmpegServer - Working directory: {WORKING_DIR}")
    try:
        # Try different possible method names
        if hasattr(mcp, 'run'):
            mcp.run()
        elif hasattr(mcp, 'start'):
            print("Using start() instead of run()")
            mcp.start()
        elif hasattr(mcp, 'serve'):
            print("Using serve() instead of run()")
            mcp.serve()
        else:
            print("Available methods:", dir(mcp))
            raise AttributeError("No known start method found for FastMCP")
    except Exception as e:
        print(f"Error starting server: {e}")
        print("Server object methods:", dir(mcp))