Ciphers.py

"""
Ciphers Module

This module contains implementations of various cryptographic ciphers.

Functions:
- Caesar, Vigenère, Affine, Hill, Substitution, and OTP cipher implementations
- Helper functions for mathematical operations

Dependencies:
- numpy: Used for numerical operations and array manipulations, particularly for Hill cipher
"""

import numpy as np
import json
import re
from typing import Dict, List, Union, Any

import numpy as np
import json
import re
from typing import Dict, List, Union, Any


def safe_parse_substitution_key(key_str: str) -> Dict[str, str]:
    """
    Safely parse substitution cipher key from string format.

    Args:
        key_str (str): String representation of substitution key

    Returns:
        Dict[str, str]: Parsed substitution key dictionary

    Raises:
        ValueError: If key format is invalid
    """
    try:
        # Remove whitespace and validate basic structure
        key_str = key_str.strip()
        if not (key_str.startswith("{") and key_str.endswith("}")):
            raise ValueError(
                "Key must be in dictionary format: {'a': 'b', 'c': 'd', ...}"
            )

        # Use json.loads for safer parsing
        # First convert Python dict format to JSON format
        json_str = key_str.replace("'", '"')
        key_dict = json.loads(json_str)

        # Validate that all keys and values are single characters
        for k, v in key_dict.items():
            if not (
                isinstance(k, str)
                and isinstance(v, str)
                and len(k) == 1
                and len(v) == 1
            ):
                raise ValueError("All keys and values must be single characters")
            if not k.isalpha() or not v.isalpha():
                raise ValueError("All keys and values must be alphabetic characters")

        return {k.lower(): v.lower() for k, v in key_dict.items()}
    except json.JSONDecodeError:
        raise ValueError("Invalid dictionary format. Use: {'a': 'b', 'c': 'd', ...}")
    except Exception as e:
        raise ValueError(f"Invalid substitution key format: {str(e)}")


def safe_parse_matrix(matrix_str: str) -> np.ndarray:
    """
    Safely parse matrix from string format.

    Args:
        matrix_str (str): String representation of matrix

    Returns:
        np.ndarray: Parsed matrix

    Raises:
        ValueError: If matrix format is invalid
    """
    try:
        # Remove whitespace and validate basic structure
        matrix_str = matrix_str.strip()
        if not (matrix_str.startswith("[") and matrix_str.endswith("]")):
            raise ValueError("Matrix must be in list format: [[1, 2], [3, 4]]")

        # Use json.loads for safer parsing
        matrix_list = json.loads(matrix_str)

        if not isinstance(matrix_list, list) or not all(
            isinstance(row, list) for row in matrix_list
        ):
            raise ValueError("Matrix must be a list of lists")

        # Validate that all elements are numbers
        for row in matrix_list:
            for element in row:
                if not isinstance(element, (int, float)):
                    raise ValueError("All matrix elements must be numbers")

        matrix = np.array(matrix_list, dtype=int)

        # Validate square matrix
        if matrix.shape[0] != matrix.shape[1]:
            raise ValueError("Matrix must be square")

        if matrix.shape[0] < 2:
            raise ValueError("Matrix must be at least 2x2")

        return matrix
    except json.JSONDecodeError:
        raise ValueError("Invalid matrix format. Use: [[1, 2], [3, 4]]")
    except Exception as e:
        raise ValueError(f"Invalid matrix format: {str(e)}")


def validate_affine_keys(a: int, b: int) -> None:
    """
    Validate affine cipher keys.

    Args:
        a (int): Multiplicative key
        b (int): Additive key

    Raises:
        ValueError: If keys are invalid
    """
    if not isinstance(a, int) or not isinstance(b, int):
        raise ValueError("Both keys must be integers")
    if gcd(a, 26) != 1:
        raise ValueError("Multiplicative key 'a' must be coprime with 26")
    if not (0 <= b < 26):
        raise ValueError("Additive key 'b' must be between 0 and 25")


"""
Caesar's Cipher: A substitution cipher that shifts the letters of the plaintext
by a fixed number of positions in the alphabet.
For example, with a shift of 3, 'A' becomes 'D', 'B' becomes 'E', and so on.
It is one of the simplest and most widely known encryption techniques.
"""


def caesar_encrypt(plaintext: str, key: int) -> str:
    """
    Encrypts the plaintext using Caesar's cipher.

    Args:
      plaintext (str): The text to be encrypted.
      key (int): The number of positions to shift each letter.

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If key is not an integer or plaintext is not a string.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")
    if not isinstance(key, int):
        raise ValueError("Key must be an integer")

    encrypted = []
    for char in plaintext:
        if char.isalpha():
            shift_base = ord("A") if char.isupper() else ord("a")
            encrypted.append(chr((ord(char) - shift_base + key) % 26 + shift_base))
        else:
            encrypted.append(char)
    return "".join(encrypted)


def caesar_decrypt(ciphertext: str, key: int) -> str:
    """
    Decrypts the ciphertext using Caesar's cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      key (int): The number of positions to shift each letter back.

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If key is not an integer or ciphertext is not a string.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")
    if not isinstance(key, int):
        raise ValueError("Key must be an integer")

    return caesar_encrypt(ciphertext, -key)


"""
Vigenère Cipher: A method of encrypting alphabetic text by using a simple form of polyalphabetic substitution.
It uses a keyword where each letter of the keyword determines the shift for the corresponding letter in the plaintext.
For example, if the keyword is "KEY" and the plaintext is "HELLO", the first letter 'H' is shifted by the position of 'K',
the second letter 'E' by the position of 'E', and so on. The keyword is repeated as necessary to match the length of the plaintext.
"""


def vigenere_encrypt(plaintext: str, keyword: str) -> str:
    """
    Encrypts the plaintext using the Vigenère cipher.

    Args:
      plaintext (str): The text to be encrypted.
      keyword (str): The keyword used for encryption.

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")
    if not isinstance(keyword, str):
        raise ValueError("Keyword must be a string")
    if not keyword.strip():
        raise ValueError("Keyword cannot be empty")
    if not keyword.isalpha():
        raise ValueError("Keyword must contain only alphabetic characters")

    encrypted = []
    keyword = keyword.lower()
    keyword_length = len(keyword)
    keyword_index = 0

    for char in plaintext:
        if char.isalpha():
            shift_base = ord("A") if char.isupper() else ord("a")
            shift = ord(keyword[keyword_index % keyword_length]) - ord("a")
            encrypted.append(chr((ord(char) - shift_base + shift) % 26 + shift_base))
            keyword_index += 1
        else:
            encrypted.append(char)
    return "".join(encrypted)


def vigenere_decrypt(ciphertext: str, keyword: str) -> str:
    """
    Decrypts the ciphertext using the Vigenère cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      keyword (str): The keyword used for decryption.

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")
    if not isinstance(keyword, str):
        raise ValueError("Keyword must be a string")
    if not keyword.strip():
        raise ValueError("Keyword cannot be empty")
    if not keyword.isalpha():
        raise ValueError("Keyword must contain only alphabetic characters")

    decrypted = []
    keyword = keyword.lower()
    keyword_length = len(keyword)
    keyword_index = 0

    for char in ciphertext:
        if char.isalpha():
            shift_base = ord("A") if char.isupper() else ord("a")
            shift = ord(keyword[keyword_index % keyword_length]) - ord("a")
            decrypted.append(chr((ord(char) - shift_base - shift) % 26 + shift_base))
            keyword_index += 1
        else:
            decrypted.append(char)
    return "".join(decrypted)


def gcd(a, b):  # For Affine Cipher
    """
    Computes the greatest common divisor of two numbers using the Euclidean algorithm.

    Args:
      a (int): The first number.
      b (int): The second number.

    Returns:
      int: The greatest common divisor of a and b.
    """
    while b:
        a, b = b, a % b
    return a


"""
Affine Cipher: A type of monoalphabetic substitution cipher, where each letter in an alphabet is mapped to its numeric equivalent,
encrypted using a simple mathematical function, and then converted back to a letter. The encryption function is:
E(x) = (a * x + b) % 26
where 'a' and 'b' are keys, and 'x' is the numeric equivalent of the plaintext letter. The decryption function is:
D(x) = a_inv * (x - b) % 26
where 'a_inv' is the modular multiplicative inverse of 'a' modulo 26.

The key 'a' must be chosen such that gcd(a, 26) = 1 to ensure that 'a' has an inverse modulo 26.
"""


def affine_encrypt(plaintext: str, a: int, b: int) -> str:
    """
    Encrypts the plaintext using the Affine cipher.

    Args:
      plaintext (str): The text to be encrypted.
      a (int): The multiplicative key (must satisfy gcd(a, 26) = 1).
      b (int): The additive key.

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")

    validate_affine_keys(a, b)

    encrypted = []
    for char in plaintext:
        if char.isalpha():
            shift_base = ord("A") if char.isupper() else ord("a")
            x = ord(char) - shift_base
            encrypted.append(chr((a * x + b) % 26 + shift_base))
        else:
            encrypted.append(char)
    return "".join(encrypted)


def affine_decrypt(ciphertext: str, a: int, b: int) -> str:
    """
    Decrypts the ciphertext using the Affine cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      a (int): The multiplicative key (must satisfy gcd(a, 26) = 1).
      b (int): The additive key.

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")

    validate_affine_keys(a, b)

    a_inv = pow(a, -1, 26)  # Modular multiplicative inverse of 'a' modulo 26
    decrypted = []
    for char in ciphertext:
        if char.isalpha():
            shift_base = ord("A") if char.isupper() else ord("a")
            y = ord(char) - shift_base
            decrypted.append(chr((a_inv * (y - b)) % 26 + shift_base))
        else:
            decrypted.append(char)
    return "".join(decrypted)


def matrix_mod_inv(matrix, modulus):  # For Hill Cipher
    """
    Computes the modular inverse of a square matrix under a given modulus.

    Args:
      matrix (numpy.ndarray): A square matrix represented as a numpy array.
      modulus (int): The modulus for the modular arithmetic.

    Returns:
      numpy.ndarray: The modular inverse of the matrix.

    Raises:
      ValueError: If the matrix is not invertible under the given modulus.
    """
    if matrix.shape[0] != matrix.shape[1]:
        raise ValueError("Only square matrices are supported.")

    # Calculate the determinant
    det = int(round(np.linalg.det(matrix))) % modulus

    # Check if the determinant has a modular inverse
    try:
        det_inv = pow(det, -1, modulus)
    except ValueError:
        raise ValueError("Matrix is not invertible under the given modulus.")

    # Compute the adjugate matrix
    adjugate = np.round(np.linalg.inv(matrix) * det).astype(int) % modulus

    # Apply the modular inverse of the determinant and modulus to the adjugate
    inverse = (det_inv * adjugate) % modulus

    return inverse


"""
Hill Cipher: A polygraphic substitution cipher based on linear algebra.
It uses a key matrix to encrypt and decrypt the text by treating the plaintext as a series of vectors.
The encryption function is:
E(P) = K * P mod 26
where 'E(P)' is the encrypted vector, 'K' is the key matrix, and 'P' is the plaintext vector.
The decryption function is:
D(C) = K_inv * C mod 26
where 'D(C)' is the decrypted vector, 'K_inv' is the modular inverse of the key matrix, and 'C' is the ciphertext vector.
"""


def hill_encrypt(plaintext: str, key_matrix: np.ndarray) -> str:
    """
    Encrypts the plaintext using the Hill cipher.

    Args:
      plaintext (str): The text to be encrypted.
      key_matrix (numpy.ndarray): The encryption key matrix.

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")
    if not isinstance(key_matrix, np.ndarray):
        raise ValueError("Key matrix must be a numpy array")

    n = key_matrix.shape[0]
    if key_matrix.shape[0] != key_matrix.shape[1]:
        raise ValueError("Key matrix must be square.")

    if not plaintext.strip():
        raise ValueError("Plaintext cannot be empty")

    plaintext = plaintext.lower().replace(" ", "")
    # Only keep alphabetic characters
    plaintext = "".join(char for char in plaintext if char.isalpha())

    if not plaintext:
        raise ValueError("Plaintext must contain at least one alphabetic character")

    while len(plaintext) % n != 0:
        plaintext += "x"  # Padding with 'x' to fit the matrix size

    plaintext_vectors = [
        [ord(char) - ord("a") for char in plaintext[i : i + n]]
        for i in range(0, len(plaintext), n)
    ]

    encrypted = []
    for vector in plaintext_vectors:
        encrypted_vector = np.dot(key_matrix, vector) % 26
        encrypted.extend(chr(num + ord("a")) for num in encrypted_vector)

    return "".join(encrypted)


def hill_decrypt(ciphertext: str, key_matrix: np.ndarray) -> str:
    """
    Decrypts the ciphertext using the Hill cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      key_matrix (numpy.ndarray): The encryption key matrix.

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")
    if not isinstance(key_matrix, np.ndarray):
        raise ValueError("Key matrix must be a numpy array")

    n = key_matrix.shape[0]
    if key_matrix.shape[0] != key_matrix.shape[1]:
        raise ValueError("Key matrix must be square.")

    if not ciphertext.strip():
        raise ValueError("Ciphertext cannot be empty")

    ciphertext = ciphertext.lower().replace(" ", "")
    # Only keep alphabetic characters
    ciphertext = "".join(char for char in ciphertext if char.isalpha())

    if not ciphertext:
        raise ValueError("Ciphertext must contain at least one alphabetic character")

    if len(ciphertext) % n != 0:
        raise ValueError(f"Ciphertext length must be divisible by {n}")

    ciphertext_vectors = [
        [ord(char) - ord("a") for char in ciphertext[i : i + n]]
        for i in range(0, len(ciphertext), n)
    ]

    try:
        key_matrix_inv = matrix_mod_inv(key_matrix, 26)
    except ValueError as e:
        raise ValueError(f"Cannot decrypt: {str(e)}")

    decrypted = []
    for vector in ciphertext_vectors:
        decrypted_vector = np.dot(key_matrix_inv, vector) % 26
        decrypted.extend(
            chr((int(round(num)) % 26) + ord("a")) for num in decrypted_vector
        )

    return "".join(decrypted)


"""
Substitution Cipher: A method of encryption where each letter in the plaintext is replaced with another letter or symbol.
The key for this cipher is a mapping of each letter in the alphabet to a unique substitution.
"""


def substitution_encrypt(plaintext: str, key: Dict[str, str]) -> str:
    """
    Encrypts the plaintext using a substitution cipher.

    Args:
      plaintext (str): The text to be encrypted.
      key (dict): A dictionary mapping each letter to its substitution.

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")
    if not isinstance(key, dict):
        raise ValueError("Key must be a dictionary")
    if not key:
        raise ValueError("Key dictionary cannot be empty")

    encrypted = []
    for char in plaintext:
        if char.isalpha():
            char_lower = char.lower()
            if char_lower in key:
                if char.isupper():
                    encrypted.append(key[char_lower].upper())
                else:
                    encrypted.append(key[char_lower])
            else:
                encrypted.append(char)  # Keep unmapped characters as-is
        else:
            encrypted.append(char)
    return "".join(encrypted)


def substitution_decrypt(ciphertext: str, key: Dict[str, str]) -> str:
    """
    Decrypts the ciphertext using a substitution cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      key (dict): A dictionary mapping each letter to its substitution.

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")
    if not isinstance(key, dict):
        raise ValueError("Key must be a dictionary")
    if not key:
        raise ValueError("Key dictionary cannot be empty")

    reverse_key = {v.lower(): k.lower() for k, v in key.items()}
    decrypted = []
    for char in ciphertext:
        if char.isalpha():
            char_lower = char.lower()
            if char_lower in reverse_key:
                if char.isupper():
                    decrypted.append(reverse_key[char_lower].upper())
                else:
                    decrypted.append(reverse_key[char_lower])
            else:
                decrypted.append(char)  # Keep unmapped characters as-is
        else:
            decrypted.append(char)
    return "".join(decrypted)


"""
OTP Cipher: A symmetric encryption algorithm that generates a random key for each message.
The key must be at least as long as the message and should only be used once.
The encryption function is:
C = P XOR K
where 'C' is the ciphertext, 'P' is the plaintext, and 'K' is the one-time pad key.
The decryption function is:
P = C XOR K
where 'P' is the plaintext, 'C' is the ciphertext, and 'K' is the one-time pad key.
"""


def otp_encrypt(plaintext: str, key: str) -> str:
    """
    Encrypts the plaintext using the One-Time Pad (OTP) cipher.

    Args:
      plaintext (str): The text to be encrypted.
      key (str): The one-time pad key (must be at least as long as the plaintext).

    Returns:
      str: The encrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(plaintext, str):
        raise ValueError("Plaintext must be a string")
    if not isinstance(key, str):
        raise ValueError("Key must be a string")
    if len(key) < len(plaintext):
        raise ValueError("Key must be at least as long as the plaintext.")
    if not plaintext.strip():
        raise ValueError("Plaintext cannot be empty")

    encrypted = []
    for p, k in zip(plaintext, key):
        encrypted.append(chr(ord(p) ^ ord(k)))
    return "".join(encrypted)


def otp_decrypt(ciphertext: str, key: str) -> str:
    """
    Decrypts the ciphertext using the One-Time Pad (OTP) cipher.

    Args:
      ciphertext (str): The text to be decrypted.
      key (str): The one-time pad key (must be at least as long as the ciphertext).

    Returns:
      str: The decrypted text.

    Raises:
      ValueError: If inputs are invalid.
    """
    if not isinstance(ciphertext, str):
        raise ValueError("Ciphertext must be a string")
    if not isinstance(key, str):
        raise ValueError("Key must be a string")
    if len(key) < len(ciphertext):
        raise ValueError("Key must be at least as long as the ciphertext.")
    if not ciphertext.strip():
        raise ValueError("Ciphertext cannot be empty")

    decrypted = []
    for c, k in zip(ciphertext, key):
        decrypted.append(chr(ord(c) ^ ord(k)))
    return "".join(decrypted)