import pyDes
import random

# Compute the Hamming weight of a byte
def hw(b):
    w = 0
    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 hw(b) % 2 == 0:
        return b ^ 0x01
    return b

def get_key(n):
    b1 = set_parity((n >> 14) % 256)
    b2 = set_parity((n >> 6) % 256)
    b3 = set_parity((n % 64) << 2)

    ret = bytearray([b1, b2, b3, 0, 0, 0, 0, 0])
    return bytes(ret)

def des_enc(data, key):
    return pyDes.des(key).encrypt(data)

def des_dec(data, key):
    return pyDes.des(key).decrypt(data)

def hack_double_des(plaintext, ciphertext):
    key_dict_e = {}
    key_dict_d = {}

    key_deck = range(pow(2, 20))
    random.shuffle(key_deck)

    for i, j in enumerate(key_deck):
        k = get_key(j)

        if i % 1024 == 0:
            print("%d/1024" % (i // 1024))

        c = des_enc(plaintext, k)
        if c in key_dict_d.keys():
            return k, key_dict_d[c]
        else:
            key_dict_e[c] = k

        m = des_dec(ciphertext, k)
        if m in key_dict_e.keys():
            return key_dict_e[m], k
        else:
            key_dict_d[m] = k

    # This should never happen
    return b'lol', b'noob'

k1 = get_key(1234)
k2 = get_key(5432)

plaintext = b'ABCDEFGH'
middle = des_enc(plaintext, k1)
ciphertext = des_enc(middle, k2)

print("k1: %r" % k1)
print("k2: %r" % k2)

print("Plaintext:\t%r" % plaintext)
print("Ciphertext:\t%r" % ciphertext)

keys = hack_double_des(plaintext, ciphertext)

print("Guess 1: %r" % keys[0])
print("Guess 2: %r" % keys[1])

print("TEST:")
res1 = des_enc(plaintext, keys[0])
res2 = des_dec(ciphertext, keys[1])
print("enc(p):\t\t%r" % res1)
print("dec(c):\t\t%r" % res2)
print("Expected:\t%r" % middle)

if res1 == res2 and keys[0] == k1 and keys[1] == k2:
    print("SUCCESS")