pip_audit_mcp.py

#!/usr/bin/env python3
"""
MCP server for pip-audit - a tool for scanning Python environments for known vulnerabilities
"""

import subprocess
import json
from typing import Dict
import gradio as gr
import os

def pip_audit_scan() -> Dict:
    """
    Scans Python environments for known vulnerabilities using pip-audit with basic settings.
    
    Returns:
        Dict: Scan results
    """
    try:
        cmd = ["pip-audit", "--format", "json"]

        print(f"Executing command: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        stdout, stderr = result.stdout, result.stderr
        return_code = result.returncode

        if return_code != 0:
            print(f"pip-audit command failed with return code {return_code}")
            print(f"Stderr: {stderr}")
            return {
                "success": False,
                "error": f"pip-audit command failed with return code {return_code}",
                "stdout": stdout,
                "stderr": stderr,
                "return_code": return_code
            }

        try:
            output_data = json.loads(stdout) if stdout else {}
            return {
                "success": True,
                "results": output_data,
                "stderr": stderr,
                "return_code": return_code
            }
        except json.JSONDecodeError as e:
            print(f"JSON parsing error: {e}")
            print(f"Raw stdout: {stdout}")
            return {
                "success": False,
                "error": "JSON parsing error: " + str(e),
                "stdout": stdout,
                "stderr": stderr,
                "return_code": return_code
            }
            
    except Exception as e:
        print(f"Error executing pip-audit: {str(e)}")
        return {
            "success": False,
            "error": f"Error executing pip-audit: {str(e)}"
        }

# Create Gradio interface
with gr.Blocks(title="Pip Audit MCP") as demo:
    gr.Markdown("# 🛡️ Pip Audit Scanner")
    gr.Markdown("Vulnerability scanning tool for Python environments with MCP support")

    with gr.Tab("Basic Scanning"):
        scan_btn = gr.Button("🔍 Run Basic Audit", variant="primary")
        scan_output = gr.JSON(label="Audit Results")

        scan_btn.click(
            fn=pip_audit_scan,
            inputs=[],
            outputs=scan_output
        )

if __name__ == "__main__":
    # Получаем настройки сервера из переменных окружения
    server_name = os.getenv("GRADIO_SERVER_NAME", "0.0.0.0")
    server_port = int(os.getenv("GRADIO_SERVER_PORT", "7863"))
    
    demo.launch(
        mcp_server=True,
        server_name=server_name,
        server_port=server_port,
        share=False
    )