Python 3 Vigenere Cipher

Question

I'm still fairly new to Python and I'm trying to see if I've efficiently used modules, functions, etc or if there's another/easier way to do something.

This Python 3 Vigenere Cipher is a rebuild of a JavaScript-based cipher and based in Windows. It accepts any alphabet character, and has demo options built in. Cipher demo uses stage 1 of the CIA Kryptos cipher.

#! python
import os
import re

## Initialize global variables
continue_cipher = ""
demo_alphabet = "KRYPTOSABCDEFGHIJLMNQUVWXZ"
demo_key = "PALIMPSEST"
demo_cipher_string = "EMUFPHZLRFAXYUSDJKZLDKRNSHGNFIVJYQTQUXQBQVYUVLLTREVJYQTMKYRDMFD"
demo_cipher_decoded = "BETWEENSUBTLESHADINGANDTHEABSENCEOFLIGHTLIESTHENUANCEOFIQLUSION"


## Visuals
def display_header():
    print("################################################")
    print("#                                              #")
    print("#            --- VIGENERE CIPHER ---           #")
    print("#                                              #")
    print("#   A simple Vigenere cipher decoder/encoder   #")
    print("#                                              #")
    print("################################################", end="\n\n")
    return
def display_results(mode, cipher_vars):
    # Clear screen for final results
    os.system('cls')

    # Display header
    display_header()

    # Decompose cipher_vars
    (alphabet, key, cipher_string, results) = cipher_vars

    print("Mode:", "Decrypt" if mode == "D" else "Encrypt", end="\n\n")
    print("Alphabet:", alphabet)
    print("Key:", key)
    print("Cipher String:", cipher_string, end="\n\n")
    print("Decoded string:" if mode == "D" else "Encoded string:", results, end="\n\n")
    return


## Validations
def string_is_alpha(input_string):
    return True if re.match("^[a-zA-Z_]*$", input_string) else False


## Cipher variables
def get_alphabet():
    global demo_alphabet

    while True:
        alphabet = input("Enter cipher alphabet: ").upper()
        if alphabet == "":
            alphabet = demo_alphabet
            break
        elif string_is_alpha(alphabet) is False:
            print("The alphabet is not valid. Alphabet should not contain spaces, digits or special characters.")
        else:
            break

    return alphabet
def get_key():
    global demo_key

    while True:
        key = input("Enter cipher key: ").upper()
        if key == "":
            key = demo_key
            break
        elif string_is_alpha(key) is False:
            print("The key is not valid. Key should not contain spaces, digits or special characters.")
        else:
            break

    return key
def get_cipher_string(mode):
    global demo_cipher_string
    global demo_cipher_decoded

    while True:
        cipher_string = input("Enter cipher string: ").upper()
        if cipher_string == "":
            cipher_string = demo_cipher_string if mode == "D" else demo_cipher_decoded
            break
        elif string_is_alpha(cipher_string) is False:
            print("The cipher string is not valid. Cipher strings should not contain spaces, digits or special characters.")
        else:
            break

    return cipher_string


## Cipher actions
def get_cipher_alphabets(alphabet, key):
    cipher_alphabets = []

    for char in key:
        char_index = alphabet.find(char)
        cipher_alphabet = alphabet[char_index:] + alphabet[:char_index]
        cipher_alphabets.append(cipher_alphabet)

    return cipher_alphabets
def start_cipher(mode, alphabet, key, cipher_string):
    mode_string = ""
    cipher_alphabets = get_cipher_alphabets(alphabet, key)
    
    cipher_alphabet_index = 0
    for char in cipher_string:
        # Reset cipher_alphabet_index to 0 when at end of cipher alphabets
        if cipher_alphabet_index == len(cipher_alphabets):
            cipher_alphabet_index = 0
        
        # Use appropriate alphabet based on mode
        # Syntax: base_alphabet[mode_alphabet.find(char)]
        if mode == "D":
            mode_string += alphabet[cipher_alphabets[cipher_alphabet_index].find(char)]
        else:
            mode_string += cipher_alphabets[cipher_alphabet_index][alphabet.find(char)]

        cipher_alphabet_index += 1

    return mode_string


## Cipher Mode
def get_cipher_mode():
    while True:
        cipher_mode = input("Choose cipher mode - [D]ecrypt or [E]ncrypt: ").upper()
        if cipher_mode != "D" and cipher_mode != "E":
            print("That is not a valid option. Please enter 'D' for decrypt and 'E' for encrypt.")
        else:
            break

    print("")
    return cipher_mode
def start_cipher_mode(mode):
    print("Press 'enter' to use demo options")
    alphabet = get_alphabet()
    key = get_key()
    cipher_string = get_cipher_string(mode)
    mode_string = start_cipher(mode, alphabet, key, cipher_string)
    return alphabet, key, cipher_string, mode_string


## Loop cipher
def get_continue_cipher():
    while True:
        continue_cipher = input("Do you want to decode/encode more? [Y/N]: ").upper()
        if continue_cipher != "Y" and continue_cipher != "N":
            print("That is not a valid option. Please enter 'Y' to continue and 'N' to quit.")
        else:
            break
    return continue_cipher


## Start vigenere cipher program
while continue_cipher != "N":
    # Clear the screen after each operation
    os.system('cls')

    # Display header
    display_header()

    # Determine cipher mode
    cipher_mode = get_cipher_mode()
    cipher_vars = start_cipher_mode(cipher_mode)

    # Display results
    display_results(cipher_mode, cipher_vars)

    continue_cipher = get_continue_cipher()

Answer 1

shebang

The shebang should be generic. You are currently calling python, which might point to python 2 on certain systems.

A generic, virtual environment friendly python shebang is:

#!/usr/bin/env python3

PEP-8

Few points from PEP-8 guide:

• Surround top-level function and class definitions with two blank lines.
• Use blank lines in functions, sparingly, to indicate logical sections.
• Constants are usually defined on a module level and written in all capital letters with underscores separating words.

PEP-484

Type hinting functions makes it easier to follow through your functions. Check out the PEP-484.

if __name__ block

Put execution logic of your script inside the if __name__ == "__main__" block. A more descriptive explanation can be checked on Stack Overflow.

Redundant logic

In your code, you have 5 different functions, only to read user input. All of them have the same job of:

1. infinite loop
2. ask for user input
3. convert to uppercase
4. validate if input is empty (or in a set of valid values)
5. in case of empty value, return a default
6. return with the value

All this could be handled by a single function:

def ask_user_input(message: str, options: List[str] = None, default: str = None, check_alpha: bool = False) -> str:
    if not any([options, default]):
        raise ValueError("Either a set of `options` for validation or a fallback `default` needed.")
    while True:
        value = input(message).upper()
        if options:
            if value in options:
                break
            else:
                print(f"Invalid value. Select one of {', '.join(options)}")
                continue
        if default is not None:
            if not value:
                value = default
                break
            elif not check_alpha:
                break
            elif not (value.isalpha() and value.isascii()):
                print("The input text should only consist of ascii alphabets.")
                continue
            else:
                break
    return value

Regex/validation

The regex for validating input allows for _, whereas the error message explicitly says no special characters. The check in the rewrite above is done using (updated based on comment below):

value.isalpha() and value.isascii()

which will perform faster than regex (unless the user keeps inputting wrong values \$ 10^ n \$ times, at which the precompiled pattern might perform slightly better).

Performance

A few things which can be changed to make the code more performant:

1. Instead of concatenating (appending) to string mode_string, push to a list and at the end, use "".join(). More details on Stack Overflow.

2. You can make your program support linux (*nix) systems too. The only dependency on windows is your system call to cls. Perhaps (taken from Stack Overflow):

   def clear():
       os.system("cls" if os.name == "nt" else "clear")
   

3. There are 2 functions with very similar names: start_cipher(mode...) and start_cipher_mode(mode). This makes it really difficult to know which one is truly starting the cipher. Perhaps, have 2 separate functions encrypt and decrypt?

4. Using modulo operation, you can remove the following conditional:

   if cipher_alphabet_index == len(cipher_alphabets):
       cipher_alphabet_index = 0
   

   and would look like:

   result.append(alphabet[cipher_alphabets[cipher_alphabet_index % alphabets_length].find(char)]
   

5. Since you only use the alphabet string to actually work with the index values of characters in it, make a dictionary. Lookups in dictionary is \$ O(1) \$ compared to \$ O(n) \$ for the .find(). This would be:

   from itertools import count
   alphabet_map = dict(zip(alphabet, count()))
   

6. From the above 2 points, it is clear that you don't really need characters/alphabets after the user input. Only the index value modulo matters. This might be difficult to understand/implement without a considerable mathematical understanding, so you can skip this for now.

Answer 2

@hjpotter covered most of my comments.

Booleans

Python has a concept of truthy and falsy values, so it is preferable just to treat values as booleans directly, rather than comparing them to True or False.

---

return True if re.match("^[a-zA-Z_]*$", input_string) else False

can be simplified to:

return re.match("^[a-zA-Z_]*$", input_string)

---

    elif string_is_alpha(alphabet) is False:

This can be simplified to:

    elif not string_is_alpha(alphabet):

If general, you rarely want to use "is" for comparison. (The main exception is comparing to None.)

Compiling regular expressions

This is almost certainly an unnecessary performance improvement, but might be useful to know for later:

The call to re.match needs to compile the regexp every time it is called. You can precompile the regexp once, and then call match on the compiled object to speed it up.

Globals

Almost every time I reach for the global keyword, it turns out to be a mistake.

I don't think you need to declare the demo identifiers as global; they should already be available to use (to read from only - if you attempt to write to them, you will define a new variable in the new scope, hiding the originals).