#!/usr/bin/env python3
"""
Entropy Impact Calculator for Shamir Coefficient Vulnerability
GitHub Issue #24 - Calculate precise entropy impact for different scenarios

This script calculates the entropy impact of the coefficient vulnerability
for various threshold values and attack scenarios.
"""

import math
import argparse
import numpy as np
import matplotlib.pyplot as plt

def calculate_entropy_reduction(t, t_prime=None):
    """
    Calculate entropy reduction for Shamir coefficient vulnerability
    
    Args:
        t: Total threshold shares required
        t_prime: Shares known by attacker (default: t-1)
    
    Returns:
        tuple: (entropy_reduction_per_byte, remaining_entropy_per_byte)
    """
    if t_prime is None:
        t_prime = t - 1
    
    # Base entropy reduction per coefficient
    base_reduction = math.log2(256/255)
    
    # Total reduction for t-1 coefficients (excluding constant term)
    total_reduction = (t - 1) * base_reduction
    
    # Remaining entropy per byte
    remaining_entropy = 8.0 - total_reduction
    
    return total_reduction, remaining_entropy

def calculate_attack_complexity(secret_bits, t, t_prime=None):
    """
    Calculate the complexity of attacking a secret with given parameters
    
    Args:
        secret_bits: Size of secret in bits
        t: Threshold shares required  
        t_prime: Shares known by attacker
    
    Returns:
        dict: Attack complexity metrics
    """
    if t_prime is None:
        t_prime = t - 1
    
    # Calculate entropy reduction per byte
    reduction_per_byte, remaining_per_byte = calculate_entropy_reduction(t, t_prime)
    
    # Calculate total remaining entropy
    secret_bytes = secret_bits / 8
    total_remaining_entropy = remaining_per_byte * secret_bytes
    
    # Attack complexity
    brute_force_without_vuln = 2 ** secret_bits
    brute_force_with_vuln = 2 ** total_remaining_entropy
    
    # Advantage factor
    advantage_factor = brute_force_without_vuln / brute_force_with_vuln
    
    return {
        'secret_bits': secret_bits,
        'threshold_t': t,
        'known_shares': t_prime,
        'entropy_reduction_per_byte': reduction_per_byte,
        'remaining_entropy_per_byte': remaining_per_byte,
        'total_original_entropy': secret_bits,
        'total_remaining_entropy': total_remaining_entropy,
        'brute_force_original': brute_force_without_vuln,
        'brute_force_vulnerable': brute_force_with_vuln,
        'attack_advantage': advantage_factor,
        'security_reduction_percent': (1 - total_remaining_entropy/secret_bits) * 100
    }

def generate_impact_report(threshold_range=(2, 17), secret_sizes=[128, 192, 256]):
    """Generate comprehensive impact report"""
    
    print("=" * 80)
    print("SHAMIR COEFFICIENT VULNERABILITY - ENTROPY IMPACT ANALYSIS")
    print("=" * 80)
    
    print(f"\n📋 Analysis Parameters:")
    print(f"   Threshold range: {threshold_range[0]} to {threshold_range[1]}")
    print(f"   Secret sizes: {secret_sizes} bits")
    print(f"   Attack scenario: Attacker knows t-1 shares")
    
    # Detailed analysis for each secret size
    for secret_bits in secret_sizes:
        print(f"\n" + "="*50)
        print(f"🔐 SECRET SIZE: {secret_bits} BITS")
        print("="*50)
        
        print(f"{'t':<3} {'Entropy Loss':<12} {'Remaining':<10} {'Attack Advantage':<15} {'Security Loss':<12}")
        print("-" * 65)
        
        for t in range(threshold_range[0], threshold_range[1]):
            metrics = calculate_attack_complexity(secret_bits, t)
            
            print(f"{t:<3} "
                  f"{metrics['entropy_reduction_per_byte']:.4f} bits  "
                  f"{metrics['total_remaining_entropy']:>6.1f} bits  "
                  f"{metrics['attack_advantage']:>13.1e}  "
                  f"{metrics['security_reduction_percent']:>9.2f}%")
    
    # Critical thresholds analysis
    print(f"\n" + "="*50)
    print("🚨 CRITICAL THRESHOLD ANALYSIS")
    print("="*50)
    
    critical_thresholds = [3, 5, 10, 16]
    
    for t in critical_thresholds:
        print(f"\n📊 Threshold t={t} (Attacker knows {t-1} shares):")
        
        for secret_bits in secret_sizes:
            metrics = calculate_attack_complexity(secret_bits, t)
            
            print(f"   {secret_bits}-bit secret:")
            print(f"      Original security: 2^{secret_bits} = {metrics['brute_force_original']:.2e}")
            print(f"      Vulnerable security: 2^{metrics['total_remaining_entropy']:.1f} = {metrics['brute_force_vulnerable']:.2e}")
            print(f"      Attack advantage: {metrics['attack_advantage']:.2e}x easier")
            print(f"      Security loss: {metrics['security_reduction_percent']:.2f}%")

def create_entropy_impact_visualization():
    """Create visualization of entropy impact"""
    
    # Threshold values
    t_values = range(2, 21)
    secret_sizes = [128, 192, 256]
    
    plt.figure(figsize=(15, 12))
    
    # Plot 1: Entropy reduction per byte vs threshold
    plt.subplot(2, 2, 1)
    entropy_reductions = [calculate_entropy_reduction(t)[0] for t in t_values]
    plt.plot(t_values, entropy_reductions, 'ro-', linewidth=2, markersize=6)
    plt.title('Entropy Reduction per Byte vs Threshold')
    plt.xlabel('Threshold t')
    plt.ylabel('Entropy Reduction (bits/byte)')
    plt.grid(True, alpha=0.3)
    
    # Plot 2: Remaining entropy for different secret sizes
    plt.subplot(2, 2, 2)
    for secret_bits in secret_sizes:
        remaining_entropies = []
        for t in t_values:
            metrics = calculate_attack_complexity(secret_bits, t)
            remaining_entropies.append(metrics['total_remaining_entropy'])
        
        plt.plot(t_values, remaining_entropies, 'o-', linewidth=2, 
                label=f'{secret_bits}-bit secret', markersize=4)
    
    plt.title('Remaining Entropy vs Threshold')
    plt.xlabel('Threshold t')
    plt.ylabel('Remaining Entropy (bits)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # Plot 3: Attack advantage (log scale)
    plt.subplot(2, 2, 3)
    for secret_bits in secret_sizes:
        advantages = []
        for t in t_values:
            metrics = calculate_attack_complexity(secret_bits, t)
            advantages.append(math.log10(metrics['attack_advantage']))
        
        plt.plot(t_values, advantages, 'o-', linewidth=2, 
                label=f'{secret_bits}-bit secret', markersize=4)
    
    plt.title('Attack Advantage vs Threshold (log10 scale)')
    plt.xlabel('Threshold t')
    plt.ylabel('Attack Advantage (log10)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # Plot 4: Security reduction percentage
    plt.subplot(2, 2, 4)
    for secret_bits in secret_sizes:
        security_losses = []
        for t in t_values:
            metrics = calculate_attack_complexity(secret_bits, t)
            security_losses.append(metrics['security_reduction_percent'])
        
        plt.plot(t_values, security_losses, 'o-', linewidth=2, 
                label=f'{secret_bits}-bit secret', markersize=4)
    
    plt.title('Security Reduction vs Threshold')
    plt.xlabel('Threshold t')
    plt.ylabel('Security Reduction (%)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('entropy_impact_analysis.png', dpi=300, bbox_inches='tight')
    print(f"\n📊 Entropy impact plots saved as 'entropy_impact_analysis.png'")

def zeronights_specific_analysis():
    """Specific analysis for ZeroNights challenge parameters"""
    
    print(f"\n" + "="*60)
    print("🎯 ZERONIGHTS CHALLENGE SPECIFIC ANALYSIS")
    print("="*60)
    
    # ZeroNights uses 128-bit entropy (mnemonic seed)
    secret_bits = 128
    threshold = 3  # Need 3 shares, attacker has 2
    
    print(f"Challenge parameters:")
    print(f"   Secret size: {secret_bits} bits (mnemonic seed)")
    print(f"   Threshold: {threshold} shares required")
    print(f"   Attacker knowledge: 2 out of 3 shares")
    
    metrics = calculate_attack_complexity(secret_bits, threshold, t_prime=2)
    
    print(f"\nVulnerability impact:")
    print(f"   Entropy reduction per byte: {metrics['entropy_reduction_per_byte']:.6f} bits")
    print(f"   Total entropy reduction: {metrics['secret_bits'] - metrics['total_remaining_entropy']:.3f} bits")
    print(f"   Remaining entropy: {metrics['total_remaining_entropy']:.3f} bits")
    print(f"   Security reduction: {metrics['security_reduction_percent']:.3f}%")
    
    print(f"\nAttack complexity:")
    print(f"   Without vulnerability: 2^{secret_bits} = {metrics['brute_force_original']:.2e}")
    print(f"   With vulnerability: 2^{metrics['total_remaining_entropy']:.1f} = {metrics['brute_force_vulnerable']:.2e}")
    print(f"   Attack advantage: {metrics['attack_advantage']:.2e}x easier")
    
    # Practical impact
    original_years = metrics['brute_force_original'] / (10**12)  # Assuming 1 TH/s
    vulnerable_years = metrics['brute_force_vulnerable'] / (10**12)
    
    print(f"\nPractical timeline (1TH/s):")
    print(f"   Without vulnerability: {original_years:.2e} years")  
    print(f"   With vulnerability: {vulnerable_years:.2e} years")
    
    return metrics

def main():
    """Main entropy impact analysis"""
    parser = argparse.ArgumentParser(description='Calculate Shamir coefficient vulnerability entropy impact')
    parser.add_argument('--threshold', type=int, default=3, help='Threshold value (default: 3)')
    parser.add_argument('--secret-bits', type=int, default=128, help='Secret size in bits (default: 128)')
    parser.add_argument('--known-shares', type=int, help='Shares known by attacker (default: threshold-1)')
    parser.add_argument('--zeronights', action='store_true', help='Run ZeroNights specific analysis')
    parser.add_argument('--full-report', action='store_true', help='Generate full impact report')
    parser.add_argument('--visualize', action='store_true', help='Create visualization plots')
    
    args = parser.parse_args()
    
    print("🔥 SHAMIR COEFFICIENT VULNERABILITY - ENTROPY IMPACT CALCULATOR")
    print("📋 GitHub Issue #24 - Critical SSS Vulnerability")
    print()
    
    if args.zeronights:
        zeronights_specific_analysis()
    elif args.full_report:
        generate_impact_report()
        if args.visualize:
            try:
                create_entropy_impact_visualization()
            except ImportError:
                print("⚠️  matplotlib not available, skipping visualization")
    else:
        # Single calculation
        t_prime = args.known_shares if args.known_shares else args.threshold - 1
        metrics = calculate_attack_complexity(args.secret_bits, args.threshold, t_prime)
        
        print(f"Single Analysis:")
        print(f"   Secret: {args.secret_bits} bits")
        print(f"   Threshold: {args.threshold}")
        print(f"   Known shares: {t_prime}")
        print(f"   Entropy reduction: {metrics['entropy_reduction_per_byte']:.6f} bits/byte")
        print(f"   Remaining entropy: {metrics['total_remaining_entropy']:.3f} bits")
        print(f"   Attack advantage: {metrics['attack_advantage']:.2e}x")
        print(f"   Security loss: {metrics['security_reduction_percent']:.3f}%")
    
    print(f"\n✅ Analysis complete!")
    print(f"💰 This vulnerability analysis supports bounty claim for Issue #24")

if __name__ == "__main__":
    main()