route_cipher.py

"""Module for route cipher encryption with column reversal."""


def encrypt(plaintext: str, key: list) -> str:
    """Encrypt text by transposing columns based on a numerical key."""
    text_list = plaintext.split()
    text_matrix = []
    encryption_matrix = []
    len_k = len(key)
    text_matrix = [text_list[i : i + len_k] for i in range(0, len(text_list), len_k)]
    if len(text_matrix[-1]) < len_k:
        text_matrix[-1] += ["\0"] * (len_k - len(text_matrix[-1]))

    for k in key:
        cipher_row = [row[abs(k) - 1] for row in text_matrix]
        if k < 0:
            cipher_row = cipher_row[::-1]
        encryption_matrix.append(cipher_row)
    cipher_text = " ".join([" ".join(r) for r in encryption_matrix])
    return cipher_text


def decrypt(ciphertext: str, key: list) -> str:
    """Decrypt text by transposing columns based on a numerical key."""
    cipher_matrix = []
    cipher_list = ciphertext.split()
    cols_count = (len(cipher_list) + len(key) - 1) // len(key)
    cipher_matrix = [
        cipher_list[i : i + cols_count] for i in range(0, len(cipher_list), cols_count)
    ]
    for k in key:
        if k < 0:
            cipher_matrix[abs(k) - 1] = cipher_matrix[abs(k) - 1][::-1]
    if len(cipher_matrix[-1]) < len(cipher_matrix[0]):
        cipher_matrix[-1] += [""] * (len(cipher_matrix[0]) - len(cipher_matrix[-1]))
    decryption_matrix = [[0 for _ in range(len(key))] for _ in range(cols_count)]

    for k in key:
        for i in range(cols_count):
            decryption_matrix[i][abs(k) - 1] = cipher_matrix[abs(k) - 1][i]
    plain_text = " ".join([" ".join(r) for r in decryption_matrix])
    return plain_text


if __name__ == "__main__":
    CIPHER_KEY = [-1, 2, -3, 4]
    INPUT_TEXT = "0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 26"

    cipher_text = encrypt(INPUT_TEXT, CIPHER_KEY)
    print(f"Encrypted text: {cipher_text}")

    plain_text = decrypt(cipher_text, CIPHER_KEY)
    print(f"Decrypted text: {plain_text}")