Agents-MCP-Hackathon/PwnGuard

1"""2PwnGuard MCP Handler Module3 4This module contains functions for handling Model Context Protocol (MCP) requests5and formatting responses according to the MCP standard.6"""7 8import json9import logging10from datetime import datetime11from typing import Dict, Any12 13from security_analyzer import SecurityAnalyzer14 15logger = logging.getLogger(__name__)16 17# Global security analyzer instance18security_analyzer = SecurityAnalyzer()19 20def analyze_mcp_request(task_id: str, user_query: str, context: str = "", history_json: str = "[]") -> str:21    """MCP-compatible security analysis function"""22    try:23        # Parse history records24        history = json.loads(history_json) if history_json else []25        26        # Execute security analysis27        analysis_result = security_analyzer.analyze_security_risk(user_query, context, history)28        29        # Generate safety response30        safety_response = security_analyzer.generate_safety_response(analysis_result, user_query)31        32        # Construct MCP response format33        mcp_response = {34            "task_id": task_id,35            "status": "success",36            "result": {37                "pwnguard_analysis": safety_response,38                "original_query": user_query,39                "processing_time": datetime.now().isoformat()40            }41        }42        43        logger.info(f"Processed MCP request: {task_id}, risk level: {analysis_result['risk_level']}")44        45        return json.dumps(mcp_response, ensure_ascii=False, indent=2)46        47    except Exception as e:48        logger.error(f"Error processing MCP request: {str(e)}")49        error_response = {50            "task_id": task_id,51            "status": "error",52            "error": str(e)53        }54        return json.dumps(error_response, ensure_ascii=False, indent=2)55 56def simple_security_check(user_query: str, use_llm_analysis: bool = False) -> str:57    """Simple security check function (for Gradio interface) - returns JSON format"""58    try:59        analysis_result = security_analyzer.analyze_security_risk(user_query, use_llm=use_llm_analysis)60        61        # Count detected risks62        detected_count = sum(1 for detected in analysis_result['detected_risks'].values() if detected)63        detected_risk_types = [risk_type for risk_type, detected in analysis_result['detected_risks'].items() if detected]64        65        # Create simplified JSON response66        simple_result = {67            "status": "success",68            "result": {69                "risk_level": analysis_result['risk_level'],70                "risk_score": f"{analysis_result['risk_score']}/{'15' if use_llm_analysis else '10'}",71                "is_safe": analysis_result['is_safe'],72                "detected_risks": detected_risk_types,73                "total_risks": detected_count,74                "analysis_mode": "LLM-Enhanced" if use_llm_analysis else "Regex-Based",75                "recommended_action": "proceed" if analysis_result['is_safe'] else "block"76            }77        }78        79        return json.dumps(simple_result, ensure_ascii=False, indent=2)80        81    except Exception as e:82        error_result = {83            "status": "error", 84            "error": str(e)85        }86        return json.dumps(error_result, ensure_ascii=False, indent=2)87 88def quick_check(user_query: str, use_llm_analysis: bool = False) -> str:89    """Quick security check with minimal JSON output"""90    try:91        analysis_result = security_analyzer.analyze_security_risk(user_query, use_llm=use_llm_analysis)92        93        # Ultra-simplified output without emojis94        quick_result = {95            "risk_level": analysis_result['risk_level'].upper(),96            "risk_score": f"{analysis_result['risk_score']}/{'15' if use_llm_analysis else '10'}",97            "is_safe": analysis_result['is_safe'],98            "detected_risks": {99                risk_type: detected 100                for risk_type, detected in analysis_result['detected_risks'].items()101            },102            "recommended_action": "proceed" if analysis_result['is_safe'] else "block"103        }104        105        return json.dumps(quick_result, ensure_ascii=False, indent=2)106        107    except Exception as e:108        return json.dumps({"error": str(e)}, ensure_ascii=False, indent=2)