test.py

#!/usr/bin/env python3
"""
Test script to compare different bitrates for radio transmission.
Tests encryption/decryption with various transmission conditions at different speeds.
"""

import argparse
import os
import sys
import numpy as np
from scipy.io import wavfile
from scipy.signal import resample
from Crypto.Random import get_random_bytes

# We'll monkey-patch the bitrate settings
import encrypt
import decrypt

# Original values
ORIGINAL_BIT_DURATION = 0.01
ORIGINAL_SAMPLE_RATE = 44100


def set_bitrate(bps, modulation='fsk2'):
    """
    Set the bitrate for testing.
    
    Common real-world radio bitrates:
    - HF Radio (SSB): 50-300 bps (we'll use 100-300 bps range)
    - VHF Packet Radio: 1200 bps (Bell 202 modem standard)
    - VHF/UHF Digital: 9600 bps
    - Amateur Radio APRS: 1200 bps
    - Marine VHF DSC: 1200 bps
    
    Args:
        bps: Target bits per second
        modulation: 'fsk2', 'fsk4', 'fsk8', or 'fsk16'
    """
    # Calculate bits per symbol based on modulation
    bits_per_symbol = {
        'fsk2': 1,
        'fsk4': 2,
        'fsk8': 3,
        'fsk16': 4
    }[modulation]
    
    # Calculate required symbol rate to achieve target bitrate
    symbols_per_second = bps / bits_per_symbol
    symbol_duration = 1.0 / symbols_per_second
    
    # Update both modules (BIT_DURATION is actually symbol duration)
    encrypt.BIT_DURATION = symbol_duration
    decrypt.BIT_DURATION = symbol_duration
    
    actual_bps = symbols_per_second * bits_per_symbol
    print(f"Bitrate set to: {actual_bps:.0f} bps ({symbols_per_second:.0f} symbols/sec × {bits_per_symbol} bits/symbol, symbol duration: {symbol_duration*1000:.2f} ms)")


def save_wav(filename, audio, sample_rate=44100):
    """Save audio array as WAV file."""
    audio_int16 = np.int16(audio * 32767)
    wavfile.write(filename, sample_rate, audio_int16)


def add_white_noise(audio, snr_db):
    """Add white noise to simulate noisy transmission."""
    signal_power = np.mean(audio ** 2)
    snr_linear = 10 ** (snr_db / 10)
    noise_power = signal_power / snr_linear
    noise = np.sqrt(noise_power) * np.random.randn(len(audio))
    return audio + noise


def downsample_and_upsample(audio, factor):
    """Simulate bitrate reduction by downsampling and upsampling."""
    downsampled = resample(audio, len(audio) // factor)
    upsampled = resample(downsampled, len(audio))
    return upsampled


def test_at_bitrate(bitrate, message, password, output_dir="test_output", mode="password", rsa_keys=None, modulation="fsk2", bandwidth_factor=None):
    """Test encryption/decryption at a specific bitrate with various conditions.
    
    Args:
        bitrate: Bits per second
        message: Message to encrypt
        password: Password for symmetric encryption
        output_dir: Directory to save test files
        mode: 'password' or 'rsa'
        rsa_keys: Tuple of (public_key, private_key) for RSA mode
        modulation: 'fsk2', 'fsk4', 'fsk8', or 'fsk16'
        bandwidth_factor: Optional int (2, 4, 8) to apply bandwidth limiting to noise tests
    """
    from encrypt import encrypt_message, add_error_correction, bytes_to_bits, modulate_to_audio, derive_key_from_password, derive_iv_from_password, encrypt_aes_key_with_rsa, prepend_metadata_to_data
    from decrypt import demodulate_fsk, bits_to_bytes, remove_error_correction, decrypt_message, extract_rsa_metadata, decrypt_aes_key_with_rsa
    
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"\n{'='*70}")
    print(f"TESTING AT {bitrate} BPS ({modulation.upper()})")
    print(f"{'='*70}")
    
    set_bitrate(bitrate, modulation)
    
    # Prepare encryption based on mode
    if mode == "rsa":
        # RSA mode - generate random AES key and encrypt it with public key
        public_key, private_key = rsa_keys
        key = get_random_bytes(32)
        iv = get_random_bytes(16)
        encrypted_aes_key = encrypt_aes_key_with_rsa(key, public_key)
        
        encrypted_data = encrypt_message(message, key, iv)
        # Prepend RSA metadata
        encrypted_data = prepend_metadata_to_data(encrypted_data, 'rsa', iv, encrypted_aes_key)
    else:
        # Password mode - derive key and IV from password
        context = 'AudioEncrypt2025'
        key = derive_key_from_password(password, context + '_Key')
        iv = derive_iv_from_password(password, context + '_IV')
        encrypted_data = encrypt_message(message, key, iv)
    protected_data = add_error_correction(encrypted_data)
    bits = bytes_to_bits(protected_data)
    audio_clean = modulate_to_audio(bits, modulation=modulation)
    
    duration = len(audio_clean) / encrypt.SAMPLE_RATE
    
    print(f"Message: '{message}'")
    print(f"Encrypted size: {len(encrypted_data)} bytes")
    print(f"With error correction: {len(protected_data)} bytes")
    print(f"Total bits: {len(bits)}")
    print(f"Audio duration: {duration:.2f} seconds")
    print(f"Efficiency: {len(message) / duration:.1f} chars/sec")
    
    results = {}
    test_scenarios = [
        ("clean", "Clean Signal", None),
        ("noise_20db", "Noise SNR=20dB", lambda a: add_white_noise(a, 20)),
        ("noise_15db", "Noise SNR=15dB", lambda a: add_white_noise(a, 15)),
        ("noise_10db", "Noise SNR=10dB", lambda a: add_white_noise(a, 10)),
        ("noise_5db", "Noise SNR=5dB", lambda a: add_white_noise(a, 5)),
        ("noise_0db", "Noise SNR=0dB", lambda a: add_white_noise(a, 0)),
        ("noise_-5db", "Noise SNR=-5dB", lambda a: add_white_noise(a, -5)),
        ("noise_-10db", "Noise SNR=-10dB", lambda a: add_white_noise(a, -10)),
        ("noise_-15db", "Noise SNR=-15dB (Extreme)", lambda a: add_white_noise(a, -15)),
    ]
    
    # Add bandwidth-limited noise scenarios if bandwidth_factor specified
    if bandwidth_factor is not None:
        bw_scenarios = [
            (f"noise_20db_bw{bandwidth_factor}x", f"Noise SNR=20dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, 20), bandwidth_factor)),
            (f"noise_15db_bw{bandwidth_factor}x", f"Noise SNR=15dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, 15), bandwidth_factor)),
            (f"noise_10db_bw{bandwidth_factor}x", f"Noise SNR=10dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, 10), bandwidth_factor)),
            (f"noise_5db_bw{bandwidth_factor}x", f"Noise SNR=5dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, 5), bandwidth_factor)),
            (f"noise_0db_bw{bandwidth_factor}x", f"Noise SNR=0dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, 0), bandwidth_factor)),
            (f"noise_-5db_bw{bandwidth_factor}x", f"Noise SNR=-5dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, -5), bandwidth_factor)),
            (f"noise_-10db_bw{bandwidth_factor}x", f"Noise SNR=-10dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, -10), bandwidth_factor)),
            (f"noise_-15db_bw{bandwidth_factor}x", f"Noise SNR=-15dB + BW/{bandwidth_factor}", 
             lambda a: downsample_and_upsample(add_white_noise(a, -15), bandwidth_factor)),
        ]
        test_scenarios.extend(bw_scenarios)
    
    for filename_suffix, test_name, processor in test_scenarios:
        audio = audio_clean.copy()
        
        if processor:
            audio = processor(audio)
        
        # Normalize
        max_val = np.max(np.abs(audio))
        if max_val > 1.0:
            audio = audio / max_val
        
        # Save to file
        output_file = os.path.join(output_dir, f"{bitrate}bps_{filename_suffix}.wav")
        save_wav(output_file, audio)
        
        # Test decryption
        try:
            demod_bits = demodulate_fsk(audio, encrypt.SAMPLE_RATE, modulation=modulation)
            received_data = bits_to_bytes(demod_bits)
            
            original_bits = bits[:len(demod_bits)]
            errors = sum(a != b for a, b in zip(original_bits, demod_bits))
            ber = errors / len(original_bits) if len(original_bits) > 0 else 0
            
            decrypted_data_with_metadata = remove_error_correction(received_data, verbose=False)
            
            # Handle RSA mode metadata extraction
            if mode == "rsa":
                _, iv_from_meta, encrypted_aes_key_received, decrypted_data = extract_rsa_metadata(decrypted_data_with_metadata)
                key = decrypt_aes_key_with_rsa(encrypted_aes_key_received, private_key)
                iv = iv_from_meta
            else:
                decrypted_data = decrypted_data_with_metadata
            
            decrypted_message = decrypt_message(decrypted_data, key, iv)
            
            success = (decrypted_message == message)
            results[test_name] = {
                'success': success,
                'ber': ber,
                'errors': errors,
                'file': output_file
            }
            
            status = "✓ PASS" if success else "✗ FAIL"
            print(f"  {test_name:.<35} {status} (BER: {ber*100:.2f}%)")
            
        except Exception as e:
            results[test_name] = {
                'success': False,
                'ber': None,
                'errors': None,
                'file': output_file
            }
            print(f"  {test_name:.<35} ✗ FAIL ({str(e)[:30]}...)")
    
    return results


def main():
    parser = argparse.ArgumentParser(
        description='Test encryption/decryption at different bitrates',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument('--mode', '-m', choices=['password', 'rsa'], default='password',
                        help='Encryption mode: password (symmetric) or rsa (asymmetric)')
    parser.add_argument('--modulation', choices=['fsk2', 'fsk4', 'fsk8', 'fsk16'], default='fsk2',
                        help='Modulation scheme: fsk2=1bit/symbol, fsk4=2bits/symbol, fsk8=3bits/symbol, fsk16=4bits/symbol (default: fsk2)')
    parser.add_argument('--bandwidth', '-bw', type=int, choices=[2, 4, 8], default=None,
                        help='Apply bandwidth limiting (bitcrushing) to noise tests: 2, 4, or 8x downsampling (optional)')
    args = parser.parse_args()
    
    print("="*70)
    print("BITRATE COMPARISON TEST SUITE")
    print("="*70)
    print(f"\nMode: {args.mode.upper()}")
    print(f"Modulation: {args.modulation.upper()}")
    if args.bandwidth:
        print(f"Bandwidth Limiting: {args.bandwidth}x downsampling (bitcrushing applied to noise tests)")
    if args.mode == 'password':
        print("Note: All symmetric tests with same password produce identical audio")
    else:
        print("Note: RSA mode generates random AES keys, so audio differs each run")
    print("\nReal-world radio transmission bitrates:")
    print("  - HF SSB Digital modes: 50-300 bps")
    print("  - VHF Packet Radio (Bell 202): 1200 bps") 
    print("  - VHF/UHF APRS: 1200 bps")
    print("  - VHF/UHF Fast modes: 9600 bps")
    print("  - Modern digital voice: 2400-4800 bps")
    print("\nWe'll test a range from slow HF to fast VHF rates.\n")
    
    test_message = "Hello World! Testing 123."
    test_password = "radio_test_2025"
    output_dir = "test_output"
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Generate RSA keys if needed
    rsa_keys = None
    if args.mode == 'rsa':
        from cryptography.hazmat.primitives.asymmetric import rsa
        from cryptography.hazmat.backends import default_backend
        
        print("Generating RSA key pair for testing...")
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        public_key = private_key.public_key()
        rsa_keys = (public_key, private_key)
        print("✓ RSA keys generated\n")
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Test different bitrates
    bitrates = [
        (100, "Ultra-slow HF SSB"),
        (200, "Very Slow HF"),
        (300, "Slow HF Digital"),
        (400, "Medium-Slow HF"),
        (500, "Medium HF"),
        (550, "Medium+ HF"),
        (600, "Medium-Fast HF/VHF"),
        (625, "Transition 625"),
        (650, "Transition 650"),
        (675, "Transition 675"),
        (700, "Fast HF/VHF"),
        (1200, "VHF Packet Radio (Standard)"),
        (2400, "Very Fast VHF"),
        (9600, "Maximum VHF/UHF"),
    ]
    
    all_results = {}
    
    for bitrate, description in bitrates:
        print(f"\n{description}")
        results = test_at_bitrate(bitrate, test_message, test_password, output_dir, 
                                 mode=args.mode, rsa_keys=rsa_keys, modulation=args.modulation,
                                 bandwidth_factor=args.bandwidth)
        all_results[bitrate] = results
    
    # Summary
    print(f"\n{'='*70}")
    print("SUMMARY - Success rate by bitrate")
    print(f"{'='*70}")
    
    for bitrate, description in bitrates:
        results = all_results[bitrate]
        passed = sum(1 for r in results.values() if r['success'])
        total = len(results)
        percentage = (passed / total * 100) if total > 0 else 0
        
        print(f"{bitrate:>5} bps ({description:.<30}) {passed}/{total} ({percentage:.0f}%)")

    
    print(f"\n{'='*70}")
    print(f"All test audio files saved to: {os.path.abspath(output_dir)}")
    print(f"{'='*70}")
    
    # Restore original values
    encrypt.BIT_DURATION = ORIGINAL_BIT_DURATION
    decrypt.BIT_DURATION = ORIGINAL_BIT_DURATION


if __name__ == '__main__':
    main()