play_fair_cipher.py

import string
from pprint import pprint


def get_pairs_of_text(text):
    '''
    This function generates pairs of elements
    from the given text.
    eg: BALLOON becomes BA LX LO ON
    '''
    pairs = []
    i = 0

    while i < len(text):
        a = text[i]
        try:
            b = text[i+1]
        except Exception as e:
            pairs.append((a, 'X'))
            break

        if a == b:
            pairs.append((a, 'X'))
            i += 1
        else:
            pairs.append((a, b))
            i += 2
    return pairs


def get_index(k, key):
    '''
    This function returns the index of the specified element
    inside the key matrix.
    '''
    for i in range(5):
        for j in range(5):
            if key[i][j] == k:
                return i, j


def find_new_pairs(pair, key, encrypt):
    '''
    This is the function which returns a tuple consisting of new characters 
    based on the pair given as input.
    '''
    a, b = pair
    ai, aj = get_index(a, key)
    bi, bj = get_index(b, key)

    if ai == bi:
        if encrypt: # move right
            aj = (aj + 1) % 5
            bj = (bj + 1) % 5
        else: # move left
            aj = (aj - 1) % 5
            bj = (bj - 1) % 5
    elif aj == bj:
        if encrypt: # move down
            ai = (ai + 1) % 5
            bi = (bi + 1) % 5
        else: # move up
            ai = (ai - 1) % 5
            bi = (bi - 1) % 5
    else:
        aj, bj = bj, aj

    return (key[ai][aj], key[bi][bj])


def encrypt(text, key):
    '''
    This is the main encrypt function
    '''

    # here we generate pairs of the original text
    pairs = get_pairs_of_text(text)
    new_pairs = []
    for pair in pairs:
        new_pairs.append(find_new_pairs(pair, key, encrypt=True))
    return ''.join(a+b for a, b in new_pairs)


def decrypt(text, key):
    '''
    This is the main decrypt function.
    '''

    # here we generate pairs of the cipher text
    pairs = get_pairs_of_text(text)
    new_pairs = []
    for pair in pairs:
        new_pairs.append(find_new_pairs(pair, key, encrypt=False))
    return ''.join(a+b for a, b in new_pairs).replace('X', '')


def create_key(key):
    '''
    This function creates a 5x5 matrix for the given key
    '''
    all_elements = []

    # this list contains those elements which are not present in the key
    list_of_non_keys = list(set(list(string.ascii_uppercase)) - set(list(key)))

    # generating a list of 25 characters by removing the one which is not present in the key
    elements = list(key + string.ascii_uppercase.replace(list_of_non_keys[0], ''))
    for el in elements:
        if not el in all_elements:
            all_elements.append(el)

    KEY = []
    # adding the elements to the matrix
    for i in range(0, 25, 5):
        KEY.append(all_elements[i:i+5])

    print("Key Matrix: ")
    pprint(KEY)

    return KEY


def playfair(text, key):
    print('Original Text: ', text)

    # the encrypted function is passed the text and the key as parameters.
    enc = encrypt(text, key=key)
    print('Cipher Text: ', enc)

    # the decrypt function is passed the cipher text and the key as parameters.
    dec = decrypt(enc, key=key)
    print('Decrypted Text: ', dec)


if __name__ == '__main__':
    text = 'theworkingbros'.replace(' ', '').upper()
    key = 'risinglight'.replace(' ', '').upper()

    # main calling function
    playfair(text, key=create_key(key))

# This code is contributed by Bhushan Borole