Various crypto algorithims simple implementations

crypto.py

__author__ = 'wasi0013'
#Algorithm resources: http://crypto.interactive-maths.com/introduction-to-cryptography.html

def caeser_encrypt(text, key):
    """
    Takes string as text and encrypts it with the key
    >>> caeser_encrypt("Hello World",2)
    'JGNNQ YQTNF'
    >>> caeser_encrypt("Hello World",0)
    'HELLO WORLD'
    """
    return "".join(i if i == " " else chr((ord(i.lower())-97 + key)%26 + 65) for i in text if i.isalpha() or i == " ")

def caeser_decrypt(text, key):
    """
    Takes string as text and decrypts it with the key
    >>> caeser_decrypt("JGNNQ YQTNF",2)
    'HELLO WORLD'
    """
    return "".join(i if i == " " else chr((ord(i.lower())-97 - key)%26 + 65) for i in text if i.isalpha() or i == " ")

def affine_encrypt(text, a, b, m=26):
    """
    Takes string as text and decrypts it with the key a,b
    >>> affine_encrypt("AFFINE CIPHER",5,8)
    'IHHWVC SWFRCP'
    """
    from fractions import gcd
    if gcd(a, m) == 1:
        return "".join(i if i == " " else chr((a*(ord(i.lower())-97) + b)%m + 65) for i in text if i.isalpha() or i == " ")

def affine_decrypt(text, a, b, m=26):
    """
    Takes string as text and decrypts it with the key a,b
    >>> affine_decrypt("IHHWVC SWFRCP",5,8)
    'AFFINE CIPHER'
    """
    from fractions import gcd
    coprimes = [i for i in range(1, m) if gcd(m, i) == 1]
    c = None
    for num in coprimes:
        if (num * a)%m == 1:
            c = num
    return "".join(i if i == " " else chr((c*(ord(i.lower())-97-b))%m + 65) for i in text if i.isalpha() or i == " ")

def railfence_encrypt(text, key):
    """
    Takes string as text and encrypts it with the key
    >>> railfence_encrypt("defend the east wall",3)
    'DNETLEEDHESWLFTAA'
    """
    text="".join(text.split())
    m = (key - 1)*2
    msg = ''
    for i in range(key):
        if i%(key - 1) == 0: msg += text[i::m]
        else:
            char_pairs = zip(text[i::m], list(text[m-i::m]) + [''])
            msg += ''.join(map(''.join, char_pairs))
    return msg.upper()
    
def transposition_encrypt(text, key):
    """
    Takes string as text and encrypts it with a keyword as key
    >>> transposition_encrypt("The tomato is a plant in the nightshade family","TOMATO")
    'TINESAXEOAHTFXHTLTHEYMAIIAIXTAPNGDLOSTNHMX'
    """
    #get column indices for the given key
    col_indices = [None]*len(key)
    temp = [i for i in key.upper()]
    temp.sort()
    for i, v in enumerate(temp):
        if col_indices[key.find(v)] is None:
            col_indices[key.find(v)] = i
        else:
            col_indices[key.find(v,key.find(v)+1)] = i

    text = "".join(text.split()).upper() + "X"*(len(text)%len(key))
    #matrix where each string is a column
    msg = [""]*len(key)
    for i, letter in enumerate(text):
        msg[i%len(key)] += letter

    msg = dict(zip(col_indices,msg))
    return "".join(msg[i] for i in sorted(msg))

def transposition_decrypt(text, key):
    """
    Takes string as text and decrypts it with a keyword as key
    >>> transposition_decrypt('TINESAXEOAHTFXHTLTHEYMAIIAIXTAPNGDLOSTNHMX',"TOMATO")
    'THETOMATOISAPLANTINTHENIGHTSHADEFAMILYXXXX'
    """
    #get column indices for the given key
    col_indices = [None]*len(key)
    temp = [i for i in key.upper()]
    temp.sort()
    for i, v in enumerate(temp):
        if col_indices[key.find(v)] is None:
            col_indices[key.find(v)] = i
        else:
            col_indices[key.find(v,key.find(v)+1)] = i
            
    #split the text into matrix of strings
    col = len(text)//len(key)
    msg = [text[i:i+col] for i in range(0, len(text), col)]
    
    #rearrange it to original matrix using key
    msg = [msg[i] for i in col_indices]
    
    plain_text = ""
    for i in range(len(key) + 1):
        for col in msg:
            plain_text += col[i]

    return plain_text

Test using python3 -m doctest crypto.py -v