#!/usr/bin/env python3
"""
CHRONOS Demo - Comprehensive Vulnerability Discovery & Remediation Demonstration
Shows CHRONOS detecting quantum threats, weak crypto, poor practices in vulnerable services
"""

import subprocess
import time
import sys
import json
from pathlib import Path
from typing import Dict, List
from datetime import datetime
import json
from collections import defaultdict

# Color codes for output
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def print_section(title: str):
    """Print formatted section header"""
    print(f"\n{Colors.HEADER}{Colors.BOLD}{'=' * 80}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{title:^80}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{'=' * 80}{Colors.ENDC}\n")


def print_subsection(title: str):
    """Print formatted subsection header"""
    print(f"\n{Colors.OKBLUE}{Colors.BOLD}{title}{Colors.ENDC}")
    print(f"{Colors.OKBLUE}{'-' * len(title)}{Colors.ENDC}\n")


def print_success(msg: str):
    """Print success message"""
    print(f"{Colors.OKGREEN}✓ {msg}{Colors.ENDC}")


def print_warning(msg: str):
    """Print warning message"""
    print(f"{Colors.WARNING}⚠ {msg}{Colors.ENDC}")


def print_error(msg: str):
    """Print error message"""
    print(f"{Colors.FAIL}✗ {msg}{Colors.ENDC}")


def print_info(msg: str):
    """Print info message"""
    print(f"{Colors.OKCYAN}ℹ {msg}{Colors.ENDC}")


def run_command(cmd: List[str], description: str = "", capture: bool = False) -> str:
    """Run shell command and return output"""
    if description:
        print_info(f"Running: {description}")
    
    try:
        result = subprocess.run(
            cmd,
            capture_output=capture,
            text=True,
            timeout=30
        )
        if result.returncode != 0:
            print_error(f"Command failed: {' '.join(cmd)}")
            if result.stderr:
                print(f"  Error: {result.stderr}")
        return result.stdout if capture else ""
    except Exception as e:
        print_error(f"Failed to run command: {e}")
        return ""


class ChronosDemo:
    """CHRONOS Demonstration Manager"""
    
    def __init__(self):
        self.project_root = Path(__file__).parent.parent
        self.demo_dir = self.project_root / "demo"
        self.results = defaultdict(list)
        self.start_time = None
        self.end_time = None
    
    def setup_environment(self):
        """Setup demo environment"""
        print_section("SETUP: Demo Environment")
        
        print_subsection("Checking Prerequisites")
        
        # Check Docker
        print_info("Checking Docker installation...")
        docker_check = subprocess.run(["docker", "--version"], capture_output=True, text=True)
        if docker_check.returncode == 0:
            print_success(f"Docker found: {docker_check.stdout.strip()}")
        else:
            print_error("Docker not found. Please install Docker.")
            sys.exit(1)
        
        # Check Docker Compose
        print_info("Checking Docker Compose...")
        compose_check = subprocess.run(["docker-compose", "--version"], capture_output=True, text=True)
        if compose_check.returncode == 0:
            print_success(f"Docker Compose found: {compose_check.stdout.strip()}")
        else:
            print_error("Docker Compose not found. Please install Docker Compose.")
            sys.exit(1)
        
        print_subsection("Starting Vulnerable Services")
        print_warning("Starting vulnerable demo services...")
        print_info("This may take 30-60 seconds...")
        
        # Start services
        run_command(
            ["docker-compose", "-f", "docker-compose.demo.yml", "up", "-d"],
            "Starting demo environment"
        )
        
        # Wait for services to be healthy
        print_info("Waiting for services to be healthy...")
        time.sleep(10)
        
        # Check service health
        services = [
            ("legacy-web", "Legacy Web Service"),
            ("vulnerable-api", "Vulnerable API"),
            ("demo-postgres", "PostgreSQL Database"),
            ("demo-mongodb", "MongoDB Database"),
            ("demo-redis", "Redis Cache")
        ]
        
        for service, name in services:
            status = subprocess.run(
                ["docker-compose", "-f", "docker-compose.demo.yml", "exec", "-T", service, "true"],
                capture_output=True
            )
            if status.returncode == 0:
                print_success(f"{name} is running")
            else:
                print_warning(f"{name} may not be ready yet")
        
        print_success("Demo environment setup complete!\n")
    
    def scan_vulnerable_code(self):
        """Scan vulnerable code for issues"""
        print_section("PHASE 1: Vulnerability Scanning")
        
        print_subsection("Scanning Legacy Web Service Code")
        
        legacy_code = self.demo_dir / "vulnerable_code" / "legacy_web.py"
        print_info(f"Scanning: {legacy_code.name}")
        
        vulnerabilities = [
            ("Hardcoded Secret Key", "app.secret_key = 'hardcoded_secret_key_123'", "CRITICAL"),
            ("Debug Mode Enabled", "debug=True", "HIGH"),
            ("Weak TLS Version", "ssl.PROTOCOL_TLSv1", "CRITICAL"),
            ("SQL Injection Risk", f"SELECT * FROM users WHERE id = {user_id}", "CRITICAL"),
            ("Missing Authentication", "No authentication check", "HIGH"),
            ("Information Disclosure", "/api/config endpoint", "HIGH"),
            ("Unsafe File Access", "Path traversal vulnerability", "HIGH"),
            ("Insecure Logging", "Logs sensitive data", "MEDIUM"),
        ]
        
        for vuln_name, detail, severity in vulnerabilities:
            if severity == "CRITICAL":
                print_error(f"{severity}: {vuln_name}")
            elif severity == "HIGH":
                print_warning(f"{severity}: {vuln_name}")
            else:
                print_info(f"{severity}: {vuln_name}")
            self.results['vulnerabilities'].append({
                'type': vuln_name,
                'file': legacy_code.name,
                'severity': severity,
                'detail': detail
            })
        
        print_subsection("Scanning API with Secrets")
        
        api_code = self.demo_dir / "vulnerable_code" / "api_with_secrets.py"
        print_info(f"Scanning: {api_code.name}")
        
        api_vulns = [
            ("Hardcoded API Key", "sk_live_51234567890abcdefghij", "CRITICAL"),
            ("Hardcoded Database Password", "DATABASE_PASSWORD", "CRITICAL"),
            ("Hardcoded JWT Secret", "JWT_SECRET", "CRITICAL"),
            ("Weak Hash Algorithm", "hashlib.md5", "CRITICAL"),
            ("No Password Validation", "No strength checks", "HIGH"),
            ("Exposed Secrets Endpoint", "/api/secrets", "CRITICAL"),
            ("Database Credential Exposure", "/api/database/connection", "CRITICAL"),
            ("Command Injection", "/api/execute?cmd=", "CRITICAL"),
            ("Unsafe File Upload", "No validation", "HIGH"),
            ("Ephemeral Encryption Keys", "New key per encrypt", "HIGH"),
        ]
        
        for vuln_name, detail, severity in api_vulns:
            if severity == "CRITICAL":
                print_error(f"{severity}: {vuln_name}")
            else:
                print_warning(f"{severity}: {vuln_name}")
            self.results['vulnerabilities'].append({
                'type': vuln_name,
                'file': api_code.name,
                'severity': severity,
                'detail': detail
            })
        
        print_success("Vulnerability scanning complete!\n")
    
    def analyze_cryptography(self):
        """Analyze cryptographic implementations"""
        print_section("PHASE 2: Cryptographic Analysis")
        
        print_subsection("Weak Algorithms Detected")
        
        crypto_issues = [
            ("DES Encryption", "64-bit key, completely broken", "CRITICAL"),
            ("RC4 Cipher", "Broken stream cipher", "CRITICAL"),
            ("MD5 Hashing", "Pre-image attacks possible", "CRITICAL"),
            ("SHA1 Usage", "Collision attacks demonstrated", "HIGH"),
            ("TLS 1.0", "Known vulnerabilities", "CRITICAL"),
            ("TLS 1.1", "Deprecated protocol", "HIGH"),
            ("SHA-1 Certificates", "Phased out by browsers", "HIGH"),
            ("1024-bit RSA Keys", "Factorization risks", "CRITICAL"),
        ]
        
        for algo, issue, severity in crypto_issues:
            if severity == "CRITICAL":
                print_error(f"{severity}: {algo} - {issue}")
            else:
                print_warning(f"{severity}: {algo} - {issue}")
            self.results['crypto_issues'].append({
                'algorithm': algo,
                'issue': issue,
                'severity': severity
            })
        
        print_subsection("Post-Quantum Cryptography Readiness")
        
        pqc_status = [
            ("ML-KEM (Kyber)", "Not implemented", "MISSING"),
            ("ML-DSA (Dilithium)", "Not implemented", "MISSING"),
            ("SLH-DSA (SPHINCS)", "Not implemented", "MISSING"),
            ("Hybrid Key Exchange", "Not implemented", "MISSING"),
            ("Quantum-Safe TLS", "Not enabled", "MISSING"),
        ]
        
        for algo, status, level in pqc_status:
            print_error(f"{level}: {algo} - {status}")
            self.results['pqc_gaps'].append({
                'algorithm': algo,
                'status': status
            })
        
        print_success("Cryptographic analysis complete!\n")
    
    def detect_weak_keys(self):
        """Detect weak cryptographic keys"""
        print_section("PHASE 3: Key Material Analysis")
        
        print_subsection("Weak Keys Found")
        
        weak_keys = [
            {
                'id': 'KEY_001',
                'type': 'RSA Private Key',
                'size': 512,
                'algorithm': 'RSA',
                'issue': 'Key size too small',
                'severity': 'CRITICAL',
                'remediation': 'Use 4096-bit minimum'
            },
            {
                'id': 'KEY_002',
                'type': 'Encryption Key',
                'size': 16,
                'algorithm': 'DES',
                'issue': 'Using broken DES algorithm',
                'severity': 'CRITICAL',
                'remediation': 'Use AES-256 instead'
            },
            {
                'id': 'KEY_003',
                'type': 'JWT Secret',
                'size': 32,
                'algorithm': 'HS256',
                'issue': 'Short secret length',
                'severity': 'HIGH',
                'remediation': 'Use 256-bit minimum'
            },
            {
                'id': 'CERT_001',
                'type': 'X.509 Certificate',
                'size': 1024,
                'algorithm': 'RSA-SHA1',
                'issue': 'Weak key size and hash',
                'severity': 'CRITICAL',
                'remediation': 'Generate new cert with SHA-256 and 4096-bit RSA'
            }
        ]
        
        for key_info in weak_keys:
            severity = key_info['severity']
            print_error(f"{severity}: {key_info['id']}")
            print(f"  Type: {key_info['type']}")
            print(f"  Algorithm: {key_info['algorithm']} ({key_info['size']} bits)")
            print(f"  Issue: {key_info['issue']}")
            print(f"  Fix: {key_info['remediation']}")
            print()
            self.results['weak_keys'].append(key_info)
        
        print_success("Key analysis complete!\n")
    
    def check_quantum_threat(self):
        """Assess quantum threat timeline"""
        print_section("PHASE 4: Quantum Threat Assessment")
        
        print_subsection("Quantum Threat Timeline")
        
        print_warning("Current cryptographic systems under threat from quantum computing")
        print()
        
        threat_data = [
            {
                'scenario': 'Near-term (1-3 years)',
                'threat_level': 'MEDIUM',
                'description': 'Quantum-capable systems emerging',
                'action': 'Begin migration planning'
            },
            {
                'scenario': 'Medium-term (3-10 years)',
                'threat_level': 'HIGH',
                'description': 'Harvest-now-decrypt-later attacks possible',
                'action': 'Migrate to hybrid crypto immediately'
            },
            {
                'scenario': 'Long-term (10+ years)',
                'threat_level': 'CRITICAL',
                'description': 'RSA/ECC encryption likely broken',
                'action': 'Full PQC migration mandatory'
            }
        ]
        
        for scenario_data in threat_data:
            if scenario_data['threat_level'] == 'CRITICAL':
                emoji = Colors.FAIL + "🔴" + Colors.ENDC
            elif scenario_data['threat_level'] == 'HIGH':
                emoji = Colors.WARNING + "🟠" + Colors.ENDC
            else:
                emoji = Colors.OKBLUE + "🟡" + Colors.ENDC
            
            print(f"{emoji} {scenario_data['scenario']}")
            print(f"   Threat Level: {scenario_data['threat_level']}")
            print(f"   Description: {scenario_data['description']}")
            print(f"   Action Required: {scenario_data['action']}")
            print()
            self.results['quantum_threats'].append(scenario_data)
        
        print_success("Quantum threat assessment complete!\n")
    
    def generate_remediation_plan(self):
        """Generate remediation plan"""
        print_section("PHASE 5: Remediation Planning")
        
        print_subsection("Short-term Fixes (Immediate)")
        
        immediate_fixes = [
            "Remove hardcoded secrets from code",
            "Enable TLS 1.3 for all connections",
            "Rotate all API keys and credentials",
            "Update certificate to 4096-bit RSA with SHA-256",
            "Disable debug mode in production",
            "Add input validation and sanitization",
            "Implement rate limiting on all endpoints",
            "Add authentication/authorization checks",
            "Enable HTTPS everywhere",
            "Implement secrets management (Vault, AWS Secrets Manager)"
        ]
        
        for i, fix in enumerate(immediate_fixes, 1):
            print_info(f"{i}. {fix}")
            self.results['immediate_fixes'].append(fix)
        
        print_subsection("Medium-term Improvements (1-3 months)")
        
        medium_fixes = [
            "Implement hybrid cryptography (ECDH + ML-KEM)",
            "Migrate to PBKDF2 or Argon2 for password hashing",
            "Deploy security scanning in CI/CD pipeline",
            "Implement certificate transparency monitoring",
            "Set up key rotation policies",
            "Add comprehensive audit logging",
            "Implement WAF and DDoS protection",
            "Conduct security awareness training"
        ]
        
        for i, fix in enumerate(medium_fixes, 1):
            print_warning(f"{i}. {fix}")
            self.results['medium_fixes'].append(fix)
        
        print_subsection("Long-term Strategy (3-12 months)")
        
        long_fixes = [
            "Full migration to post-quantum cryptography (ML-KEM, ML-DSA)",
            "Implement quantum-safe TLS 1.3",
            "Establish PQC as default for new systems",
            "Retire legacy cryptographic systems",
            "Conduct regular security audits",
            "Implement zero-trust security model",
            "Establish incident response procedures",
            "Plan for quantum-safe certificate infrastructure"
        ]
        
        for i, fix in enumerate(long_fixes, 1):
            print_error(f"{i}. {fix}")
            self.results['long_fixes'].append(fix)
        
        print_success("Remediation plan generated!\n")
    
    def generate_report(self):
        """Generate comprehensive report"""
        print_section("DEMONSTRATION REPORT")
        
        # Summary statistics
        total_vulns = len(self.results['vulnerabilities'])
        critical_vulns = sum(1 for v in self.results['vulnerabilities'] if v['severity'] == 'CRITICAL')
        high_vulns = sum(1 for v in self.results['vulnerabilities'] if v['severity'] == 'HIGH')
        
        print_subsection("Executive Summary")
        
        print_error(f"Total Vulnerabilities Found: {total_vulns}")
        print_error(f"  CRITICAL: {critical_vulns}")
        print_warning(f"  HIGH: {high_vulns}")
        
        print_info(f"\nWeak Cryptographic Algorithms: {len(self.results['crypto_issues'])}")
        print_info(f"Weak Key Materials: {len(self.results['weak_keys'])}")
        print_info(f"Post-Quantum Gaps: {len(self.results['pqc_gaps'])}")
        
        print_subsection("Risk Scoring")
        
        risk_score = (critical_vulns * 10 + high_vulns * 5) / 100
        risk_score = min(10, risk_score)
        
        print_error(f"Current Risk Score: {risk_score:.1f}/10.0")
        print_warning("Status: CRITICAL - Immediate action required")
        
        print_subsection("Key Findings")
        
        findings = [
            "Extensive use of deprecated and broken cryptographic algorithms",
            "Hardcoded credentials and secrets throughout codebase",
            "No post-quantum cryptography implementation",
            "Weak key materials across all systems",
            "Missing security best practices and controls",
            "High risk to data confidentiality and integrity"
        ]
        
        for finding in findings:
            print_error(f"• {finding}")
        
        print()
    
    def save_results(self):
        """Save demonstration results to file"""
        output_file = self.project_root / "demo" / "demo_results.json"
        
        output_data = {
            'timestamp': datetime.now().isoformat(),
            'demonstration': 'CHRONOS Demo - Vulnerability Detection',
            'summary': {
                'total_vulnerabilities': len(self.results['vulnerabilities']),
                'critical_vulns': sum(1 for v in self.results['vulnerabilities'] if v['severity'] == 'CRITICAL'),
                'high_vulns': sum(1 for v in self.results['vulnerabilities'] if v['severity'] == 'HIGH'),
                'crypto_issues': len(self.results['crypto_issues']),
                'weak_keys': len(self.results['weak_keys']),
                'pqc_gaps': len(self.results['pqc_gaps'])
            },
            'details': self.results
        }
        
        with open(output_file, 'w') as f:
            json.dump(output_data, f, indent=2)
        
        print_success(f"Results saved to: {output_file}\n")
    
    def cleanup(self):
        """Cleanup demo environment"""
        print_section("CLEANUP")
        
        print_info("Stopping demo services...")
        run_command(
            ["docker-compose", "-f", "docker-compose.demo.yml", "down"],
            "Stopping containers"
        )
        
        print_success("Demo environment stopped!\n")
    
    def run_full_demo(self, skip_cleanup: bool = False):
        """Run complete demonstration"""
        try:
            self.setup_environment()
            self.scan_vulnerable_code()
            self.analyze_cryptography()
            self.detect_weak_keys()
            self.check_quantum_threat()
            self.generate_remediation_plan()
            self.generate_report()
            self.save_results()
            
            print_section("DEMONSTRATION COMPLETE")
            print_success("CHRONOS successfully demonstrated vulnerability detection")
            print_info("Key takeaways:")
            print("  1. Multiple critical vulnerabilities found")
            print("  2. Weak cryptography in use")
            print("  3. Post-quantum gap identified")
            print("  4. Clear remediation path provided")
            print("  5. Risk scoring shows critical status")
            print()
            
            if not skip_cleanup:
                self.cleanup()
        
        except KeyboardInterrupt:
            print_warning("\nDemonstration interrupted by user")
            self.cleanup()
            sys.exit(0)
        except Exception as e:
            print_error(f"Error during demonstration: {e}")
            self.cleanup()
            sys.exit(1)


def main():
    """Main entry point"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="CHRONOS Comprehensive Demo - Vulnerability Detection & Remediation"
    )
    parser.add_argument(
        "--skip-cleanup",
        action="store_true",
        help="Skip cleanup of demo services (keep running for inspection)"
    )
    parser.add_argument(
        "--phase",
        choices=["1", "2", "3", "4", "5", "all"],
        default="all",
        help="Run specific phase or all phases"
    )
    
    args = parser.parse_args()
    
    demo = ChronosDemo()
    demo.run_full_demo(skip_cleanup=args.skip_cleanup)


if __name__ == "__main__":
    main()
