'''
Substitution encryption and decryption, after
Miller and Ranum, chapter 3 and Jon Beck
'''

import random


# Note that the following function works in this setting, but
# in general contains an error if index is -1
def removeChar(string, index):
    return string[:index] + string[index+1:]

# randomly generate a permutation of alphabet where each character
# appears exactly once.
def randomKeyGen():
    alphabet = "abcdefghijklmnopqrstuvwxyz"
    key = ""
    for i in range(len(alphabet)):
        idx = random.randint(0,25-i)
        key = key + alphabet[idx]
        alphabet = removeChar(alphabet,idx)
    return key

# This function returns a new string with all duplicates removed.
def removeDuplicateCharacters(stringWithDuplicates):
    stringWithoutDuplicates = ''
    for character in stringWithDuplicates:
        # Add character to the new string if it isn't already there.
        if character not in stringWithoutDuplicates:
            stringWithoutDuplicates += character

    return stringWithoutDuplicates

# This returns a new string with all characters in charactersToRemove
# removed from originalString
def removeMatches(originalString, charactersToRemove):
    newString = ''
    for character in originalString:
        if character not in charactersToRemove:
            newString += character

    return newString

# This generates an encryption key from a password by: removing duplicates
# from the password, and creating a new string that breaks the alphabet
# at the last character in the password.
def generateKeyFromPassword(password):
    alphabet = 'abcdefghijklmnopqrstuvwxyz '
    password = removeDuplicateCharacters(password)
    lastCharacter = password[-1]
    lastIndex = alphabet.find(lastCharacter)
    afterString = removeMatches(alphabet[lastIndex+1:], password)
    beforeString = removeMatches(alphabet[:lastIndex], password)
    return password + afterString + beforeString
    
# This is just substitutionEncrypt that calls the password making
# code from above.
def substitutionEncrypt(plainText, password):
    alphabet = 'abcdefghijklmnopqrstuvwxyz '
    plainText = plainText.lower()
    key = generateKeyFromPassword(password)

    cipherText = ''
    for character in plainText:
        index = alphabet.find(character)
        cipherText += key[index]

    return cipherText

# Decrypt messages encoded as above.
def substitutionDecrypt(cipherText, password):
    alphabet = 'abcdefghijklmnopqrstuvwxyz '
    key = generateKeyFromPassword(password)

    plaintext = ''
    for ch in cipherText:
        index = key.find(ch)
        plaintext += alphabet[index]

    return plaintext

# a general function to encrypt a message using key
def substitutionEncryptWithKey(plainText, key):

    # This is the alphabet of characters that are allowed in the plainText
    alphabet = "abcdefghijklmnopqrstuvwxyz "
    # Convert the plainText message to all lower case.
    plainText = plainText.lower()

    # initialize the cipherText variable
    cipherText = ""
    for ch in plainText:
        # find the plainText character in the alphabet
        idx = alphabet.find(ch)

        # replace it with the character in the equivalent position
        # in the key
        # Remember that + here is doing string concatenation.
        cipherText = cipherText + key[idx]
        
    return cipherText