#!/usr/bin/env python3
"""
MCP Security Auditor CLI - Command-line interface for MCP security analysis.

Usage:
    mcp-audit scan <path> [options]
    mcp-audit report <scan-id> [options]
    mcp-audit ci <path> [options]
"""

import argparse
import sys
import os
import json
import asyncio
from pathlib import Path
from typing import Optional
from datetime import datetime

from .core.scanner import MCPSecurityScanner
from .core.config import AuditConfig
from .core.report import ReportGenerator
from .core.models import ScanResult
from .analyzers import (
    StaticCodeAnalyzer,
    PermissionAnalyzer,
    NetworkAnalyzer,
    DependencyAnalyzer,
    InjectionPatternAnalyzer,
    ConfigurationAnalyzer,
    SecretsAnalyzer
)


def create_parser() -> argparse.ArgumentParser:
    """Create the argument parser with all subcommands."""
    parser = argparse.ArgumentParser(
        prog="mcp-audit",
        description="MCP Security Auditor - Evaluate security of MCP servers before deployment",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Scan a local MCP server directory
  mcp-audit scan ./my-mcp-server

  # Scan with specific severity threshold
  mcp-audit scan ./my-mcp-server --severity high

  # Generate SARIF report for CI/CD
  mcp-audit scan ./my-mcp-server --format sarif --output report.sarif

  # Generate comprehensive HTML report
  mcp-audit scan ./my-mcp-server --format html --output security-report.html

  # CI mode with exit code based on findings
  mcp-audit ci ./my-mcp-server --fail-on high

  # Output SIEM-compatible logs
  mcp-audit scan ./my-mcp-server --format siem --siem-format cef

  # Scan from Git URL
  mcp-audit scan https://github.com/user/mcp-server.git
        """
    )
    parser.add_argument(
        "--version", "-v",
        action="version",
        version="%(prog)s 1.0.0"
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose output"
    )
    parser.add_argument(
        "--quiet", "-q",
        action="store_true",
        help="Suppress non-essential output"
    )
    
    subparsers = parser.add_subparsers(dest="command", help="Available commands")
    
    # Scan command
    scan_parser = subparsers.add_parser(
        "scan",
        help="Scan an MCP server for security vulnerabilities"
    )
    scan_parser.add_argument(
        "target",
        help="Path to MCP server directory or Git URL"
    )
    scan_parser.add_argument(
        "--severity", "-s",
        choices=["critical", "high", "medium", "low", "info"],
        default="info",
        help="Minimum severity level to report (default: info)"
    )
    scan_parser.add_argument(
        "--format", "-f",
        choices=["json", "html", "sarif", "markdown", "siem", "text"],
        default="text",
        help="Output format (default: text)"
    )
    scan_parser.add_argument(
        "--output", "-o",
        type=str,
        help="Output file path (default: stdout for text/json)"
    )
    scan_parser.add_argument(
        "--config", "-c",
        type=str,
        help="Path to custom configuration file"
    )
    scan_parser.add_argument(
        "--skip-analyzers",
        type=str,
        nargs="+",
        choices=["static", "permissions", "network", "dependencies", "injection", "config", "secrets", "readonly"],
        help="Skip specific analyzers"
    )
    scan_parser.add_argument(
        "--include-analyzers",
        type=str,
        nargs="+",
        choices=["static", "permissions", "network", "dependencies", "injection", "config", "secrets", "readonly"],
        help="Run only specific analyzers"
    )
    scan_parser.add_argument(
        "--siem-format",
        choices=["cef", "leef", "json-syslog", "splunk"],
        default="cef",
        help="SIEM output format when --format is siem (default: cef)"
    )
    scan_parser.add_argument(
        "--no-color",
        action="store_true",
        help="Disable colored output"
    )
    scan_parser.add_argument(
        "--rules",
        type=str,
        help="Path to custom rules file (YAML)"
    )
    
    # CI command
    ci_parser = subparsers.add_parser(
        "ci",
        help="Run in CI/CD mode with exit codes based on findings"
    )
    ci_parser.add_argument(
        "target",
        help="Path to MCP server directory"
    )
    ci_parser.add_argument(
        "--fail-on",
        choices=["critical", "high", "medium", "low"],
        default="high",
        help="Fail (exit 1) if findings at or above this severity (default: high)"
    )
    ci_parser.add_argument(
        "--format", "-f",
        choices=["json", "sarif", "text"],
        default="sarif",
        help="Output format (default: sarif)"
    )
    ci_parser.add_argument(
        "--output", "-o",
        type=str,
        help="Output file path"
    )
    ci_parser.add_argument(
        "--baseline",
        type=str,
        help="Path to baseline file to suppress known issues"
    )
    
    # Report command
    report_parser = subparsers.add_parser(
        "report",
        help="Generate report from previous scan results"
    )
    report_parser.add_argument(
        "scan_file",
        help="Path to scan results JSON file"
    )
    report_parser.add_argument(
        "--format", "-f",
        choices=["html", "markdown", "sarif", "json", "text"],
        default="html",
        help="Report format (default: html)"
    )
    report_parser.add_argument(
        "--output", "-o",
        type=str,
        required=True,
        help="Output file path"
    )
    report_parser.add_argument(
        "--template",
        type=str,
        help="Custom report template"
    )
    
    # Init command
    init_parser = subparsers.add_parser(
        "init",
        help="Initialize configuration file"
    )
    init_parser.add_argument(
        "--output", "-o",
        type=str,
        default=".mcp-audit.yaml",
        help="Configuration file path (default: .mcp-audit.yaml)"
    )
    
    # Certify command
    certify_parser = subparsers.add_parser(
        "certify",
        help="Generate security certification for MCP server"
    )
    certify_parser.add_argument(
        "target",
        help="Path to MCP server directory"
    )
    certify_parser.add_argument(
        "--level",
        choices=["basic", "standard", "enterprise"],
        default="standard",
        help="Certification level (default: standard)"
    )
    certify_parser.add_argument(
        "--output", "-o",
        type=str,
        required=True,
        help="Output certification file path"
    )
    
    return parser


async def run_scan(args) -> int:
    """Execute a security scan."""
    config = AuditConfig.load(args.config if hasattr(args, 'config') and args.config else None)
    
    # Determine which analyzers to run
    analyzers_to_run = None
    if hasattr(args, 'include_analyzers') and args.include_analyzers:
        analyzers_to_run = args.include_analyzers
    elif hasattr(args, 'skip_analyzers') and args.skip_analyzers:
        all_analyzers = ["static", "permissions", "network", "dependencies", "injection", "config", "secrets"]
        analyzers_to_run = [a for a in all_analyzers if a not in args.skip_analyzers]
    
    scanner = MCPSecurityScanner(
        config=config,
        analyzers=analyzers_to_run,
        verbose=args.verbose if hasattr(args, 'verbose') else False
    )
    
    # Run scan
    if not args.quiet:
        print(f"🔍 Scanning MCP server at: {args.target}")
        print(f"   Minimum severity: {args.severity}")
        print()
    
    try:
        results = await scanner.scan(args.target)
    except Exception as e:
        print(f"❌ Scan failed: {e}", file=sys.stderr)
        return 1
    
    # Filter by severity
    severity_levels = ["critical", "high", "medium", "low", "info"]
    min_severity_idx = severity_levels.index(args.severity)
    filtered_findings = [
        f for f in results.findings
        if severity_levels.index(f.severity) <= min_severity_idx
    ]
    results.findings = filtered_findings
    
    # Generate report
    generator = ReportGenerator(results)
    
    output_format = args.format
    siem_format = getattr(args, 'siem_format', 'cef')
    
    if output_format == "text":
        report = generator.to_text(use_color=not getattr(args, 'no_color', False))
    elif output_format == "json":
        report = generator.to_json()
    elif output_format == "sarif":
        report = generator.to_sarif()
    elif output_format == "html":
        report = generator.to_html()
    elif output_format == "markdown":
        report = generator.to_markdown()
    elif output_format == "siem":
        report = generator.to_siem(format_type=siem_format)
    else:
        report = generator.to_text()
    
    # Output
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(report)
        if not args.quiet:
            print(f"📄 Report saved to: {args.output}")
    else:
        print(report)
    
    return 0


async def run_ci(args) -> int:
    """Execute a CI/CD security scan with exit codes."""
    config = AuditConfig.load(None)
    
    scanner = MCPSecurityScanner(config=config, verbose=False)
    
    print(f"🔍 CI Security Scan: {args.target}")
    
    try:
        results = await scanner.scan(args.target)
    except Exception as e:
        print(f"❌ Scan failed: {e}", file=sys.stderr)
        return 2  # Exit code 2 for errors
    
    # Load baseline if provided
    baseline_ids = set()
    if args.baseline and Path(args.baseline).exists():
        with open(args.baseline, encoding='utf-8') as f:
            baseline_data = json.load(f)
            baseline_ids = set(baseline_data.get("suppressed_findings", []))
    
    # Filter out baseline findings
    results.findings = [f for f in results.findings if f.id not in baseline_ids]
    
    # Generate report
    generator = ReportGenerator(results)
    
    if args.format == "sarif":
        report = generator.to_sarif()
    elif args.format == "json":
        report = generator.to_json()
    else:
        report = generator.to_text(use_color=False)
    
    if args.output:
        with open(args.output, 'w', encoding='utf-8') as f:
            f.write(report)
        print(f"📄 Report saved to: {args.output}")
    else:
        print(report)
    
    # Check if we should fail
    severity_levels = ["critical", "high", "medium", "low"]
    fail_threshold_idx = severity_levels.index(args.fail_on)
    
    failing_findings = [
        f for f in results.findings
        if f.severity in severity_levels and severity_levels.index(f.severity) <= fail_threshold_idx
    ]
    
    if failing_findings:
        print(f"\n❌ Found {len(failing_findings)} findings at or above '{args.fail_on}' severity")
        return 1
    
    print(f"\n✅ No findings at or above '{args.fail_on}' severity")
    return 0


def run_init(args) -> int:
    """Initialize a configuration file."""
    default_config = """# MCP Security Auditor Configuration
# Version: 1.0

# Severity threshold for reporting
severity_threshold: info

# Analyzers to run
analyzers:
  static: true
  permissions: true
  network: true
  dependencies: true
  injection: true
  config: true
  secrets: true
  readonly: true  # Critical for database MCPs

# Custom rules
rules:
  # Patterns considered dangerous in code
  dangerous_patterns:
    - pattern: "eval\\\\("
      severity: critical
      description: "Dynamic code execution detected"
    - pattern: "exec\\\\("
      severity: critical
      description: "System command execution detected"
    - pattern: "subprocess\\\\.call"
      severity: high
      description: "Subprocess execution detected"
  
  # Allowed domains for network analysis
  allowed_domains:
    - "api.anthropic.com"
    - "modelcontextprotocol.io"
  
  # Required permissions that must be declared
  required_permissions: []
  
  # Forbidden permissions
  forbidden_permissions:
    - "system.shell"
    - "filesystem.root"

# Dependency scanning
dependencies:
  check_vulnerabilities: true
  vulnerability_sources:
    - osv
    - nvd

# Read-only mode enforcement
readonly:
  # Require all tools to have readOnlyHint annotation
  require_readonly_annotation: true
  # Fail on any write operations detected
  strict_mode: false
  # Allowed write operations (if strict_mode is false)
  allowed_write_tools: []
  
# Secrets detection
secrets:
  patterns:
    - name: "AWS Access Key"
      pattern: "AKIA[0-9A-Z]{16}"
      severity: critical
    - name: "API Key"
      pattern: "[a-zA-Z0-9_-]*api[_-]?key[a-zA-Z0-9_-]*\\\\s*[=:]\\\\s*['\"][^'\"]{16,}['\"]"
      severity: high
    - name: "Private Key"
      pattern: "-----BEGIN (RSA |EC |DSA |OPENSSH )?PRIVATE KEY-----"
      severity: critical

# CI/CD settings
ci:
  fail_on: high
  output_format: sarif
  
# SIEM integration
siem:
  format: cef
  vendor: "MCPAuditor"
  product: "MCP Security Auditor"
  version: "1.0"

# Report settings  
report:
  include_remediation: true
  include_code_snippets: true
  max_snippet_lines: 10
"""
    
    output_path = Path(args.output)
    if output_path.exists():
        response = input(f"File {args.output} already exists. Overwrite? [y/N]: ")
        if response.lower() != 'y':
            print("Aborted.")
            return 1
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(default_config)
    
    print(f"✅ Configuration file created: {args.output}")
    return 0


async def run_certify(args) -> int:
    """Generate security certification."""
    config = AuditConfig.load(None)
    scanner = MCPSecurityScanner(config=config, verbose=False)
    
    print(f"🔍 Running certification scan: {args.target}")
    print(f"   Certification level: {args.level}")
    
    try:
        results = await scanner.scan(args.target)
    except Exception as e:
        print(f"❌ Certification failed: {e}", file=sys.stderr)
        return 1
    
    # Determine certification thresholds
    thresholds = {
        "basic": {"critical": 0, "high": 3, "medium": 10},
        "standard": {"critical": 0, "high": 0, "medium": 5},
        "enterprise": {"critical": 0, "high": 0, "medium": 0}
    }
    
    threshold = thresholds[args.level]
    
    # Count findings by severity
    counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
    for finding in results.findings:
        counts[finding.severity] = counts.get(finding.severity, 0) + 1
    
    # Check if certification passes
    certified = (
        counts["critical"] <= threshold["critical"] and
        counts["high"] <= threshold["high"] and
        counts["medium"] <= threshold["medium"]
    )
    
    certification = {
        "certified": certified,
        "level": args.level,
        "target": args.target,
        "scan_date": datetime.utcnow().isoformat(),
        "expiry_date": (datetime.utcnow().replace(year=datetime.utcnow().year + 1)).isoformat(),
        "findings_summary": counts,
        "thresholds": threshold,
        "scanner_version": "1.0.0",
        "certification_id": f"MCP-CERT-{datetime.utcnow().strftime('%Y%m%d%H%M%S')}",
        "risk_score": results.risk_score
    }
    
    with open(args.output, 'w', encoding='utf-8') as f:
        json.dump(certification, f, indent=2)
    
    if certified:
        print(f"\n✅ CERTIFIED at {args.level} level")
        print(f"   Certification ID: {certification['certification_id']}")
        print(f"   Risk Score: {results.risk_score}/100")
        print(f"   Valid until: {certification['expiry_date'][:10]}")
    else:
        print(f"\n❌ CERTIFICATION FAILED for {args.level} level")
        print(f"   Findings exceed threshold:")
        for sev in ["critical", "high", "medium"]:
            if counts[sev] > threshold[sev]:
                print(f"     - {sev}: {counts[sev]} (max: {threshold[sev]})")
    
    print(f"\n📄 Certification saved to: {args.output}")
    return 0 if certified else 1


def run_report(args) -> int:
    """Generate a report from previously saved scan results."""
    scan_file = Path(args.scan_file)
    
    if not scan_file.exists():
        print(f"❌ Scan file not found: {args.scan_file}", file=sys.stderr)
        return 1
    
    # Load scan results
    try:
        with open(scan_file, 'r', encoding='utf-8') as f:
            scan_data = json.load(f)
        results = ScanResult.from_dict(scan_data)
    except json.JSONDecodeError as e:
        print(f"❌ Invalid JSON in scan file: {e}", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"❌ Failed to load scan results: {e}", file=sys.stderr)
        return 1
    
    # Generate report
    generator = ReportGenerator(results)
    
    output_format = args.format
    if output_format == "html":
        report = generator.to_html()
    elif output_format == "markdown":
        report = generator.to_markdown()
    elif output_format == "sarif":
        report = generator.to_sarif()
    elif output_format == "json":
        report = generator.to_json()
    elif output_format == "text":
        report = generator.to_text(use_color=False)
    else:
        report = generator.to_html()
    
    # Write output
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(report)
    
    print(f"✅ Report generated: {args.output}")
    print(f"   Format: {output_format}")
    print(f"   Findings: {len(results.findings)}")
    print(f"   Risk Score: {results.risk_score}/100")
    return 0


def main():
    """Main entry point."""
    parser = create_parser()
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        sys.exit(0)
    
    # Run the appropriate command
    if args.command == "scan":
        exit_code = asyncio.run(run_scan(args))
    elif args.command == "ci":
        exit_code = asyncio.run(run_ci(args))
    elif args.command == "report":
        exit_code = run_report(args)
    elif args.command == "init":
        exit_code = run_init(args)
    elif args.command == "certify":
        exit_code = asyncio.run(run_certify(args))
    else:
        parser.print_help()
        exit_code = 0
    
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
