Cryptography - HW2.1 (decipher.py)

import time
import sys
import oracle as O
import itertools

ORACLE_FALSE = 0
ORACLE_TRUE = 1

LOG = True

''' Save ranges for most common characters'''
# list of ranges
ords = [[ord(' ')],                                 # Space
        range(ord('a'), ord('z') + 1),              # LowerCase Letters
        range(ord('A'), ord('Z') + 1),              # UpperCase Letters
        range(ord('0'), ord('9') + 1),              # Numbers
        [44, 46, 33, 63, 40, 41, 39, 34],           # Punctuation Marks
        range(58, 63), range(35, 39),               # other common ranges
        [42, 43, 45, 47, 64],
        range(91, 97), range(123, 127)]

# convert to one whole range
common_chars_range = []
for x in ords:
    common_chars_range += x

''' Range for last block [1..15] '''
last_char_odds = range(1, 16)

''' Ranges left with no occurances '''
common_chars_set = set(common_chars_range)
all_chars_set = set(list(range(256)))
left_chars_set = all_chars_set - common_chars_set
left_chars_range = list(left_chars_set)


def _log(s):
    # only log when LOG variable is set on True
    if LOG:
        print s


def remove_pad(block):
    # in the last byte is number of padded bytes
    bytes_to_remove = block[-1]
    # remove last n bytes
    block_without_pad = block[: - bytes_to_remove]

    return block_without_pad


def convert_to_string(result):
    # remove padding
    result[-1] = remove_pad(result[-1])

    plain_text = ''
    for chunk in result:
        # convert integers to characters
        chunk_chars = list(map(chr, chunk))
        # get current string
        cur_str = ''.join(chunk_chars)
        # append to result
        plain_text += cur_str

    return plain_text


def decrypt_byte(IV, block, byte_ind, chunk_len, range_to_iter, to_xor=0):
    n = (chunk_len - byte_ind) ^ to_xor

    # function that takes range to iterate over
    def _iter_over_range(current_range):
        for curr_byte in current_range:
            IV[byte_ind] = curr_byte ^ n

            # check if it is correct decoding
            if O.Oracle_Send((IV+block), 2) == ORACLE_FALSE:
                continue

            # if not last byte return answer
            if byte_ind != chunk_len - 1:
                return curr_byte

            # if it's last byte increase IV[-2] byte by 1 (for checking purposes)
            IV_new = IV[:-2] + [(-IV[-2]+1) % 256] + [curr_byte ^ n]
            # if with this IV its true than last byte guess was rigth
            if O.Oracle_Send((IV_new + block), 2) == ORACLE_TRUE:
                return curr_byte
        return None

    byte_found = _iter_over_range(range_to_iter)
    # if we found result in common characters return result
    if byte_found:
        return byte_found
    # at this point we need to iterate over uncommon characters
    return _iter_over_range(left_chars_range)


def decrypt_chunk(IV, cur, chunk_len, is_last):
    # Initialize decrypted bytes with 0's
    decrypted_bytes = [0] * chunk_len

    # EXTRA (to work faster)
    # if it's last block then it has padding, so let's find only last block
    # and we will know that last n bytes will be same
    if is_last:
        # decrypt current byte
        cur_byte = decrypt_byte(
            IV[:], cur, chunk_len-1, chunk_len, last_char_odds, IV[(chunk_len-1)])

        # update i'th byte of
        n = decrypted_bytes[chunk_len-1] = cur_byte

        # assign values of n to last n elements of array
        for i in range(n):
            decrypted_bytes[-(i+1)] = n
        _log("Last Block --> Change {} Bytes : {}".format(n, decrypted_bytes))

    # Iterate and evaluate for each byt
    for i in reversed(range(chunk_len - decrypted_bytes[-1])):
        # Change IV (xor in such way that last n byte contains integer "n" in them)
        for j in range(i + 1, chunk_len):
            IV[j] = IV[j] ^ decrypted_bytes[j] ^ (chunk_len - i)

        # decrypt current byte
        cur_byte = decrypt_byte(
            IV[:], cur, i, chunk_len, common_chars_range, IV[i])

        # update i'th byte of
        decrypted_bytes[i] = cur_byte

        # change IV back to normal (as it was before)
        for j in range(i + 1, chunk_len):
            IV[j] = IV[j] ^ decrypted_bytes[j] ^ (chunk_len - i)

        _log('Change Byte #{:02d} : {}'.format(i, decrypted_bytes))
    return decrypted_bytes


def from_hex_to_integers(hex_cipher):
    int_array = []
    for i in range(0, len(hex_cipher), 2):
        # get current pair
        cur_hex = hex_cipher[i:i + 2]
        # convert to integer
        cur_int = int(cur_hex, 16)
        # append to list
        int_array += [cur_int]
    return int_array


def split_into_chunks(arr, n):
    chunk_arr = []
    for i in range(0, len(arr), n):
        # get current chunk
        chunk = arr[i:i + n]
        # append to list
        chunk_arr += [chunk]
    return chunk_arr


def break_block_cipher(ciphertext):
    ''' Step 0: Initialize connection '''
    O.Oracle_Connect()

    ''' Step 1: Convert pairs of hex numbers to array of integers'''
    int_array = from_hex_to_integers(ciphertext)
    _log('Integer Array : \n{}\n'.format(int_array))

    ''' Step 2: Split into chunks of 16 bytes '''
    chunks_list = split_into_chunks(int_array, 16)
    _log('Integer Chunks Array : \n{}\n'.format(chunks_list))

    ''' Step 3: Decrypt each chunk '''
    decrypted_chunks = []
    for i in range(len(chunks_list) - 1):
        _log("Decrypting Chunk #{}".format(i))
        # get current and next chunks
        IV = chunks_list[i]
        chunk = chunks_list[i + 1]
        # decrypt current chunk using next
        is_last = (i == len(chunks_list)-2)
        decrypted_chunk = decrypt_chunk(IV, chunk, 16, is_last)
        # append to result
        decrypted_chunks += [decrypted_chunk]
        _log('')

    ''' Step 4: Convert list of lists of ints to string to get plain text'''
    plain_text = convert_to_string(decrypted_chunks)

    ''' Step 5: Disconnect '''
    O.Oracle_Disconnect()

    return plain_text


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print('Should pass exactly 1 argument: python decipher.py <ciphertext_filename>')
        exit()

    '''Input '''
    # filename
    ciphertext_filename = sys.argv[1]
    # open and read file
    with open(ciphertext_filename, 'r') as ciphertext_file:
        ciphertext = ciphertext_file.read()

    ''' Break Code '''
    plain_text = break_block_cipher(ciphertext)

    ''' Output '''
    print(plain_text)