#!/usr/bin/python
import sys, getopt, binascii, random, time, os.path, time
from Crypto.Cipher import DES

dictionary = {}

# From http://stackoverflow.com/a/843846/1544337
# 1 for odd parity, 0 for even parity
# Was used only in a previous version of nthByte
# def parity(b):
#     c = 0
#     while b != 0:
#         c += 1
#         b &= b - 1
#     return c % 2

# Find the nth possibility for a byte with odd parity
oddBytes = [1,2,4,7,8,11,13,14,16,19,21,22,25,26,28,31,32,35,37,38,41,42,44,47,49,50,52,55,56,59,61,62,64,67,69,70,73,74,76,79,81,82,84,87,88,91,93,94,97,98,100,103,104,107,109,110,112,115,117,118,121,122,124,127,128,131,133,134,137,138,140,143,145,146,148,151,152,155,157,158,161,162,164,167,168,171,173,174,176,179,181,182,185,186,188,191,193,194,196,199,200,203,205,206,208,211,213,214,217,218,220,223,224,227,229,230,233,234,236,239,241,242,244,247,248,251,253,254]
def nthByte(n):
    return oddBytes[n]
    # c = -1                        # This is the old version. The new version uses a faster lookup table.
    # b = 0
    # while c != n:
    #     b += 1
    #     if parity(b) == 1:
    #         c += 1
    # return b

# Find the nth key in which all bytes have odd parity
def nthKey(n):
    key = ''
    for b in range(7,-1,-1):
        key += chr(nthByte((n >> b*7) & 0x7f))
    return key

# Create a dictionary for the first 2^l keys and a given plaintext
def createDictionary(plaintext, l):
    for n in range(0, pow(2, l)):
        encryption = DES.new(nthKey(n), DES.MODE_ECB).encrypt(plaintext)
        if encryption in dictionary:
            dictionary[encryption].append(n)
        else:
            dictionary[encryption] = [n]

# Read the dictionary from a file if it exists, or generate it
def readOrMakeDictionary(plaintext, l):
    if os.path.isfile('dict-' + binascii.b2a_hex(plaintext) + '-' + str(l) + '.txt'):
        with open('dict-' + binascii.b2a_hex(plaintext) + '-' + str(l) + '.txt') as f:
            for line in f:
                dictionary[binascii.a2b_hex(line[:16])] = [int(x) for x in line[17:-1].split(',')]
        return False
    else:
        createDictionary(plaintext, l)
        return True

# Save the dictionary to a file
def saveDictionary(plaintext, l):
    dictionary_file = open('dict-' + binascii.b2a_hex(plaintext) + '-' + str(l) + '.txt', 'w')
    for encryption in dictionary:
        dictionary_file.write(binascii.b2a_hex(encryption) + ':' + ','.join(str(x) for x in dictionary[encryption]) + '\n')
    print 'Saved dictionary as dict-' + binascii.b2a_hex(plaintext) + '-' + str(l) + '.txt'

def main(argv):
    l = 24
    plaintext = ''
    ciphertext = ''
    save_dictionary = False

    # Parse arguments. Use -h for help.
    try:
        opts, args = getopt.getopt(argv, "hl:p:c:s")
    except getopt.GetoptError:
        usage()
        sys.exit(2)

    for opt, arg in opts:
        if opt == '-h':
            usage()
            sys.exit()
        elif opt == '-l':
            l = int(arg)
        elif opt == '-p':
            plaintext = binascii.a2b_hex(arg)
        elif opt == '-c':
            ciphertext = binascii.a2b_hex(arg)
        elif opt == '-s':
            save_dictionary = True

    assert plaintext != ''
    assert ciphertext != ''

    # Make the dictionary if it doesn't exist yet
    print 'Making dictionary (p:' + binascii.b2a_hex(plaintext) + ';l:' + str(l) + ')...',
    timer = time.clock()
    if not readOrMakeDictionary(plaintext, l):
        save_dictionary = False
    print str(time.clock() - timer) + 's'

    if save_dictionary:
        saveDictionary(plaintext, l)

    # Find matches
    print 'Finding matches...',
    timer = time.clock()
    time_key = time_decryption = time_matching = 0
    matches = []
    for k in range(0, pow(2,l)):
        # Generate key
        time_t = time.clock()
        key = nthKey(k)
        time_key += time.clock() - time_t
        
        # Decrypt
        time_t = time.clock()
        des = DES.new(key, DES.MODE_ECB)
        decryption = des.decrypt(ciphertext)
        time_decryption += time.clock() - time_t
        
        # Find matches
        time_t = time.clock()
        if decryption in dictionary:
            [matches.append({'k1': nthKey(i), 'k2': key}) for i in dictionary[decryption]]
        time_matching += time.clock() - time_t
    print str(time.clock() - timer) + 's'

    print 'Key generation:      ' + str(time_key) + 's'
    print 'Decryption:          ' + str(time_decryption) + 's'
    print 'Matching dictionary: ' + str(time_matching) + 's'

    for match in matches:
        print 'k1: ' + binascii.b2a_hex(match['k1']) + '; k2: ' + binascii.b2a_hex(match['k2'])

def usage():
    print 'Usage: break.py -p <plaintext> -c <ciphertext> [-l <keylength>] [-s]'
    print '  plaintext  : plaintext of the known-plaintext attack'
    print '  ciphertext : ciphertext of the known-plaintext attack'
    print '  keylength  : pick keys from the first 2^l keys (default: 24)'
    print '  -s         : save the dictionary for this plaintext and keylength to a file'

if __name__ == "__main__":
    main(sys.argv[1:])