""" monoalphabetic.py
    c odenthal
    2005/05/30

    This module contains some classical monoalphabetic encoding
    algorithms. All the algorithms assume an alphabet of 26
    letters. The common plan is to construct a translation
    table that will convert plain text to cipher text. The
    core function will take the translation table and wrap
    it in some housekeeping functions that will, for example,
    strip all the none letter characters from the plain text
    message. A function is returned that will do the enciphering.
"""

import string
import stringtools

UPPERCASE = string.ascii_uppercase
LOWERCASE = string.ascii_lowercase
ALPHINDEX = range(26)

char2int = dict( zip(UPPERCASE, ALPHINDEX) )
int2char = dict( zip(ALPHINDEX, UPPERCASE) )

def monoalphabetic(frm,to,strip=False):
    """ Takes 'frm' and 'to' - two strings of equal length -
        and returns encipher and decipher tables for use
        with the 'string' method 'translate'. The encipher
        table replaces the characters in 'frm' with the
        corresponding characters in 'to'. The decipher
        table reverses the translation.
    """
    frm = frm.lower()
    to = to.upper()
    encTable = string.maketrans(frm,to)
    decTable = string.maketrans(to,frm)

    def encipher(strng):
        strng = strng.lower()
        if strip:
            strng = stringtools.stripPunc(strng)
        return strng.translate(encTable)

    def decipher(strng):
        strng = strng.upper()
        return strng.translate(decTable)

    return encipher,decipher

def standard(shift=3,reverse=False,strip=False):
    """ A standard cyclic shift of 'shift' positions
        modulo 26. The default produces the Caesar
        cipher. An encoder and decoder are returned.
    """
    frm = LOWERCASE
    to  = UPPERCASE[shift:] + UPPERCASE[:shift]
    if reverse:
        to = to[::-1]
    return monoalphabetic(frm,to,strip=strip)

def mixed(pKey="", cKey="", shift=0):
    """ A mixed substitution constructed with
        keyword 'pKey' in the plain text alphabet
        and keyword 'cKey' in the cipher text
        alphabet. The cipher text alphabet is
        shifted by 'shift' positions to the right.
        An encoder and decoder are returned.
    """
    # Clean up 'pKey' and get 'frm' alphabet.--------------

    # 'pKey' should be lower case.
    pKey = pKey.lower()

    # Remove punctuation from 'pKey'.
    pKey = stringtools.stripPunc(pKey)

    # Add the entire alphabet and then remove
    # duplicates to get the 'frm' alphabet.
    frm = stringtools.stripDups(pKey + LOWERCASE)

    # Clean up 'cKey' and get 'to' alphabet.---------------
    # 'cKey' should be upper case.
    cKey = cKey.upper()

    # Remove punctuation from 'cKey'.
    cKey = stringtools.stripPunc(cKey)

    # Add the entire alphabet and then remove
    # duplicates to get the 'to' alphabet.
    to = stringtools.stripDups(cKey + UPPERCASE)

    # Shift the 'to' alphabet to the left by 'shift' positions.
    to  = to[shift:] + to[:shift]

    # The 'monoalphabetic' function will now produce
    # the encryption and decryption functions.
    return monoalphabetic(frm,to)

def affine(a=1,b=0):
    """ An affine transformation  P -> aP+b mod 26.
        'a' must be relatively prime to 26.
        An 'encoder' and a 'decoder' are returned.
    """
    frm = LOWERCASE
    enc = lambda p: int2char[(a*char2int[p] + b) % 26]
    to = "".join(map(enc,UPPERCASE))

    return monoalphabetic(frm,to)