Attack Vigenère cipher

attack_Vigenère.py

CIPHER = '''
KCCPKBGUFDPHQTYAVINRRTMVGRKDNBVFDETDGILTXRGUDDKOTFMBPVGEGLTGCKQRACQCWDNAWCRXIZ
AKFTLEWRPTYCQKYVXCHKFTPONCQQRHJVAJUWETMCMSPKQDYHJVDAHCTRLSVSKCGCZQQDZXGSFRLSWC
WSJTBHAFSIASPRJAHKJRJUMVGKMITZHFPDISPZLVLGWTFPLKKEBDPGCEBSHCTJRWXBAFSPEZQNRWXC
VYCGAONWDDKACKAWBBIKFTIOVKCGGHJVLNHIFFSQESVYCLACNVRWBBIREPBBVFEXOSCDYGZWPFDTKF
QIYCWHJVLNHIQIBTKHJVNPIST
'''
CIPHER = ''.join(CIPHER.split())
frequency = {
    'A': 0.08167,
    'B': 0.01492,
    'C': 0.02782,
    'D': 0.04253,
    'E': 0.12702,
    'F': 0.02228,
    'G': 0.02015,
    'H': 0.06094,
    'I': 0.06966,
    'J': 0.00153,
    'K': 0.00772,
    'L': 0.04025,
    'M': 0.02406,
    'N': 0.06749,
    'O': 0.07507,
    'P': 0.01929,
    'Q': 0.00095,
    'R': 0.05987,
    'S': 0.06327,
    'T': 0.09056,
    'U': 0.02758,
    'V': 0.00978,
    'W': 0.02360,
    'X': 0.00150,
    'Y': 0.01974,
    'Z': 0.00074
}

RECTA = [[chr((i + j) % 26 + 65) for i in range(0, 26)] for j in range(0, 26)]


def calPi(cipher):
    d = {}
    for i in range(0, 26):
        d[chr(i+65)] = 0
    for c in cipher:
        d[c] += 1
    l = len(cipher)
    s = sum([d[c]*(d[c] - 1) for c in d])
    return s/(l*(l-1))


def guess_key(cipher, ic_storage, key_limit):
    for t in range(1, key_limit):
        cipher_list = []
        for i in range(0, t):
            cipher_list.append(cipher[i::t])
        pi_values = [calPi(s) for s in cipher_list]
        average = sum(pi_values)/len(pi_values)
        print(t, "Average:", average, "Diff:", abs(average - 0.065))
        ic_storage.append([t, abs(average - 0.065)])
    ic_storage = sorted(ic_storage, key=lambda x: x[1])
    return ic_storage


def decrypt(cipher, key):
    plain = ''
    for i in range(0, len(cipher)):
        column = RECTA[ord(key[i % len(key)]) - 65]
        row_id = column.index(cipher[i])
        plain += chr(row_id + 65)
    return plain


if __name__ == '__main__':
    ic_storage = []
    ic_storage = guess_key(CIPHER, ic_storage, 10)
    key = ''
    for ic in ic_storage:
        t = ic[0]
        cipher_list = []
        key = ''
        for i in range(0, t):
            cipher_list.append(CIPHER[i::t])
        print("Possible key: ", end='')
        for s in cipher_list:
            d = {}
            for i in range(0, 26):
                d[chr(i+65)] = 0
            for c in s:
                d[c] += 1
            probability = []
            l = len(s)
            for k in range(0, 26):
                p = sum([d[chr((i+k) % 26+65)]/l*(frequency[chr(i+65)])
                        for i in range(0, 26)])
                probability.append([p, k])
            probability = sorted(probability, key=lambda x: x[0], reverse=True)
            key += chr(probability[0][1]+65)
            print(chr(probability[0][1]+65), end='')
        print()
        plain_text = decrypt(CIPHER, key)
        print("Possible plain text:")
        print(plain_text)
        flag = input("Do you want to continue? (y/[n]): ")
        if flag != 'n':
            break
    print("Key:", key)