Rot-n/Caesar Cipher

rot_cipher.py

"""rot_cipher.py - A rot-n/Caesar cipher tool that takes user input plaintext and outputs ciphertext
                 - Provides encrypt() and decrypt()"""


# Store alphabet and the length for later use
ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
ALPHABET_LEN = len(ALPHABET)


def encrypt(plaintext: str, rot: int) -> str:
    """Apply rot-n encryption to plaintext and output

    :param plaintext: input plaintext
    :param rot: rotation amount (i.e. the key)
    :return: output ciphertext
    """

    # Go through each character in input text and rotate
    ciphertext: str = ''
    for char in plaintext.upper():
        # Only rotate characters that exist in alphabet
        if char in ALPHABET:
            # Get "theoretical index", as if alphabet wrapped around/was endlessly repeating
            idx = ALPHABET.index(char) + rot

            # Check if index is out of range. If so, subtract until we get in range
            while idx >= ALPHABET_LEN:
                idx -= ALPHABET_LEN

            # Grab rotated char and add to ciphertext
            ciphertext += ALPHABET[idx]

        # Otherwise, just add as-is
        else:
            ciphertext += char

    return ciphertext


def decrypt(ciphertext: str, rot: int) -> str:
    """Apply rot-n decryption to ciphertext and output

    :param ciphertext: input ciphertext
    :param rot: rotation amount (i.e. the key)
    :return: output plaintext
    """

    # Go through each character in input text and reverse-rotate
    plaintext: str = ''
    for char in ciphertext.upper():
        # Only rotate characters that exist in alphabet
        if char in ALPHABET:
            # Get "theoretical index", as if alphabet wrapped around/was endlessly repeating
            idx = ALPHABET.index(char) - rot

            # Check if index is out of range. If so, add until we get in range
            while idx <= 0:
                idx += ALPHABET_LEN

            # Grab rotated char and add to ciphertext
            plaintext += ALPHABET[idx]

        # Otherwise, just add as-is
        else:
            plaintext += char

    return plaintext


def main():
    """Main script that asks user for input and provides output to user"""

    print('Welcome to Caesar/ROT-n Cipher\n')

    # Ask for encrypt or decrypt mode
    mode = input(f'''Select mode: 'e' for encrypt, 'd' for decrypt, 'q' to exit: ''').strip()
    if mode == 'q':
        print('Bye!')
        exit()
    elif mode not in ['e', 'd']:
        print('Invalid mode.')
        exit()

    # Ask for input text (named '_in' as 'in' is a reserved keyword)
    _in = input(f'''Enter text that you want to {'en' if mode == 'e' else 'de'}crypt (non-alphabetical characters will be ignored): ''')

    # Ask for input rotation amount (exit if not a number)
    try:
        rot = int(input('Enter numeric key (rotation amount): '))
    except ValueError:
        print('Key must be a number.')
        exit()

    # Exit if rotation amount is negative
    if rot < 0:
        print('Key must be greater than 0.')
        exit()

    # Detect whether user wants to encrypt or decrypt, run corresponding function, and print result
    if mode == 'e':
        print(encrypt(_in, rot))
    elif mode == 'd':
        print(decrypt(_in, rot))


# If script is run (not imported), continuously loop through main function
if __name__ == '__main__':
    while True:
        main()