import pyDes
import random

# Compute the Hamming weight of a byte
def hw(b):
    w = 0
    # For each bit in a byte, shift one to the
    # right and add 1 if the least significant
    # bit is set.
    for i in range(8):
        w += (b >> i) & 1
    return w

# Sets the key parity bit of a DES key byte
def set_parity(b):
    # If Hamming weight is even,
    # flip least-significant bit
    if hw(b) % 2 == 0:
        return b ^ 0x01
    # Else, return b as-is.
    return b

# Take a 20-bit number and generate a DES key with
# parity bits in significant bytes, followed by
# zero-padding.
def get_key(n):
    # 20 = 2*8+4. Subtracting parity bits gives
    # effective key bits: 2*7+6, so we shift
    # each whole byte by 7 * the number of less
    # significant bytes + effective length of least
    # significant byte, take modulo 2^7 = 128 and
    # finally shift one to the left to make space
    # for the parity bit.

    # Byte 3 (least-significant byte): take modulo
    # 2^l, where l is the size of the partial byte
    # in bits, and shift by 8-l. This gives 8-6=2.
    b3 = set_parity((n % 64) << 2)
    # Byte 2: Shift by combined length of less
    # signficiant bytes: 6
    b2 = set_parity(((n >> 6) % 128) << 1)
    # Byte 3: Shift by combined length of less
    # signficiant bytes: 7+6=13
    b1 = set_parity(((n >> 13) % 128) << 1)

    # Assemble the generated bytes into a
    # bytearray object.
    key = bytearray([b1, b2, b3, 0, 0, 0, 0, 0])

    # Return the key as a bytes object for DES
    # library compatibility.
    return bytes(key)

# Wrapper function for single DES encryption.
def des_enc(data, key):
    # Initialize DES cipher object using the
    # provided key and use it to encrypt the
    # provided data.
    return pyDes.des(key).encrypt(data)

# Wrapper function for single DES decryption.
def des_dec(data, key):
    # Initialize DES cipher object using the
    # provided key and use it to decrypt the
    # provided data.
    return pyDes.des(key).decrypt(data)

# This is the heart of the operation. This
# function takes a plaintext-ciphertext pair,
# along with a key size (in bits), and performs
# a meet-in-the-middle attack.
def hack_double_des(plaintext, ciphertext, key_size=56):
    # Initialize two dictionary objects, one for
    # each side of the attack. These will contain
    # the keys already tested, indexed by the
    # values resulting from encryption and
    # decryption, respectively.
    key_dict_e = {}
    key_dict_d = {}

    # TODO: implement an LC pseudo-random range generator function
    # To efficiently iterate through huge ranges without having to
    # allocate so much memory.
    key_deck = range(pow(2, key_size))
    #random.shuffle(key_deck)

    # Compute the size of a thousandth step over the keyspace.
    # This is used to display the progress in terms of keyspace
    # coverage.
    permille = pow(2, key_size) // 1000

    # Randomly cycle through the keyspace.
    for i, j in enumerate(key_deck):
        # Generate DES key from current random value. This is the
        # key guess for the current iteration.
        k = get_key(j)

        # If the current index is a whole multiple of a thousandth
        # step over the keyspace, print the current coverage
        # percentage.
        if i % permille == 0:
            print("{}% of keyspace covered".format((i / permille)/10.0))

        # Encrypt the plaintext using the current key guess.
        c = des_enc(plaintext, k)
        # If the same intermediate value has been found by decrypting
        # the ciphertext, return the current key guess as k1 and the
        # key stored at the index of the intermediate value as k2.
        if c in key_dict_d.keys():
            return k, key_dict_d[c]
        # Else, Store the current key guess at the index of the
        # intermediate value.
        else:
            key_dict_e[c] = k

        # Decrypt the ciphertext using the current key guess.
        m = des_dec(ciphertext, k)
        # If the same intermediate value has been found by encrypting
        # the plaintext, return the key stored at the index of the
        # intermediate value as k1 and the current key guess as k2.
        if m in key_dict_e.keys():
            return key_dict_e[m], k
        # Else, Store the current key guess at the index of the
        # intermediate value.
        else:
            key_dict_d[m] = k

    # Abandon hope; all is lost. This should never happen.
    return b'lol     ', b'nope    '

# Define some integers in the range [1 .. 2^20] and use them
# to generate DES keys.
k1 = get_key(193)
k2 = get_key(5523)

# Define the plaintext for testing. This is known to the attacker.
plaintext = b'ABCDEFGH'
# Encrypt once using k1 and store for later use.
middle = des_enc(plaintext, k1)
# Encrypt again using k2 and store as ciphertext. This is known to
# the attacker.
ciphertext = des_enc(middle, k2)

# Print the actual keys for ease of analysis.
print("Actual k1: %r" % k1)
print("Actual k2: %r" % k2)

# Print the plaintext and ciphertext for ease of analysis.
print("Plaintext:\t%r" % plaintext)
print("Ciphertext:\t%r" % ciphertext)

# Hack the planet! But mostly double DES, using a 20-bit keyspace.
keys = hack_double_des(plaintext, ciphertext, 20)

# Print the keys identified by the hack function for ease of analysis.
print("Guess k1: %r" % keys[0])
print("Guess k2: %r" % keys[1])

# Generate the intermediate value produced using each of the
# identified keys.
res1 = des_enc(plaintext, keys[0])
res2 = des_dec(ciphertext, keys[1])

# Assert that the intermediate values are identical.
assert res1 == res2
# Assert that the first identified key is identical to k1.
assert keys[0] == k1
# Assert that the second identified key is identical to k2.
assert keys[1] == k2

# Inform the user that all tests have passed.
print("All tests passed!")