diff_crypt.py

import random
from collections import defaultdict


def generate_xor_pairs(target_xor, count):
    pairs = []
    used_numbers = set()

    while len(pairs) < count:
        plaintext1 = random.randint(0, 0xFFFF)
        if plaintext1 in used_numbers:
            continue
        plaintext2 = plaintext1 ^ target_xor
        if plaintext2 not in used_numbers:
            # ENCRYPT BEFORE STORING
            ciphertext1 = encrypt(plaintext1)
            ciphertext2 = encrypt(plaintext2)

            pairs.append((ciphertext1, ciphertext2))
            used_numbers.add(plaintext1)
            used_numbers.add(plaintext2)

    return pairs


def substitute(input, sbox):
    blocks = [(input >> shift) & 0xF for shift in (12, 8, 4, 0)]
    s_blocks = [sbox[block] for block in blocks]
    subbed = (s_blocks[0] << 12) | (s_blocks[1] << 8) | (s_blocks[2] << 4) | s_blocks[3]
    return subbed


def permute(input, pt):
    bits = f"{input:016b}"  # 16 bit bin representation
    # ex. table's first elem = 15, the 16tth bit is put into the 0th position
    permuted_bits = "".join(bits[pt[i]] for i in range(16))
    return int(permuted_bits, 2)


def en_round(prev_text, round_key, perm_bool):
    result = prev_text ^ round_key
    result = substitute(result, SBOX)
    if perm_bool:
        result = permute(result, PT)
    return result


def encrypt(plaintext):
    result = plaintext
    for i in range(3):
        result = en_round(result, ROUND_KEYS[i], True)
    result = en_round(result, ROUND_KEYS[3], False)
    result = result ^ ROUND_KEYS[4]
    return result


def d_round(prev_text, inv_sbox, inv_pt, round_key, perm_bool):
    result = prev_text ^ round_key  # inv round keys
    if perm_bool:
        result = permute(result, inv_pt)  # inv perm table
    result = substitute(result, inv_sbox)  # inv sbox
    return result


# ONE ROUND
def partial_decrypt(ciphertext1, ciphertext2, inv_sbox, inv_pt, mod_round_keys):
    ciphertext1 = d_round(ciphertext1, inv_sbox, inv_pt, mod_round_keys[4], False)
    ciphertext2 = d_round(ciphertext2, inv_sbox, inv_pt, mod_round_keys[4], False)
    partial_xor = ciphertext1 ^ ciphertext2
    return partial_xor


def generate_ddt(sbox):
    ddt = [[0] * 16 for x in [0] * 16]
    for i in range(16):
        for j in range(16):
            in_diff = i ^ j
            out_diff = sbox[i] ^ sbox[j]
            ddt[in_diff][out_diff] += 1
    return ddt


def find_ideal_diff_trails(ddt, permutation_table):
    results = {}
    final_results = {}
    for i in range(20000):

        results, round_XOR, round_prob = diff_trail(i, ddt, permutation_table)

        blocks = [(round_XOR >> shift) & 0xF for shift in (12, 8, 4, 0)]
        if blocks.count(0) == 2:
            final_results[i] = [round_XOR, round_prob]
            for j in range(4):
                if blocks[j] == 0:
                    final_results[i].append(j)
    # so final_results[input XOR] = [output XOR, probability, inactive nibble 1, inactive nibble 2]

    top_30 = sorted(final_results.items(), key=lambda x: x[1][1], reverse=True)[:30]

    first_input, first_list = top_30[0]
    print(
        f"1.   input XOR: {hex(first_input)},   output XOR position at u4: {hex(first_list[0])},   probability: {first_list[1]}"
    )
    for input, list in top_30[1:]:
        # find next highest trail that has opposite active boxes to the first
        if (
            list[2] != first_list[2]
            and list[2] != first_list[3]
            and list[3] != first_list[2]
            and list[3] != first_list[3]
        ):
            print(
                f"2.   input XOR: {hex(input)},   output XOR position at u4: {hex(list[0])},   probability: {list[1]}"
            )
            last_input = input
            last_list = list
            break
    return first_input, first_list, last_input, last_list


def diff_trail(i, ddt, permutation_table):

    results = {}

    results[i] = [[], []]
    results[i][0].append(i)  # adding x
    results[i][1].append(1)
    round_XOR = i
    round_prob = 1
    for j in range(3):

        results[i][0].append(round_XOR)  # adding u XOR
        results[i][0].append(round_prob)  # adding u prob

        blocks = [(round_XOR >> shift) & 0xF for shift in (12, 8, 4, 0)]

        box = 0
        s_outs = []
        for input_XOR in blocks:  # S-box
            output_XOR = ddt[input_XOR].index(max(ddt[input_XOR]))
            prob = ddt[input_XOR][output_XOR] / 16
            round_prob *= prob
            s_outs.append(output_XOR)
            box += 1
        round_XOR = (s_outs[0] << 12) | (s_outs[1] << 8) | (s_outs[2] << 4) | s_outs[3]
        results[i][0].append(round_XOR)
        results[i][1].append(round_prob)

        round_XOR = permute(round_XOR, permutation_table)
        results[i][0].append(round_XOR)
        results[i][1].append(round_prob)

    return results, round_XOR, round_prob


def find_partial_key(input_xor, target_xor, key_space, inv_sbox, inv_pt):

    pairs = generate_xor_pairs(input_xor, 10000)
    count_table = defaultdict(int)

    for ciphertext1, ciphertext2 in pairs:
        for partial_key in key_space:
            modified_round_keys = ROUND_KEYS[:]
            # replace 5th key (which is being search for) with partial key
            modified_round_keys[4] = partial_key

            partial_xor = partial_decrypt(
                ciphertext1, ciphertext2, inv_sbox, inv_pt, modified_round_keys
            )

            if partial_xor == target_xor:
                count_table[partial_key] += 1

    total_keys_found = sum(count_table.values())
    table = sorted(
        ((key, count / total_keys_found) for key, count in count_table.items()),
        key=lambda x: x[1],
        reverse=True,
    )

    print(f"\n {'Key Byte':>10} | {'Probability':>10}")
    print("-" * 23)
    for key_byte, prob in table[:28]:  # CHANGE to get more partial keys + probabilities
        print(f"{hex(key_byte):>10} | {prob:>10.4f}")

    return table[0]


def gen_keys(inactive_nib1, inactive_nib2):
    # generates partial keys
    key_space = []
    indices = [0, 1, 2, 3]
    # remove inactives, leaving only the active nibbles
    indices = [indices[i] for i in range(4) if i not in [inactive_nib1, inactive_nib2]]
    shift1 = 12 - (indices[0] * 4)
    shift2 = 12 - (indices[1] * 4)
    for nibble1 in range(16):
        for nibble2 in range(16):
            partial_key = (nibble1 << shift1) | (nibble2 << shift2)
            key_space.append(partial_key)
    return key_space


# CONSTANTS

SBOX = [7, 2, 6, 9, 1, 3, 13, 0, 12, 5, 10, 15, 8, 4, 14, 11]

PT = [3, 12, 0, 5, 10, 7, 14, 1, 9, 2, 11, 6, 8, 4, 13, 15]

ROUND_KEYS = [12345, 32766, 2001, 15000, 5432]


def main():
    inv_sbox = [SBOX.index(i) for i in range(16)]
    inv_pt = [PT.index(i) for i in range(16)]

    print("\n", "///////" * 3, "S-BOX USED: ", "///////" * 3, "\n")

    print(SBOX)

    print("\n", "///////" * 3, "GENERATING THE DDT", "///////" * 3, "\n")

    ddt = generate_ddt(SBOX)
    for row in ddt:
        print(row)
    print()

    print(
        "\n",
        "///////" * 3,
        "FINDING PROBABLE DIFFERENTIAL CHARACTERISTICS",
        "///////" * 3,
        "\n",
    )
    probs = {}
    for i in range(16):
        for j in range(16):
            if ddt[i][j] >= 6:
                probs[i] = [j, ddt[i][j]]
                print(f"{ddt[i][j]}/16:   {i} --> {j}")

    print(
        "\n",
        "///////" * 3,
        "FINDING MOST PROBABLE DIFFERENTIAL TRAILS",
        "///////" * 3,
        "\n",
    )

    dt_input1, dt_list1, dt_input2, dt_list2 = find_ideal_diff_trails(ddt, PT)

    print("\n", "///////" * 3, "FINDING KEY BYTES", "///////" * 3, "\n")

    key_sp_1 = gen_keys(dt_list1[2], dt_list1[3])
    key_sp_2 = gen_keys(dt_list2[2], dt_list2[3])

    kb1 = find_partial_key(dt_input1, dt_list1[0], key_sp_1, inv_sbox, inv_pt)
    kb2 = find_partial_key(dt_input2, dt_list2[0], key_sp_2, inv_sbox, inv_pt)

    print("\n", "///////" * 3, "THE PARTIAL KEYS FOUND ARE", "///////" * 3, "\n")
    print(f"{hex(kb1[0])}   with a probability of   {kb1[1]}")
    print(f"{hex(kb2[0])}   with a probability of   {kb2[1]}")

    print("\n", "///////" * 3, "THUS, THE FIFTH KEY MUST BE ", "///////" * 3, "\n")
    print(hex(kb1[0] | kb2[0]))
    print("aka: ", kb1[0] | kb2[0])


main()