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

# 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

def main(argv):
    l = 24

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

    for opt, arg in opts:
        if opt == '-h':
            usage()
            sys.exit()
        elif opt == '-l':
            l = int(arg)

    # generate two keys
    r = random.SystemRandom()
    k1 = nthKey(r.randint(0, pow(2, l) - 1))
    k2 = nthKey(r.randint(0, pow(2, l) - 1))
    
    print 'k_1: %s' % binascii.b2a_hex(k1)
    print 'k_2: %s' % binascii.b2a_hex(k2)

    round1 = DES.new(k1, DES.MODE_ECB)
    round2 = DES.new(k2, DES.MODE_ECB)

    # Some random plaintext - ciphertext pairs to verify the working of this code
    print 'plaintext        : ciphertext'
    for n in range(0,10):
        pt = '{:016x}'.format(r.randint(0, pow(2, 64) -1))
        ct = round2.encrypt(round1.encrypt(binascii.a2b_hex(pt)))
        print pt + ' : ' + binascii.b2a_hex(ct)

def usage():
    print 'Usage: 2des-demo.py [-l <keylength>]'
    print '  keylength : pick keys from the first 2^l keys (default: 24)'

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