decrypt.py

#!/usr/bin/env python3
"""
Decrypt and demodulate audio WAV files to recover encrypted text messages.
Supports AES-256 symmetric decryption and RSA asymmetric decryption with error correction decoding.
"""

import argparse
import hashlib
import os
import sys
import numpy as np
from scipy.io import wavfile
from scipy.signal import butter, filtfilt
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
import reedsolo
from cryptography.hazmat.primitives.asymmetric import padding as rsa_padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.backends import default_backend

# Audio parameters (must match encrypt.py)
SAMPLE_RATE = 44100
BIT_DURATION = 0.01

# Modulation schemes - frequency mappings
FREQS_FSK2 = [1200, 2400]
FREQS_FSK4 = [1200, 1800, 2400, 3000]
FREQS_FSK8 = [1000, 1400, 1800, 2200, 2600, 3000, 3400, 3800]
FREQS_FSK16 = [800, 1000, 1200, 1400, 1600, 1800, 2000, 2200,
               2400, 2600, 2800, 3000, 3200, 3400, 3600, 3800]

MODULATION_SCHEMES = {
    'fsk2': (FREQS_FSK2, 1),
    'fsk4': (FREQS_FSK4, 2),
    'fsk8': (FREQS_FSK8, 3),
    'fsk16': (FREQS_FSK16, 4)
}

# Reed-Solomon error correction
RS_DECODER = reedsolo.RSCodec(64)


def derive_key_from_password(password: str, context: str = 'AudioEncryptKey2025') -> bytes:
    """Derive a 256-bit key from a password using PBKDF2 with custom context."""
    return hashlib.pbkdf2_hmac('sha256', password.encode(), context.encode(), 100000, dklen=32)


def derive_iv_from_password(password: str, context: str = 'AudioEncryptIV2025') -> bytes:
    """Derive a 128-bit IV from a password using PBKDF2 with custom context."""
    return hashlib.pbkdf2_hmac('sha256', password.encode(), context.encode(), 50000, dklen=16)


def extract_rsa_metadata(data: bytes) -> tuple:
    """Extract RSA metadata from the beginning of data.
    Returns: (mode, iv, encrypted_aes_key, remaining_data)
    """
    if len(data) < 1:
        return None, None, None, data
    
    mode_byte = data[0]
    if mode_byte != 1:  # Not RSA mode
        return None, None, None, data
    
    # RSA format: [mode(1)][iv_len(2)][iv][key_len(2)][encrypted_key][data]
    if len(data) < 5:
        raise Exception("Invalid RSA metadata format")
    
    iv_len = int.from_bytes(data[1:3], 'big')
    iv = data[3:3+iv_len]
    
    key_len_pos = 3 + iv_len
    key_len = int.from_bytes(data[key_len_pos:key_len_pos+2], 'big')
    encrypted_aes_key = bytes(data[key_len_pos+2:key_len_pos+2+key_len])  # Convert to bytes
    
    remaining_data = bytes(data[key_len_pos+2+key_len:])  # Convert to bytes
    
    return 'asymmetric', bytes(iv), encrypted_aes_key, remaining_data


def decrypt_aes_key_with_rsa(encrypted_key: bytes, private_key) -> bytes:
    """Decrypt AES key using RSA private key."""
    aes_key = private_key.decrypt(
        encrypted_key,
        rsa_padding.OAEP(
            mgf=rsa_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return aes_key


def bandpass_filter(data, lowcut, highcut, fs, order=5):
    """Apply bandpass filter to audio data."""
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return filtfilt(b, a, data)


def demodulate_fsk(audio: np.ndarray, sample_rate: int, modulation: str = 'fsk2') -> list:
    """Demodulate FSK audio to bits.
    
    Args:
        audio: Audio signal
        sample_rate: Sample rate
        modulation: 'fsk2', 'fsk4', 'fsk8', or 'fsk16'
    """
    if modulation not in MODULATION_SCHEMES:
        raise ValueError(f"Unknown modulation: {modulation}")
    
    frequencies, bits_per_symbol = MODULATION_SCHEMES[modulation]
    samples_per_symbol = int(sample_rate * BIT_DURATION)
    
    # Skip preamble (100 symbols)
    preamble_samples = 100 * samples_per_symbol
    if len(audio) > preamble_samples:
        audio = audio[preamble_samples:]
    
    # Create bandpass filters for each frequency
    filtered_audio = []
    for freq in frequencies:
        bandwidth = 200 if modulation == 'fsk2' else 150  # Tighter for higher FSK
        filtered = bandpass_filter(audio, freq - bandwidth, freq + bandwidth, sample_rate)
        filtered_audio.append(filtered)
    
    # Demodulate by comparing energy in each frequency band
    num_symbols = len(audio) // samples_per_symbol
    bits = []
    
    for i in range(num_symbols):
        start_idx = i * samples_per_symbol
        end_idx = start_idx + samples_per_symbol
        
        if end_idx > len(audio):
            break
        
        # Calculate energy in each frequency band
        energies = []
        for filtered in filtered_audio:
            energy = np.sum(filtered[start_idx:end_idx] ** 2)
            energies.append(energy)
        
        # Find frequency with maximum energy
        symbol_value = np.argmax(energies)
        
        # Convert symbol to bits
        symbol_bits = []
        for _ in range(bits_per_symbol):
            symbol_bits.insert(0, symbol_value & 1)
            symbol_value >>= 1
        bits.extend(symbol_bits)
    
    return bits


def bits_to_bytes(bits: list) -> bytes:
    """Convert list of bits to bytes."""
    # Pad bits to multiple of 8
    while len(bits) % 8 != 0:
        bits.append(0)
    
    data = bytearray()
    for i in range(0, len(bits), 8):
        byte = 0
        for j in range(8):
            byte = (byte << 1) | bits[i + j]
        data.append(byte)
    
    return bytes(data)


def remove_error_correction(data: bytes, verbose: bool = True) -> bytes:
    """Remove Reed-Solomon error correction and recover original data."""
    try:
        result = RS_DECODER.decode(data)
        # Handle both tuple and bytes return types
        if isinstance(result, tuple):
            return result[0]
        return result
    except reedsolo.ReedSolomonError as e:
        if verbose:
            print(f"Warning: Error correction failed - data may be corrupted: {e}")
        # Try to return partially corrected data
        try:
            result = RS_DECODER.decode(data, erase_pos=None)
            if isinstance(result, tuple):
                return result[0]
            return result
        except:
            raise Exception("Data is too corrupted to recover")


def decrypt_message(encrypted_data: bytes, key: bytes, iv: bytes) -> str:
    """Decrypt message using AES-256-CBC."""
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted = cipher.decrypt(encrypted_data)
    unpadded = unpad(decrypted, AES.block_size)
    return unpadded.decode('utf-8')


def main():
    parser = argparse.ArgumentParser(
        description='Decrypt and demodulate audio WAV files to recover text messages',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Symmetric decryption using a password:
  python decrypt.py -input encrypted.wav -password "mysecret"
  
  # Symmetric decryption using explicit key:
  python decrypt.py -input encrypted.wav -key 0123456789abcdef...
  
  # Asymmetric decryption using RSA private key:
  python decrypt.py -input encrypted.wav -privatekey private_key.pem
        """
    )
    
    parser.add_argument('-input', '-i', required=True, help='Input WAV file to decrypt')
    parser.add_argument('-password', '-p', help='Password for key derivation (symmetric)')
    parser.add_argument('-key', '-k', help='AES-256 key (64 hex characters, symmetric)')
    parser.add_argument('-iv', help='Initialization vector (32 hex characters, required with -key)')
    parser.add_argument('-privatekey', '-priv', help='RSA private key file (asymmetric)')
    parser.add_argument('-context', '-c', help='Custom context string for password derivation (must match encryption context)')
    parser.add_argument('--modulation', choices=['fsk2', 'fsk4', 'fsk8', 'fsk16'], default='fsk2',
                        help='Modulation scheme (must match encryption, default: fsk2)')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.input):
        print(f"Error: Input file not found: {args.input}")
        sys.exit(1)
    
    # Load audio file
    print(f"Loading audio file: {args.input}")
    sample_rate, audio_data = wavfile.read(args.input)
    
    if sample_rate != SAMPLE_RATE:
        print(f"Warning: Sample rate mismatch (expected {SAMPLE_RATE}, got {sample_rate})")
    
    # Convert to float
    audio = audio_data.astype(np.float32) / 32768.0
    print(f"Audio duration: {len(audio) / sample_rate:.2f} seconds")
    
    # Demodulate FSK
    print(f"\nDemodulating {args.modulation.upper()} audio...")
    bits = demodulate_fsk(audio, sample_rate, modulation=args.modulation)
    print(f"Demodulated {len(bits)} bits")
    
    # Convert bits to bytes
    protected_data = bits_to_bytes(bits)
    print(f"Received {len(protected_data)} bytes")
    
    # Remove error correction
    print("\nApplying error correction...")
    try:
        decrypted_with_metadata = remove_error_correction(protected_data, verbose=True)
        print(f"✓ Error correction successful")
    except Exception as e:
        print(f"✗ Error correction failed: {e}")
        sys.exit(1)
    
    # Check if this is RSA mode (has embedded metadata)
    mode, iv_from_meta, encrypted_aes_key, encrypted_data = extract_rsa_metadata(decrypted_with_metadata)
    
    # Determine decryption key and IV
    if mode == 'asymmetric':
        # RSA asymmetric decryption
        if not args.privatekey:
            print("\nError: This message was encrypted with RSA. Provide -privatekey to decrypt.")
            sys.exit(1)
        
        print("\nDetected RSA asymmetric encryption (hybrid mode)")
        print(f"IV (hex): {iv_from_meta.hex()}")
        
        # Load private key
        try:
            with open(args.privatekey, 'rb') as f:
                private_key = serialization.load_pem_private_key(
                    f.read(),
                    password=None,
                    backend=default_backend()
                )
        except Exception as e:
            print(f"Error loading private key: {e}")
            sys.exit(1)
        
        # Decrypt the AES key with RSA private key
        try:
            key = decrypt_aes_key_with_rsa(encrypted_aes_key, private_key)
            iv = iv_from_meta
            print(f"✓ AES key decrypted using RSA private key")
        except Exception as e:
            print(f"Error decrypting AES key: {e}")
            print("Make sure you're using the correct private key.")
            sys.exit(1)
    
    elif args.password:
        # Password-based decryption with derived IV
        context = args.context if args.context else 'AudioEncrypt2025'
        key = derive_key_from_password(args.password, context + '_Key')
        iv = derive_iv_from_password(args.password, context + '_IV')
        encrypted_data = decrypted_with_metadata  # No metadata to strip
        print("\nUsing password-based decryption (AES-256-CBC)")
        print(f"Key and IV derived from password with context: '{context}'")
    elif args.key:
        # Explicit key decryption - IV must be provided
        if not args.iv:
            print("\nError: When using -key, you must also provide -iv")
            sys.exit(1)
        
        try:
            key = bytes.fromhex(args.key)
            if len(key) != 32:
                print("Error: Key must be 32 bytes (64 hex characters) for AES-256")
                sys.exit(1)
        except ValueError:
            print("Error: Invalid hex key")
            sys.exit(1)
        
        try:
            iv = bytes.fromhex(args.iv)
            if len(iv) != 16:
                print("Error: IV must be 16 bytes (32 hex characters)")
                sys.exit(1)
        except ValueError:
            print("Error: Invalid hex IV")
            sys.exit(1)
        
        encrypted_data = decrypted_with_metadata  # No metadata to strip
        print("\nUsing explicit key and IV (AES-256-CBC)")
    else:
        print("\nError: Provide -password, -key + -iv, or -privatekey to decrypt")
        sys.exit(1)
    
    # Decrypt message
    print("\nDecrypting message...")
    try:
        message = decrypt_message(encrypted_data, key, iv)
        print(f"\n{'='*60}")
        print(f"DECRYPTED MESSAGE:")
        print(f"{'='*60}")
        print(f"{message}")
        print(f"{'='*60}")
    except Exception as e:
        print(f"✗ Decryption failed: {e}")
        print("Possible causes:")
        print("  - Wrong password/key")
        print("  - Corrupted audio data")
        print("  - Transmission errors beyond error correction capacity")
        sys.exit(1)


if __name__ == '__main__':
    main()