Cryptography - HW2.2 (forge.py)

import itertools
import oracle as O
import time
import sys

CHUNK_LENGTH = 16


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 verify(message, tag):
    ret = O.Vrfy(message, len(message), tag)
    if ret == 1:
        print "Tag verified!"
    else:
        print "Tag Rejected!"


def break_CBC_MAC(message):
    ''' Step 0 : Connect to server '''
    O.Oracle_Connect()

    ''' Step 1 : Convert characters to integers'''
    message = list(map(ord, message))

    ''' Step 2 : Divide into chunks'''
    chunks = split_into_chunks(message, CHUNK_LENGTH)

    '''' Step 3 : For each pair evaluate new Mac '''
    # starting mac filled with zeros
    tag = bytearray([0]*CHUNK_LENGTH*2)
    for i in range(0, len(chunks), 2):
        # get pairs of chunks
        curr_chunk = chunks[i]
        next_chunk = chunks[i + 1]

        # create new msg with xoring to old tag
        xord_chunk = [curr_chunk[i] ^ tag[i] for i in range(CHUNK_LENGTH)]

        # evaluate new tag
        tag = O.Mac(xord_chunk + next_chunk, 2 * CHUNK_LENGTH)

    ''' Step 5: Verify Our Tag'''
    verify(message, tag)

    ''' Step 6: Disconnect from server '''
    O.Oracle_Disconnect()

    return tag


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

    '''Input '''
    # filename
    message_filename = sys.argv[1]
    # open and read file
    with open(message_filename, 'r') as message_file:
        message = message_file.read()

    ''' Break Code '''
    plain_text = break_CBC_MAC(message)

    ''' Output '''
    print(plain_text)