encrypt.py

#!/usr/bin/env python3
"""
Encrypt and modulate text messages to audio WAV files for transmission over radio.
Supports AES-256 symmetric encryption and RSA asymmetric encryption with error correction coding.
"""

import argparse
import hashlib
import os
import sys
import numpy as np
from scipy.io import wavfile
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad
import reedsolo
from cryptography.hazmat.primitives.asymmetric import rsa, padding as rsa_padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.backends import default_backend

# Audio parameters
SAMPLE_RATE = 44100  # Hz
BIT_DURATION = 0.01  # seconds per symbol (100 symbols/sec)
AMPLITUDE = 0.8

# Modulation schemes - frequency mappings
# FSK2 (Binary): 1 bit per symbol, 2 frequencies
FREQS_FSK2 = [1200, 2400]  # 0, 1

# FSK4 (Quadrature): 2 bits per symbol, 4 frequencies  
FREQS_FSK4 = [1200, 1800, 2400, 3000]  # 00, 01, 10, 11

# FSK8: 3 bits per symbol, 8 frequencies
FREQS_FSK8 = [1000, 1400, 1800, 2200, 2600, 3000, 3400, 3800]  # 000-111

# FSK16: 4 bits per symbol, 16 frequencies
FREQS_FSK16 = [800, 1000, 1200, 1400, 1600, 1800, 2000, 2200,
               2400, 2600, 2800, 3000, 3200, 3400, 3600, 3800]  # 0000-1111

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

# Default modulation
DEFAULT_MODULATION = 'fsk2'
CURRENT_MODULATION = DEFAULT_MODULATION

# Reed-Solomon error correction (can correct up to 32 symbol errors)
RS_ENCODER = 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.
    
    Args:
        password: The password to derive from
        context: Custom context string (allows different groups/networks to use different namespaces)
    """
    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.
    
    Args:
        password: The password to derive from
        context: Custom context string (allows different groups/networks to use different namespaces)
    """
    return hashlib.pbkdf2_hmac('sha256', password.encode(), context.encode(), 50000, dklen=16)


def encrypt_message(message: str, key: bytes, iv: bytes) -> bytes:
    """Encrypt message using AES-256-CBC."""
    cipher = AES.new(key, AES.MODE_CBC, iv)
    padded_message = pad(message.encode('utf-8'), AES.block_size)
    return cipher.encrypt(padded_message)


def add_error_correction(data: bytes) -> bytes:
    """Add Reed-Solomon error correction codes."""
    return RS_ENCODER.encode(data)


def bytes_to_bits(data: bytes) -> list:
    """Convert bytes to list of bits."""
    bits = []
    for byte in data:
        for i in range(7, -1, -1):
            bits.append((byte >> i) & 1)
    return bits


def modulate_to_audio(bits: list, modulation: str = 'fsk2') -> np.ndarray:
    """Modulate bits to audio using FSK (Frequency Shift Keying).
    
    Args:
        bits: List of bits to modulate
        modulation: 'fsk2', 'fsk4', 'fsk8', or 'fsk16'
    """
    if modulation not in MODULATION_SCHEMES:
        raise ValueError(f"Unknown modulation: {modulation}. Use: {list(MODULATION_SCHEMES.keys())}")
    
    frequencies, bits_per_symbol = MODULATION_SCHEMES[modulation]
    samples_per_symbol = int(SAMPLE_RATE * BIT_DURATION)
    
    # Pad bits to multiple of bits_per_symbol
    while len(bits) % bits_per_symbol != 0:
        bits.append(0)
    
    # Group bits into symbols
    num_symbols = len(bits) // bits_per_symbol
    symbols = []
    for i in range(num_symbols):
        symbol_bits = bits[i * bits_per_symbol:(i + 1) * bits_per_symbol]
        # Convert bits to symbol index
        symbol_value = 0
        for bit in symbol_bits:
            symbol_value = (symbol_value << 1) | bit
        symbols.append(symbol_value)
    
    # Generate audio
    total_samples = len(symbols) * samples_per_symbol
    audio = np.zeros(total_samples)
    t = np.linspace(0, BIT_DURATION, samples_per_symbol, endpoint=False)
    
    for i, symbol in enumerate(symbols):
        freq = frequencies[symbol]
        start_idx = i * samples_per_symbol
        end_idx = start_idx + samples_per_symbol
        audio[start_idx:end_idx] = AMPLITUDE * np.sin(2 * np.pi * freq * t)
    
    # Add preamble (alternating symbols for synchronization)
    preamble_symbols = [0, len(frequencies) - 1] * 50  # 100 symbols
    preamble_audio = np.zeros(len(preamble_symbols) * samples_per_symbol)
    for i, symbol in enumerate(preamble_symbols):
        freq = frequencies[symbol]
        start_idx = i * samples_per_symbol
        end_idx = start_idx + samples_per_symbol
        preamble_audio[start_idx:end_idx] = AMPLITUDE * np.sin(2 * np.pi * freq * t)
    
    return np.concatenate([preamble_audio, audio])


def encrypt_aes_key_with_rsa(aes_key: bytes, public_key) -> bytes:
    """Encrypt AES key using RSA public key (hybrid encryption)."""
    encrypted_key = public_key.encrypt(
        aes_key,
        rsa_padding.OAEP(
            mgf=rsa_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return encrypted_key


def prepend_metadata_to_data(data: bytes, mode: str, iv: bytes, encrypted_aes_key: bytes = None) -> bytes:
    """Prepend metadata to encrypted data for RSA mode.
    Format: [mode_byte(1)][iv_length(2)][iv][key_length(2)][encrypted_key][data]
    """
    if mode == 'rsa':
        mode_byte = b'\x01'
        iv_len = len(iv).to_bytes(2, 'big')
        key_len = len(encrypted_aes_key).to_bytes(2, 'big')
        return mode_byte + iv_len + iv + key_len + encrypted_aes_key + data
    else:
        # No metadata needed for password/key modes
        return data





def main():
    parser = argparse.ArgumentParser(
        description='Encrypt and modulate text messages to audio WAV files',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Symmetric encryption using a password:
  python encrypt.py -message "Secret message" -password "mysecret" -output encrypted.wav
  
  # Symmetric encryption using explicit key and IV (hex):
  python encrypt.py -message "Secret message" -key 0123456789abcdef... -iv fedcba9876543210... -output encrypted.wav
  
  # Asymmetric encryption using RSA public key:
  python encrypt.py -message "Secret message" -publickey public_key.pem -output encrypted.wav
  
  # Generate RSA key pair:
  python encrypt.py -genkeys -keyname my_keypair
        """
    )
    
    parser.add_argument('-message', '-m', help='Message to encrypt')
    parser.add_argument('-output', '-o', default='encrypted.wav', help='Output WAV file (default: encrypted.wav)')
    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)')
    parser.add_argument('-publickey', '-pub', help='RSA public key file (asymmetric)')
    parser.add_argument('-genkeys', action='store_true', help='Generate RSA key pair')
    parser.add_argument('-keyname', default='keypair', help='Name for generated key files (default: keypair)')
    parser.add_argument('-context', '-c', help='Custom context string for password derivation (optional, allows different networks to use unique namespaces)')
    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)')
    
    args = parser.parse_args()
    
    # Handle key generation
    if args.genkeys:
        print("Generating RSA key pair (2048-bit)...")
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        public_key = private_key.public_key()
        
        # Save private key
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        private_filename = f"{args.keyname}_private.pem"
        with open(private_filename, 'wb') as f:
            f.write(private_pem)
        print(f"✓ Private key saved: {private_filename}")
        
        # Save public key
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        public_filename = f"{args.keyname}_public.pem"
        with open(public_filename, 'wb') as f:
            f.write(public_pem)
        print(f"✓ Public key saved: {public_filename}")
        print(f"\n⚠ Keep the private key secure!")
        print(f"Share the public key with anyone who wants to send you encrypted messages.")
        return
    
    if not args.message:
        print("Error: -message is required (unless using -genkeys)")
        sys.exit(1)
    
    # Determine encryption key and IV
    encrypted_aes_key = None
    
    if args.publickey:
        # RSA asymmetric encryption (hybrid: RSA for AES key, AES for message)
        print("Using RSA asymmetric encryption (hybrid mode)")
        
        # Load public key
        try:
            with open(args.publickey, 'rb') as f:
                public_key = serialization.load_pem_public_key(
                    f.read(),
                    backend=default_backend()
                )
        except Exception as e:
            print(f"Error loading public key: {e}")
            sys.exit(1)
        
        # Generate random AES key and IV
        key = get_random_bytes(32)  # AES-256
        iv = get_random_bytes(16)
        
        # Encrypt the AES key with RSA public key
        encrypted_aes_key = encrypt_aes_key_with_rsa(key, public_key)
        
        key_mode = 'rsa'
        print(f"Generated AES-256 key (will be encrypted with RSA)")
        print(f"IV (hex): {iv.hex()}")
        
    elif args.password:
        # Derive key and IV from password (no metadata needed!)
        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')
        print("Using password-based encryption (AES-256-CBC)")
        print(f"Key and IV derived from password with context: '{context}'")
        key_mode = 'password'
    elif args.key:
        # Use provided key and IV
        if not args.iv:
            print("Error: When using -key, you must also provide -iv")
            print("This ensures both parties use the same IV for decryption.")
            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)
        
        print("Using explicit key and IV (AES-256-CBC)")
        key_mode = 'key'
    else:
        print("Error: Provide -password, -key + -iv, or -publickey")
        sys.exit(1)
    
    print(f"\nEncrypting message: '{args.message}'")
    
    # Encrypt the message
    encrypted_data = encrypt_message(args.message, key, iv)
    print(f"Encrypted data size: {len(encrypted_data)} bytes")
    
    # For RSA mode, prepend metadata to encrypted data
    if key_mode == 'rsa':
        encrypted_data = prepend_metadata_to_data(encrypted_data, 'rsa', iv, encrypted_aes_key)
        print(f"With embedded RSA metadata: {len(encrypted_data)} bytes")
    
    # Add error correction
    protected_data = add_error_correction(encrypted_data)
    print(f"With error correction: {len(protected_data)} bytes")
    
    # Convert to bits
    bits = bytes_to_bits(protected_data)
    print(f"Total bits to transmit: {len(bits)}")
    
    # Modulate to audio
    print(f"\nModulating to audio ({args.modulation.upper()})...")
    audio = modulate_to_audio(bits, modulation=args.modulation)
    
    # Calculate actual bitrate
    _, bits_per_symbol = MODULATION_SCHEMES[args.modulation]
    symbol_rate = 1.0 / BIT_DURATION
    actual_bitrate = symbol_rate * bits_per_symbol
    
    # Normalize and convert to int16
    audio_int16 = np.int16(audio * 32767)
    
    # Save WAV file
    wavfile.write(args.output, SAMPLE_RATE, audio_int16)
    print(f"\n✓ Audio file saved: {args.output}")
    print(f"  Duration: {len(audio) / SAMPLE_RATE:.2f} seconds")
    print(f"  Sample rate: {SAMPLE_RATE} Hz")
    print(f"  Modulation: {args.modulation.upper()} ({bits_per_symbol} bits/symbol)")
    print(f"  Symbol rate: {symbol_rate:.0f} symbols/sec")
    print(f"  Bitrate: {actual_bitrate:.0f} bps")
    
    # Mode-specific messages
    if key_mode == 'rsa':
        print(f"\n✓ Message encrypted with RSA public key")
        print(f"Only the holder of the private key can decrypt this message.")
        print(f"Metadata embedded in audio - no separate file needed!")
    elif key_mode == 'password':
        print(f"\n✓ Password-based encryption complete")
        print(f"No metadata file needed - just share the WAV and password!")
    else:  # key mode
        print(f"\n⚠ Keep your key and IV safe:")
        print(f"  Key: {key.hex()}")
        print(f"  IV:  {iv.hex()}")


if __name__ == '__main__':
    main()