# ===================================================================
# MODULE 10, LESSON 5: QUANTUM COMPUTING & POST-QUANTUM CRYPTOGRAPHY
# ===================================================================

import pandas as pd
import numpy as np
import hashlib
import time
import json
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import warnings
warnings.filterwarnings('ignore')

print("="*70)
print("QUANTUM RISK ASSESSMENT & PQC SIMULATION")
print("="*70)

# ----------------------------------------------------------------
# PART A: CRYPTOGRAPHIC INVENTORY (Identifying Vulnerable Assets)
# ----------------------------------------------------------------
print("\n" + "-"*60)
print("PART A: Cryptographic Asset Inventory")
print("-"*60)

crypto_inventory = pd.DataFrame({
    'Asset': [
        'Customer Authentication (Mobile App)',
        'API Gateway (REST/TLS)',
        'SWIFT Payment Signing',
        'Internal Email Encryption',
        'Blockchain Private Keys (CBDC)',
        'Loan Document Signing',
        'Data-at-Rest Encryption (HDFS)'
    ],
    'Algorithm': [
        'ECC-256',
        'RSA-2048 + ECDHE',
        'ECDSA-256',
        'RSA-4096',
        'Secp256k1 (ECDSA)',
        'RSA-2048',
        'AES-256-GCM'
    ],
    'Quantum_Vulnerable': [
        True,  # ECC broken by Shor
        True,  # RSA broken by Shor
        True,  # ECC broken by Shor
        True,  # RSA broken by Shor
        True,  # ECC broken by Shor
        True,  # RSA broken by Shor
        False  # AES-256 only weakened by Grover
    ],
    'PQC_Replacement': [
        'CRYSTALS-Dilithium',
        'CRYSTALS-Kyber (ML-KEM)',
        'Falcon-512',
        'SPHINCS+',
        'CRYSTALS-Dilithium',
        'CRYSTALS-Dilithium',
        'AES-256 (double key size)'
    ],
    'Migration_Priority': ['High', 'Critical', 'High', 'Medium', 'Critical', 'Medium', 'Low']
})

print("Cryptographic Inventory with Quantum Risk Assessment:")
print(crypto_inventory.to_string(index=False))

# ----------------------------------------------------------------
# PART B: SIMULATING SHOR'S ALGORITHM IMPACT (Key Breaking Time)
# ----------------------------------------------------------------
print("\n" + "-"*60)
print("PART B: Simulating Key Breaking Time with Quantum Computers")
print("-"*60)

def quantum_break_time(bits, qubit_speed=1000000):
    """
    Simulates time to break RSA/ECC using Shor's Algorithm.
    Qubit_speed: operations per second (1M = realistic for 2030).
    """
    # Shor's algorithm complexity: O(log^3 n) operations
    ops = (bits ** 3)  # Simplified
    seconds = ops / qubit_speed
    hours = seconds / 3600
    days = hours / 24
    years = days / 365
    return years

algorithms = [
    {'name': 'RSA-1024', 'bits': 1024},
    {'name': 'RSA-2048', 'bits': 2048},
    {'name': 'RSA-4096', 'bits': 4096},
    {'name': 'ECC-256', 'bits': 256},
    {'name': 'ECC-384', 'bits': 384},
]

break_data = []
for algo in algorithms:
    years = quantum_break_time(algo['bits'])
    break_data.append({
        'Algorithm': algo['name'],
        'Key Size (bits)': algo['bits'],
        'Time to Break (Years)': round(years, 2),
        'Risk': 'High' if years < 5 else ('Medium' if years < 20 else 'Low')
    })

break_df = pd.DataFrame(break_data)
print("Estimated Time to Break Classical Cryptography (2030 Quantum Speed):")
print(break_df.to_string(index=False))

# ----------------------------------------------------------------
# PART C: SIMULATING POST-QUANTUM HYBRID ENCRYPTION
# ----------------------------------------------------------------
print("\n" + "-"*60)
print("PART C: Hybrid Classical + PQC Encryption Simulation")
print("-"*60)

class HybridQuantumSafeCrypto:
    """
    Simulates a hybrid encryption scheme: Classical (RSA) + PQC (Lattice-based).
    In production, you would use Kyber or Dilithium libraries.
    """
    def __init__(self):
        # Generate Classical RSA Key (Simulated)
        self.rsa_private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        self.rsa_public_key = self.rsa_private_key.public_key()
        
        # Simulate PQC Public/Private Key (Lattice-based)
        # In reality: use 'pqcrypto' or 'liboqs' library
        self.pqc_public_key = "PQC_PUBLIC_KEY_" + hashlib.sha256(b"lattice_seed").hexdigest()[:16]
        self.pqc_private_key = "PQC_PRIVATE_KEY_" + hashlib.sha256(b"lattice_secret").hexdigest()[:16]
        
        print("Hybrid Crypto System Initialized: RSA-2048 + Lattice-based PQC (Simulated)")

    def hybrid_encrypt(self, plaintext):
        """Encrypts a message using both RSA and PQC (Hybrid)."""
        print(f"\nEncrypting: '{plaintext}'")
        
        # 1. Classical RSA Encryption (Simulated)
        rsa_ciphertext = f"RSA_{plaintext}_encrypted"
        print(f"  - RSA-2048 Ciphertext: {rsa_ciphertext[:20]}...")
        
        # 2. PQC Encryption (Simulated Lattice-based)
        pqc_ciphertext = f"PQC_{plaintext}_encrypted"
        print(f"  - PQC (Lattice) Ciphertext: {pqc_ciphertext[:20]}...")
        
        # 3. Combined Payload
        hybrid_payload = {
            'rsa_part': rsa_ciphertext,
            'pqc_part': pqc_ciphertext,
            'algorithm': 'RSA-2048 + CRYSTALS-Kyber (Sim)'
        }
        return hybrid_payload

    def hybrid_decrypt(self, hybrid_payload):
        """Decrypts using both systems (fallback if one fails)."""
        print("\nDecrypting Hybrid Payload...")
        
        # Simulate RSA decryption (valid)
        rsa_plaintext = hybrid_payload['rsa_part'].replace('RSA_', '').replace('_encrypted', '')
        
        # Simulate PQC decryption (valid)
        pqc_plaintext = hybrid_payload['pqc_part'].replace('PQC_', '').replace('_encrypted', '')
        
        # Verify both decryptions match (consensus)
        if rsa_plaintext == pqc_plaintext:
            print(f"  ✅ Both algorithms agree on plaintext: '{rsa_plaintext}'")
            return rsa_plaintext
        else:
            print("  ⚠️ Consensus failed! Potential quantum attack or corruption.")
            return None

# Instantiate and test hybrid crypto
hybrid = HybridQuantumSafeCrypto()
message = "WIRE_TRANSFER_12345_AMOUNT_50000"

encrypted = hybrid.hybrid_encrypt(message)
print(f"\nHybrid Payload Structure: {json.dumps(encrypted, indent=2)}")

decrypted = hybrid.hybrid_decrypt(encrypted)
print(f"\nFinal Decrypted Message: {decrypted}")

# ----------------------------------------------------------------
# PART D: QUANTUM READINESS ROADMAP
# ----------------------------------------------------------------
print("\n" + "-"*60)
print("PART D: Quantum Readiness Roadmap for Banks")
print("-"*60)

quantum_roadmap = {
    "Phase 1 (2025-2026) – Discovery": {
        "Focus": "Inventory and assessment.",
        "Activities": [
            "Complete cryptographic inventory (all keys, certificates, HSMs).",
            "Identify critical assets (payment systems, customer auth).",
            "Start crypto-agility design (configuration-driven algorithms)."
        ],
        "Success Metrics": ["100% crypto asset inventory complete", "Crypto-agility proof-of-concept"]
    },
    "Phase 2 (2027-2028) – Hybrid Deployment": {
        "Focus": "Implement hybrid classical/PQC in non-critical systems.",
        "Activities": [
            "Upgrade TLS to support hybrid key exchange (Kyber + ECDH).",
            "Pilot PQC signing for internal document management.",
            "Test PQC libraries (Open Quantum Safe, liboqs)."
        ],
        "Success Metrics": ["TLS 1.3 with PQC support in test environments", "Internal PQC pilot complete"]
    },
    "Phase 3 (2029-2030) – Critical Systems Migration": {
        "Focus": "Move core banking systems to PQC.",
        "Activities": [
            "Migrate payment signing to CRYSTALS-Dilithium.",
            "Update HSM hardware for PQC support.",
            "Implement PQC for SWIFT/ISO 20022 messages."
        ],
        "Success Metrics": ["80% critical systems PQC-ready", "HSMs with PQC support"]
    },
    "Phase 4 (2031+) – Full PQC & Continuous Monitoring": {
        "Focus": "Complete migration and ongoing assessment.",
        "Activities": [
            "Complete full PQC migration (all systems).",
            "Implement continuous quantum threat monitoring.",
            "Participate in industry PQC standards evolution."
        ],
        "Success Metrics": ["100% PQC migration complete", "Quantum-safe bank certification"]
    }
}

for phase, details in quantum_roadmap.items():
    print(f"\n{phase}:")
    print(f"  Focus: {details['Focus']}")
    print("  Activities:")
    for activity in details['Activities']:
        print(f"    • {activity}")
    print("  Success Metrics:")
    for metric in details['Success Metrics']:
        print(f"    • {metric}")

# ----------------------------------------------------------------
# SECTION 5: SUMMARY FOR THE DATA PRACTITIONER
# ----------------------------------------------------------------
print("\n" + "="*70)
print("LESSON 5 SUMMARY FOR THE DATA PRACTITIONER")
print("="*70)

print("""
1. Quantum computers will break RSA and ECC within hours (Shor's Algorithm).
2. Banks must migrate to NIST-standardized PQC algorithms (Kyber, Dilithium, Falcon).
3. 'Store Now, Decrypt Later' (SNDL) attacks are a real threat today – historical data is at risk.
4. Implementation Strategy:
   - Phase 1: Inventory all crypto assets.
   - Phase 2: Implement crypto-agile architecture.
   - Phase 3: Hybrid deployment (Classical + PQC).
   - Phase 4: Full PQC migration.
5. As a Data Practitioner, your role includes:
   - Identifying where sensitive data is encrypted.
   - Ensuring data pipelines use PQC-ready libraries.
   - Monitoring NIST standards evolution.
6. Action: Conduct a cryptographic inventory of your organization's data storage and communication channels.
""")

print("="*70)
print("END OF LESSON 5 – MODULE 10")
print("="*70)