# Vigenere encryption and decryption functions
# By Eliot Ball

# == Helper functions ==
def NumberToLetter(n):
    "Returns the nth letter in the alphabet, with 0 = 26 = A and so on"
    return "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[n % 26]

def LetterToNumber(l):
    "Returns the location in the alphabet of l, with 0 = 26 = A and so on\nReturns -1 if l is not a letter"
    return "ABCDEFGHIJKLMNOPQRSTUVWXYZ".find(l)

# == Encryption and decryption functions ==
def VigEnc(plaintext, key):
    "Encrypts the given plaintext using the vigenere cipher with the given key"
    # Convert the plaintext and key to upper case
    plaintext = plaintext.upper()
    key = key.upper()
    # Prepare the variables for encryption
    ciphertext = ""
    keypos = 0
    # Iterate through each character in the text encrypting it
    for character in plaintext:
        # Check that it is a letter
        if LetterToNumber(character) != -1:
            # Add the changed character onto the ciphertext
            ciphertext += NumberToLetter(LetterToNumber(character) + LetterToNumber(key[keypos]))
            # Advance through the key, taking the reaminder when dividing by the key length so that it
            # goes back to the start of the key when the end is reached (a % b means remainder of a / b)
            keypos = (keypos + 1) % len(key)
        else:
            # Just add the character on as it is
            ciphertext += character
    # Return the ciphertext
    return ciphertext

def VigDec(ciphertext, key):
    "Decrypts the given ciphertext using the vigenere cipher with the given key"
    # Same as VigEnc but shifting in the opposite direction
    ciphertext = ciphertext.upper()
    key = key.upper()
    plaintext = ""
    keypos = 0
    for character in ciphertext:
        if LetterToNumber(character) != -1:
            plaintext += NumberToLetter(LetterToNumber(character) - LetterToNumber(key[keypos]))
            keypos = (keypos + 1) % len(key)
        else:
            plaintext += character
    return plaintext

# == Cracking tools ==
def Frequencies(text):
    "Returns a list of the percentage of the total character taken up by each letter of the alphabet"
    # Prepare the alphabet
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    # Set up the frequency results
    results = [0 for i in xrange(len(alphabet))]
    # Loop through the alphabet counting the frequencies of each letter
    for letter in xrange(len(alphabet)):
        # Loop through the text counting
        for char in text:
            # If it's the same, add one on to the result for that letter
            if alphabet[letter] == char:
                results[letter] += 1
    # Divide all of the values by the length of the text and multiply by 100 to get the percentage
    for i in xrange(len(results)):
        results[i] = (float(results[i]) / len(text)) * 100.0
    # Return the results
    return results

def Abnormality(text):
    "Returns a value measuring the 'abnormality' of a piece of text\ncompared with normal English (in terms of letter frequency)"
    # Prepare the normal frequencies
    normal = [8.167, 1.492, 2.782, 4.253, 12.702, 2.228, 2.015, 6.094, 6.966, 0.153, 0.772, 4.025, 2.406, 6.749, 7.507, 1.929, 0.095, 5.987, 6.327, 9.056, 2.758, 0.978, 2.360, 0.150, 1.974, 0.074]
    # Get the actual frequencies in the text
    freq = Frequencies(text)
    # Prepare the result
    result = 0.0
    # Iterate through the frequencies summing the squares of the differences (squared to remove negatives)
    for i in xrange(len(normal)):
        result += (freq[i] - normal[i])**2
    # Return the result
    return result

def LeastAbnormal(text):
    "Returns the shift that gives the least abnormal value (same as solving the caesar cipher)"
    # Prepare the alphabet
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    # Prepare the store the best one
    best = ["*", 99999999999]
    # Loop through the alphabet performing each shift
    for letter in alphabet:
        shifted = VigDec(text, letter)
        if Abnormality(shifted) < best[1]:
            best[0] = letter
            best[1] = Abnormality(shifted)
    return best

def GetSubText(text, n, k):
    "Returns a string containing every nth character from the given text starting at the kth character"
    # Prepare result
    result = ""
    # Loop through text
    for i in xrange(len(text)):
        # If this is in the subtext, add it on
        # (that is, if i - k divides n evenly
        if (i - k) % n == 0:
            result += text[i]
    # Return the result
    return result

def VigCrack(ciphertext, keylength):
    "Attempts to crack a ciphertext encrypted with the Vigenere cipher and return the key, given the length of the key\nWorks considerably better for longer ciphertexts"
    # Prepare the key
    key = ""
    # Loop through each of the subtexts (the parts of the text encrypted with each character of the key)
    for i in xrange(keylength):
        # Get the subtext
        subtext = GetSubText(ciphertext, keylength, i)
        # Determine the character that works best for this
        nextchar = LeastAbnormal(subtext)[0]
        # Add this character to the key
        key += nextchar
    return key