Caesar cipher

caesar_cipher.py

#!/usr/bin/env python3

# Written by Lari Liuhamo

import string
from typing import Generator

def caesar_cipher(message: str, shift: int, alphabet=string.ascii_uppercase) -> str:
    """
    Performs the Caesar cipher on a given message (string) by shifting all characters using the given shift.
    Defaults to the English alphabet, an alternative alphabet can be given as an optional third argument.
    
    Returns the shifted message.
    
    UPDATE 2021/02/03 - Character case information is now preserved
    """
    
    result = ""

    for char in message:
        upper_char = char.upper()
        if upper_char in alphabet:
            position = alphabet.find(upper_char)
            new_position = (position + shift) % len(alphabet)
            new_char = alphabet[new_position]
            result += new_char if char.isupper() else new_char.lower()
        else:
            result += char

    return result


def brute_force_gen(message: str, alphabet=string.ascii_uppercase) -> Generator[str, None, None]:
    """
    Yields all possible cipher results for a given string using the given alphabet, defaulting to the English alphabet.
    The output text format is the (positive) shifted number and the resulting message.
    
    UPDATE 2021/02/03 - Added dynamic shift number formatting support for larger alphabets
    """
    
    format_digits = len(str(len(alphabet)))
    
    for idx, _ in enumerate(alphabet):
        yield f"{idx:{format_digits}}: {caesar_cipher(message, idx, alphabet=alphabet)}"
        
def brute_force(message: str, alphabet=string.ascii_uppercase):
    """
    Prints all possible cipher results for a given string using the given alphabet, defaulting to the English alphabet.
    The print format is the (positive) shifted number and the resulting message.
    """
    
    for result in brute_force_gen(message, alphabet=alphabet):
        print(result)
        
        
def caesar_cipher_finnish(message: str, shift: int) -> str:
    """
    A shortcut for messages that use the full Finnish alphabet
    """
    
    alphabet = string.ascii_uppercase + "ÅÄÖ"
    return caesar_cipher(message, shift, alphabet=alphabet)

# Examples:
# encrypted = caesar_cipher("Hello, world!", 9)
# decrypted = caesar_cipher(encrypted, -9)
# 
# brute_force(encrypted)
# 
# enc2 = caesar_cipher("VDPPRQUÄBVWB", -3, "ABCDEFGHIJKLMNOPQRSTUVWXYZÅÄÖ")