test_ciphers.py

#!/usr/bin/env python3
"""
Test script for the Ciphers module to verify all functions work correctly.
"""

import numpy as np
import Ciphers as cp


def test_caesar():
    """Test Caesar cipher."""
    print("Testing Caesar Cipher...")
    message = "HELLO WORLD"
    key = 3

    encrypted = cp.caesar_encrypt(message, key)
    decrypted = cp.caesar_decrypt(encrypted, key)

    print(f"Original: {message}")
    print(f"Encrypted: {encrypted}")
    print(f"Decrypted: {decrypted}")

    assert decrypted == message, "Caesar cipher test failed"
    print("✓ Caesar cipher test passed\n")


def test_vigenere():
    """Test Vigenère cipher."""
    print("Testing Vigenère Cipher...")
    message = "HELLO WORLD"
    key = "KEY"

    encrypted = cp.vigenere_encrypt(message, key)
    decrypted = cp.vigenere_decrypt(encrypted, key)

    print(f"Original: {message}")
    print(f"Key: {key}")
    print(f"Encrypted: {encrypted}")
    print(f"Decrypted: {decrypted}")

    assert decrypted == message, "Vigenère cipher test failed"
    print("✓ Vigenère cipher test passed\n")


def test_affine():
    """Test Affine cipher."""
    print("Testing Affine Cipher...")
    message = "HELLO"
    a, b = 5, 8

    encrypted = cp.affine_encrypt(message, a, b)
    decrypted = cp.affine_decrypt(encrypted, a, b)

    print(f"Original: {message}")
    print(f"Keys: a={a}, b={b}")
    print(f"Encrypted: {encrypted}")
    print(f"Decrypted: {decrypted}")

    assert decrypted == message, "Affine cipher test failed"
    print("✓ Affine cipher test passed\n")


def test_substitution():
    """Test Substitution cipher."""
    print("Testing Substitution Cipher...")
    message = "HELLO"
    key = {"h": "m", "e": "n", "l": "o", "o": "p"}

    encrypted = cp.substitution_encrypt(message, key)
    decrypted = cp.substitution_decrypt(encrypted, key)

    print(f"Original: {message}")
    print(f"Key: {key}")
    print(f"Encrypted: {encrypted}")
    print(f"Decrypted: {decrypted}")

    assert decrypted == message, "Substitution cipher test failed"
    print("✓ Substitution cipher test passed\n")


def test_hill():
    """Test Hill cipher."""
    print("Testing Hill Cipher...")
    message = "HELLO"
    # Use a matrix that has an inverse modulo 26
    key_matrix = np.array([[3, 2], [5, 7]])

    encrypted = cp.hill_encrypt(message, key_matrix)
    decrypted = cp.hill_decrypt(encrypted, key_matrix)

    print(f"Original: {message}")
    print(f"Key matrix:\n{key_matrix}")
    print(f"Encrypted: {encrypted}")
    print(f"Decrypted: {decrypted}")

    # Remove padding for comparison
    decrypted_clean = decrypted.rstrip("x").upper()
    message_clean = message.upper()

    assert (
        decrypted_clean == message_clean
    ), f"Hill cipher test failed: {decrypted_clean} != {message_clean}"
    print("✓ Hill cipher test passed\n")


def test_otp():
    """Test One-Time Pad cipher."""
    print("Testing OTP Cipher...")
    message = "HELLO"
    key = "MYSECRETKEY"

    encrypted = cp.otp_encrypt(message, key)
    decrypted = cp.otp_decrypt(encrypted, key)

    print(f"Original: {message}")
    print(f"Key: {key}")
    print(f"Encrypted: {repr(encrypted)}")  # Use repr to show special characters
    print(f"Decrypted: {decrypted}")

    assert decrypted == message, "OTP cipher test failed"
    print("✓ OTP cipher test passed\n")


def test_safe_parsing():
    """Test safe parsing functions."""
    print("Testing Safe Parsing Functions...")

    # Test substitution key parsing
    key_str = "{'a': 'b', 'c': 'd'}"
    parsed_key = cp.safe_parse_substitution_key(key_str)
    expected = {"a": "b", "c": "d"}
    assert (
        parsed_key == expected
    ), f"Substitution key parsing failed: {parsed_key} != {expected}"

    # Test matrix parsing
    matrix_str = "[[1, 2], [3, 4]]"
    parsed_matrix = cp.safe_parse_matrix(matrix_str)
    expected_matrix = np.array([[1, 2], [3, 4]])
    assert np.array_equal(parsed_matrix, expected_matrix), "Matrix parsing failed"

    print("✓ Safe parsing tests passed\n")


def main():
    """Run all tests."""
    print("Running Cipher Tests...\n")

    try:
        test_caesar()
        test_vigenere()
        test_affine()
        test_substitution()
        test_hill()
        test_otp()
        test_safe_parsing()

        print("🎉 All tests passed successfully!")

    except Exception as e:
        print(f"❌ Test failed: {e}")
        return False

    return True


if __name__ == "__main__":
    main()