#!/usr/bin/env python3
"""
PQC Readiness Analyzer MCP Server

This MCP server analyzes Git commits for cryptographic algorithm usage
and calculates PQC (Post-Quantum Cryptography) readiness scores.
"""

import asyncio
import json
import re
import subprocess
import sys
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict

# MCP SDK imports
try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent
except ImportError:
    print("Error: MCP SDK not installed. Run: pip install mcp", file=sys.stderr)
    sys.exit(1)


@dataclass
class AlgorithmDetection:
    """Represents a detected cryptographic algorithm"""
    algorithm: str
    algorithm_type: str  # RSA, ECDSA, DILITHIUM, FALCON, HYBRID
    family: str  # CLASSICAL, PQC, HYBRID
    pqc_safe: bool
    file_path: str
    line_number: int
    context: str


@dataclass
class PQCReadinessScore:
    """PQC Readiness analysis result"""
    commit_id: str
    total_algorithms: int
    pqc_safe_count: int
    pqc_unsafe_count: int
    classical_count: int
    hybrid_count: int
    pqc_readiness_percentage: float
    quantum_risk_level: str  # LOW, MEDIUM, HIGH, CRITICAL
    detected_algorithms: List[Dict[str, Any]]
    recommendations: List[str]


class CryptoAlgorithmDetector:
    """Detects cryptographic algorithms in code"""
    
    # Algorithm patterns for detection
    ALGORITHM_PATTERNS = {
        # Classical algorithms
        'RSA': {
            'patterns': [
                r'RSA[_-]?3072',
                r'RSA[_-]?2048',
                r'RSA[_-]?4096',
                r'KeyPairGenerator\.getInstance\(["\']RSA["\']',
                r'RSAKeyPairGenerator',
                r'AlgorithmType\.RSA',
            ],
            'family': 'CLASSICAL',
            'pqc_safe': False
        },
        'ECDSA': {
            'patterns': [
                r'ECDSA',
                r'EC[_-]?P[_-]?256',
                r'secp256r1',
                r'prime256v1',
                r'ECDSAKeyPairGenerator',
                r'AlgorithmType\.ECDSA',
            ],
            'family': 'CLASSICAL',
            'pqc_safe': False
        },
        # PQC algorithms
        'DILITHIUM': {
            'patterns': [
                r'Dilithium[23]?',
                r'DILITHIUM[23]?',
                r'DilithiumKeyPairGenerator',
                r'AlgorithmType\.DILITHIUM',
            ],
            'family': 'PQC',
            'pqc_safe': True
        },
        'FALCON': {
            'patterns': [
                r'Falcon[_-]?512',
                r'Falcon[_-]?1024',
                r'FALCON[_-]?512',
                r'FalconKeyPairGenerator',
                r'AlgorithmType\.FALCON',
            ],
            'family': 'PQC',
            'pqc_safe': True
        },
        'SPHINCS': {
            'patterns': [
                r'SPHINCS\+',
                r'SPHINCSPlus',
            ],
            'family': 'PQC',
            'pqc_safe': True
        },
        # Hybrid algorithms
        'HYBRID': {
            'patterns': [
                r'HYBRID[_-]?RSA[_-]?DILITHIUM',
                r'HybridKeyPairGenerator',
                r'AlgorithmType\.HYBRID',
            ],
            'family': 'HYBRID',
            'pqc_safe': True
        },
        # Hash algorithms (for context)
        'SHA256': {
            'patterns': [
                r'SHA[_-]?256',
                r'SHA256withRSA',
                r'SHA256withECDSA',
            ],
            'family': 'HASH',
            'pqc_safe': True  # SHA-256 is quantum-resistant for hashing
        },
        'SHA3': {
            'patterns': [
                r'SHA3[_-]?256',
                r'SHA3[_-]?512',
            ],
            'family': 'HASH',
            'pqc_safe': True
        }
    }
    
    def detect_algorithms(self, content: str, file_path: str) -> List[AlgorithmDetection]:
        """Detect cryptographic algorithms in code content"""
        detections = []
        lines = content.split('\n')
        
        for line_num, line in enumerate(lines, 1):
            for algo_name, algo_info in self.ALGORITHM_PATTERNS.items():
                for pattern in algo_info['patterns']:
                    if re.search(pattern, line, re.IGNORECASE):
                        detection = AlgorithmDetection(
                            algorithm=algo_name,
                            algorithm_type=algo_name,
                            family=algo_info['family'],
                            pqc_safe=algo_info['pqc_safe'],
                            file_path=file_path,
                            line_number=line_num,
                            context=line.strip()
                        )
                        detections.append(detection)
                        break  # Only detect once per line per algorithm
        
        return detections


class GitCommitAnalyzer:
    """Analyzes Git commits for cryptographic changes"""
    
    def __init__(self, repo_path: str):
        self.repo_path = repo_path
        self.detector = CryptoAlgorithmDetector()
    
    def get_commit_diff(self, commit_id: str) -> str:
        """Get the diff for a specific commit"""
        try:
            result = subprocess.run(
                ['git', 'show', commit_id],
                cwd=self.repo_path,
                capture_output=True,
                text=True,
                check=True
            )
            return result.stdout
        except subprocess.CalledProcessError as e:
            raise Exception(f"Failed to get commit diff: {e.stderr}")
    
    def get_commit_files(self, commit_id: str) -> List[str]:
        """Get list of files changed in a commit"""
        try:
            result = subprocess.run(
                ['git', 'diff-tree', '--no-commit-id', '--name-only', '-r', commit_id],
                cwd=self.repo_path,
                capture_output=True,
                text=True,
                check=True
            )
            return [f.strip() for f in result.stdout.split('\n') if f.strip()]
        except subprocess.CalledProcessError as e:
            raise Exception(f"Failed to get commit files: {e.stderr}")
    
    def get_file_content_at_commit(self, commit_id: str, file_path: str) -> str:
        """Get file content at a specific commit"""
        try:
            result = subprocess.run(
                ['git', 'show', f'{commit_id}:{file_path}'],
                cwd=self.repo_path,
                capture_output=True,
                text=True,
                check=True
            )
            return result.stdout
        except subprocess.CalledProcessError:
            return ""
    
    def analyze_commit(self, commit_id: str) -> PQCReadinessScore:
        """Analyze a commit for PQC readiness"""
        # Get changed files
        changed_files = self.get_commit_files(commit_id)
        
        # Filter for relevant files (Java, Python, etc.)
        relevant_extensions = ['.java', '.py', '.js', '.ts', '.go', '.cpp', '.c', '.h']
        relevant_files = [
            f for f in changed_files 
            if any(f.endswith(ext) for ext in relevant_extensions)
        ]
        
        # Detect algorithms in changed files
        all_detections = []
        for file_path in relevant_files:
            content = self.get_file_content_at_commit(commit_id, file_path)
            if content:
                detections = self.detector.detect_algorithms(content, file_path)
                all_detections.extend(detections)
        
        # Calculate statistics
        total_algorithms = len(all_detections)
        pqc_safe_count = sum(1 for d in all_detections if d.pqc_safe)
        pqc_unsafe_count = total_algorithms - pqc_safe_count
        
        classical_count = sum(1 for d in all_detections if d.family == 'CLASSICAL')
        hybrid_count = sum(1 for d in all_detections if d.family == 'HYBRID')
        pqc_count = sum(1 for d in all_detections if d.family == 'PQC')
        
        # Calculate PQC readiness percentage
        if total_algorithms > 0:
            pqc_readiness_percentage = (pqc_safe_count / total_algorithms) * 100
        else:
            pqc_readiness_percentage = 0.0
        
        # Determine quantum risk level
        if pqc_readiness_percentage >= 90:
            quantum_risk_level = "LOW"
        elif pqc_readiness_percentage >= 70:
            quantum_risk_level = "MEDIUM"
        elif pqc_readiness_percentage >= 40:
            quantum_risk_level = "HIGH"
        else:
            quantum_risk_level = "CRITICAL"
        
        # Generate recommendations
        recommendations = self._generate_recommendations(
            classical_count, pqc_count, hybrid_count, pqc_readiness_percentage
        )
        
        return PQCReadinessScore(
            commit_id=commit_id,
            total_algorithms=total_algorithms,
            pqc_safe_count=pqc_safe_count,
            pqc_unsafe_count=pqc_unsafe_count,
            classical_count=classical_count,
            hybrid_count=hybrid_count,
            pqc_readiness_percentage=round(pqc_readiness_percentage, 2),
            quantum_risk_level=quantum_risk_level,
            detected_algorithms=[asdict(d) for d in all_detections],
            recommendations=recommendations
        )
    
    def _generate_recommendations(
        self, classical_count: int, pqc_count: int, 
        hybrid_count: int, readiness_pct: float
    ) -> List[str]:
        """Generate recommendations based on analysis"""
        recommendations = []
        
        if classical_count > 0 and pqc_count == 0:
            recommendations.append(
                "⚠️ CRITICAL: No PQC algorithms detected. Consider migrating to "
                "Dilithium3 or Falcon-512 for quantum resistance."
            )
        
        if readiness_pct < 50:
            recommendations.append(
                "🔴 HIGH PRIORITY: PQC readiness is below 50%. Immediate action "
                "required to prepare for quantum threats."
            )
        elif readiness_pct < 80:
            recommendations.append(
                "🟡 MEDIUM PRIORITY: PQC readiness is moderate. Plan migration "
                "strategy for remaining classical algorithms."
            )
        
        if classical_count > 0:
            recommendations.append(
                f"📊 Found {classical_count} classical algorithm(s). Consider "
                "hybrid approach (RSA+Dilithium) for gradual migration."
            )
        
        if hybrid_count > 0:
            recommendations.append(
                f"✅ Good: {hybrid_count} hybrid algorithm(s) detected. This provides "
                "a balanced transition path to full PQC."
            )
        
        if pqc_count > 0:
            recommendations.append(
                f"✅ Excellent: {pqc_count} PQC algorithm(s) detected. Continue "
                "adopting NIST-standardized algorithms."
            )
        
        if readiness_pct >= 90:
            recommendations.append(
                "🎉 Outstanding: PQC readiness is excellent! Maintain this level "
                "and monitor for new NIST standards."
            )
        
        return recommendations


# Initialize MCP server
app = Server("pqc-analyzer")


@app.list_tools()
async def list_tools() -> list[Tool]:
    """List available tools"""
    return [
        Tool(
            name="analyze_commit_pqc_readiness",
            description=(
                "Analyze a Git commit for PQC (Post-Quantum Cryptography) readiness. "
                "Detects cryptographic algorithms used in the commit and calculates "
                "a PQC readiness score, including counts of PQC-safe vs unsafe algorithms."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "commit_id": {
                        "type": "string",
                        "description": "Git commit ID (SHA) to analyze"
                    },
                    "repo_path": {
                        "type": "string",
                        "description": "Path to the Git repository (optional, defaults to current directory)",
                        "default": "."
                    }
                },
                "required": ["commit_id"]
            }
        ),
        Tool(
            name="get_pqc_algorithm_info",
            description=(
                "Get detailed information about supported PQC algorithms, "
                "including their security levels, key sizes, and NIST standardization status."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "algorithm": {
                        "type": "string",
                        "description": "Algorithm name (optional, returns all if not specified)",
                        "enum": ["DILITHIUM", "FALCON", "SPHINCS", "RSA", "ECDSA", "HYBRID"]
                    }
                }
            }
        )
    ]


@app.call_tool()
async def call_tool(name: str, arguments: Any) -> list[TextContent]:
    """Handle tool calls"""
    
    if name == "analyze_commit_pqc_readiness":
        commit_id = arguments.get("commit_id")
        repo_path = arguments.get("repo_path", ".")
        
        if not commit_id:
            return [TextContent(
                type="text",
                text="Error: commit_id is required"
            )]
        
        try:
            analyzer = GitCommitAnalyzer(repo_path)
            result = analyzer.analyze_commit(commit_id)
            
            # Format the result as a detailed report
            report = f"""
# PQC Readiness Analysis Report

## Commit Information
- **Commit ID**: {result.commit_id}
- **Repository**: {repo_path}

## Summary Statistics
- **Total Algorithms Detected**: {result.total_algorithms}
- **PQC-Safe Algorithms**: {result.pqc_safe_count}
- **PQC-Unsafe Algorithms**: {result.pqc_unsafe_count}
- **Classical Algorithms**: {result.classical_count}
- **Hybrid Algorithms**: {result.hybrid_count}

## PQC Readiness Score
- **Score**: {result.pqc_readiness_percentage}%
- **Quantum Risk Level**: {result.quantum_risk_level}

## Detected Algorithms
"""
            if result.detected_algorithms:
                for i, algo in enumerate(result.detected_algorithms, 1):
                    report += f"\n### {i}. {algo['algorithm']} ({algo['family']})\n"
                    report += f"- **File**: `{algo['file_path']}`\n"
                    report += f"- **Line**: {algo['line_number']}\n"
                    report += f"- **PQC Safe**: {'✅ Yes' if algo['pqc_safe'] else '❌ No'}\n"
                    report += f"- **Context**: `{algo['context']}`\n"
            else:
                report += "\nNo cryptographic algorithms detected in this commit.\n"
            
            report += "\n## Recommendations\n"
            for rec in result.recommendations:
                report += f"\n{rec}\n"
            
            return [TextContent(type="text", text=report)]
            
        except Exception as e:
            return [TextContent(
                type="text",
                text=f"Error analyzing commit: {str(e)}"
            )]
    
    elif name == "get_pqc_algorithm_info":
        algorithm = arguments.get("algorithm")
        
        algorithm_info = {
            "DILITHIUM": {
                "name": "Dilithium",
                "type": "Digital Signature",
                "family": "PQC",
                "nist_status": "NIST Standardized (FIPS 204)",
                "security_level": "NIST Level 2, 3, 5",
                "key_sizes": "Dilithium2 (2420 bytes), Dilithium3 (3293 bytes), Dilithium5 (4595 bytes)",
                "signature_size": "2420-4595 bytes",
                "quantum_safe": True,
                "recommended": True,
                "use_cases": ["TLS certificates", "Code signing", "Document signing"]
            },
            "FALCON": {
                "name": "Falcon",
                "type": "Digital Signature",
                "family": "PQC",
                "nist_status": "NIST Standardized (FIPS 206)",
                "security_level": "NIST Level 1, 5",
                "key_sizes": "Falcon-512 (897 bytes), Falcon-1024 (1793 bytes)",
                "signature_size": "666-1280 bytes",
                "quantum_safe": True,
                "recommended": True,
                "use_cases": ["Constrained environments", "IoT devices", "Embedded systems"]
            },
            "SPHINCS": {
                "name": "SPHINCS+",
                "type": "Digital Signature",
                "family": "PQC",
                "nist_status": "NIST Standardized (FIPS 205)",
                "security_level": "NIST Level 1, 3, 5",
                "key_sizes": "32-64 bytes (public key)",
                "signature_size": "7856-49856 bytes",
                "quantum_safe": True,
                "recommended": True,
                "use_cases": ["Long-term security", "Stateless signatures", "Backup signatures"]
            },
            "RSA": {
                "name": "RSA",
                "type": "Digital Signature / Encryption",
                "family": "CLASSICAL",
                "nist_status": "Classical Standard",
                "security_level": "Classical 112-128 bits",
                "key_sizes": "2048-4096 bits",
                "signature_size": "256-512 bytes",
                "quantum_safe": False,
                "recommended": False,
                "use_cases": ["Legacy systems", "Hybrid transition"],
                "quantum_risk": "Vulnerable to Shor's algorithm"
            },
            "ECDSA": {
                "name": "ECDSA",
                "type": "Digital Signature",
                "family": "CLASSICAL",
                "nist_status": "Classical Standard",
                "security_level": "Classical 128 bits (P-256)",
                "key_sizes": "256-521 bits",
                "signature_size": "64-132 bytes",
                "quantum_safe": False,
                "recommended": False,
                "use_cases": ["Legacy systems", "Hybrid transition"],
                "quantum_risk": "Vulnerable to Shor's algorithm"
            },
            "HYBRID": {
                "name": "Hybrid (RSA+Dilithium)",
                "type": "Digital Signature",
                "family": "HYBRID",
                "nist_status": "Experimental",
                "security_level": "Transitional",
                "key_sizes": "Combined classical + PQC",
                "signature_size": "Combined",
                "quantum_safe": True,
                "recommended": True,
                "use_cases": ["Migration path", "Dual security", "Compatibility bridge"]
            }
        }
        
        if algorithm:
            if algorithm in algorithm_info:
                info = algorithm_info[algorithm]
                report = f"""
# {info['name']} Algorithm Information

## Basic Information
- **Type**: {info['type']}
- **Family**: {info['family']}
- **NIST Status**: {info['nist_status']}
- **Security Level**: {info['security_level']}

## Technical Details
- **Key Sizes**: {info['key_sizes']}
- **Signature Size**: {info['signature_size']}
- **Quantum Safe**: {'✅ Yes' if info['quantum_safe'] else '❌ No'}
- **Recommended**: {'✅ Yes' if info['recommended'] else '⚠️ No (Legacy)'}

## Use Cases
"""
                for use_case in info['use_cases']:
                    report += f"- {use_case}\n"
                
                if 'quantum_risk' in info:
                    report += f"\n## ⚠️ Quantum Risk\n{info['quantum_risk']}\n"
                
                return [TextContent(type="text", text=report)]
            else:
                return [TextContent(
                    type="text",
                    text=f"Algorithm '{algorithm}' not found. Available: DILITHIUM, FALCON, SPHINCS, RSA, ECDSA, HYBRID"
                )]
        else:
            # Return all algorithms
            report = "# Supported Cryptographic Algorithms\n\n"
            for algo_name, info in algorithm_info.items():
                status = "✅ PQC-Safe" if info['quantum_safe'] else "❌ Quantum-Vulnerable"
                report += f"## {info['name']} ({info['family']})\n"
                report += f"- **Status**: {status}\n"
                report += f"- **NIST**: {info['nist_status']}\n"
                report += f"- **Recommended**: {'Yes' if info['recommended'] else 'No (Legacy)'}\n\n"
            
            return [TextContent(type="text", text=report)]
    
    return [TextContent(type="text", text=f"Unknown tool: {name}")]


def main():
    """Main entry point for the MCP server"""
    asyncio.run(run_server())


async def run_server():
    """Run the MCP server"""
    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options()
        )


if __name__ == "__main__":
    main()