Cryptography - HW1.6 (Challenge-06)

from base64 import b64decode
from itertools import zip_longest
import math
import binascii


def hamming_distance(s1, s2):
    diffs = 0
    for ch1, ch2 in zip(s1, s2):
        if ch1 != ch2:
            diffs += 1
    return diffs


def score(s):
    # frequencies taken from web
    freq = {'a': 6517, 'b': 1242, 'c': 2173, 'd': 3498, 'e': 10414, 'f': 1978, 'g': 1586, 'h': 4928, 'i': 5580, 'j': 90,
            'k': 505, 'l': 3314, 'm': 2021, 'n': 5645, 'o': 5963, 'p': 1376, 'q': 86, 'r': 4975, 's': 5157, 't': 7293,
            'u': 2251, 'v': 829, 'w': 1712, 'x': 136, 'y': 1459, 'z': 78, ' ': 19181}
    # convert string to lowercase
    s = s.lower()
    # sum up frequencies
    result = sum([freq[ch] if (ch in freq) else 0 for ch in s])

    return result


def solve_multiple_xor(transposed_blocks):
    arr = []
    # iterate and solve for each block
    for x in transposed_blocks:
        # conert to characters and join to make string
        l = ''.join(list(map(chr, x)))
        # solve and append to array of solutions
        arr.append(solve_single_xor(l))

    return arr


def solve_single_xor(s):
    # iterate over all characters
    result_arr = []
    for ind in range(256):
        # xor each char
        result = ''
        for t in s:
            r = ord(t) ^ ind
            result += chr(r)
        # if given text is built with regex return answer
        cur_score = score(result)
        result_arr += [(cur_score, result)]
    # choose max from results
    final_result = max(result_arr)
    # return only string answer
    return final_result[1]


def transpose(chunks):
    # transpose matrix (padded with None values to fill)
    result = list(zip_longest(*chunks))
    # remove None values from list
    result = [list(filter(lambda a:a != None, x)) for x in result]
    # return list of lists
    return result


def _eval_score_for_key(cipher, K_len):
    score = 0
    # for each pair find hamming distance
    for i in range((len(cipher) // K_len) - 1):
        # get pairs
        s1 = cipher[i * K_len:(i + 1) * K_len]
        s2 = cipher[(i + 1) * K_len:(i + 2) * K_len]
        # evaluate current score (normalize with key size)
        result = hamming_distance(s1, s2) / K_len
        # append to score
        score += result
    # normalize final score with total pairs evaluated
    score /= (len(cipher) // K_len) - 1
    # return score and current K_len
    return (score, K_len)


def get_key_size(s):
    # decode and encode with given base64
    s = binascii.a2b_base64(s.encode()).decode()
    # given ranges
    R_MIN = 2
    R_MAX = 41
    # calculate score for each key size between 2 and 41
    scores_array = [_eval_score_for_key(s, x) for x in range(R_MIN, R_MAX)]
    # we say that it's minimum will be our result
    _, K_len = min(scores_array)
    # return key size
    return K_len


def break_repeating_key_xor(cipher):
    """ Step 1: Find Key Size """
    K_len = get_key_size(cipher)

    """ Step 2: Decode from Base64"""
    cipher = b64decode(cipher)

    """ Step 3: break the ciphertext into blocks of KEYSIZE length. """
    blocks = [cipher[i:i + K_len]
              for i in range(0, len(cipher), K_len)]

    """ Step 4: Now transpose the blocks:
                make a block that is the first byte of every block,
                and a block that is the second byte of every block, and so on... """
    transposed_blocks = transpose(blocks)

    """ Step 5: Solve each block as if it was single-character XOR. """
    key = solve_multiple_xor(transposed_blocks)

    """ Step 6: Transpose back to normal """
    result = transpose(key)

    """ Step 7: convert list of lists to string """
    result = list(map(lambda x: ''.join(x), result))
    result = ''.join(result)

    return result


if __name__ == "__main__":
    # input key
    cipher = input()
    # break cipher
    plain_text = break_repeating_key_xor(cipher)
    # ouput decrypted
    print(plain_text)