Eugene Tan

"""
DES cipher module

Usage:
>>> plaintext = 0x0123456789abcdef
>>> key = 0x3b3898371520f75e

Encryption:
>>> ciphertext = bin_array_to_int(des(plaintext, key, debug=False))
>>> hex(ciphertext)
'0xaa39b9777efc3c14L'

Decryption:
>>> decrypted = des(ciphertext, key, decrypt=True, debug=False)
>>> hex(bin_array_to_int(decrypted))
'0x123456789abcdefL'

Note that python removes zeros in front
"""

# initial permutation
IP = [
    58, 50, 42, 34, 26, 18, 10, 2,
    60, 52, 44, 36, 28, 20, 12, 4,
    62, 54, 46, 38, 30, 22, 14, 6,
    64, 56, 48, 40, 32, 24, 16, 8,
    57, 49, 41, 33, 25, 17,  9, 1,
    59, 51, 43, 35, 27, 19, 11, 3,
    61, 53, 45, 37, 29, 21, 13, 5,
    63, 55, 47, 39, 31, 23, 15, 7
]

# final permutation IP^-1
FP = [
    40, 8, 48, 16, 56, 24, 64, 32,
    39, 7, 47, 15, 55, 23, 63, 31,
    38, 6, 46, 14, 54, 22, 62, 30,
    37, 5, 45, 13, 53, 21, 61, 29,
    36, 4, 44, 12, 52, 20, 60, 28,
    35, 3, 43, 11, 51, 19, 59, 27,
    34, 2, 42, 10, 50, 18, 58, 26,
    33, 1, 41,  9, 49, 17, 57, 25
]

# expansion
E = [
    32,  1,  2,  3,  4,  5,
     4,  5,  6,  7,  8,  9,
     8,  9, 10, 11, 12, 13,
    12, 13, 14, 15, 16, 17,
    16, 17, 18, 19, 20, 21,
    20, 21, 22, 23, 24, 25,
    24, 25, 26, 27, 28, 29,
    28, 29, 30, 31, 32, 1
]

# permutation choice 1
PC1 = [
    57, 49, 41, 33, 25, 17,  9,
     1, 58, 50, 42, 34, 26, 18,
    10,  2, 59, 51, 43, 35, 27,
    19, 11,  3, 60, 52, 44, 36,
    63, 55, 47, 39, 31, 23, 15,
     7, 62, 54, 46, 38, 30, 22,
    14,  6, 61, 53, 45, 37, 29,
    21, 13,  5, 28, 20, 12,  4
]

# permutation choice 2
PC2 = [
    14, 17, 11, 24,  1,  5,
     3, 28, 15,  6, 21, 10,
    23, 19, 12,  4, 26,  8,
    16,  7, 27, 20, 13,  2,
    41, 52, 31, 37, 47, 55,
    30, 40, 51, 45, 33, 48,
    44, 49, 39, 56, 34, 53,
    46, 42, 50, 36, 29, 32
]

###########
# S-Boxes #
###########
S1 = [
    [14,  4, 13, 1,  2, 15, 11,  8,  3, 10,  6, 12,  5,  9, 0,  7],
     [0, 15,  7, 4, 14,  2, 13,  1, 10,  6, 12, 11,  9,  5, 3,  8],
     [4,  1, 14, 8, 13,  6,  2, 11, 15, 12,  9,  7,  3, 10, 5,  0],
    [15, 12,  8, 2,  4,  9,  1,  7,  5, 11,  3, 14, 10,  0, 6, 13]
]

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

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

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

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

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

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

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

SBOX = [S1, S2, S3, S4, S5, S6, S7, S8]

# P-Box
P = [
    16,  7, 20, 21,
    29, 12, 28, 17,
     1, 15, 23, 26,
     5, 18, 31, 10,
     2,  8, 24, 14,
    32, 27,  3,  9,
    19, 13, 30,  6,
    22, 11,  4, 25
]

def s(input):
    """Run sboxes on input, returning the result

    >>> input = pad(int_to_bin_array(0b001001100001110100011001001011111001101000011010), 48)
    >>> pprint(s(input), 4)
    '1110 1101 0010 0001 0111 0110 1100 0000'
    """
    buf = input
    sbox_result = []
    for i in range(len(SBOX)):
        s_input, buf = split(buf, 6)
        s_val = sbox(s_input, SBOX[i])
        sbox_result += pad(int_to_bin_array(s_val), 4)
    return sbox_result

def sbox(input, s_table):
    """Substitute block according to given sbox

    >>> sbox(list('111001'), S8)
    3
    """
    col = bin_array_to_int(input[1:5])
    row = bin_array_to_int(input[0] + input[-1])
    return s_table[row][col]

def permutate(bits, perm_table):
    """Permutate block according to given table

    Uses perm tables which start from index 1

    >>> val = pad(int_to_bin_array(0x675a69675e5a6b5a), 64)
    >>> hex(bin_array_to_int(permutate(val, IP)))
    '0xffb2194d004df6fbL'
    """
    #init y
    y = []
    for i in range(len(perm_table)):
        y.append(0)

    # permutate
    for i in range(len(perm_table)):
        index = perm_table[i]
        y[i] = bits[index - 1]
    return y

def bin_array_to_int(bits):
    """Converts a bit array to an integer

    >>> bin_array_to_int(list('1011'))
    11
    """
    return int(''.join(bits), 2)

def int_to_bin_array(x):
    """Converts an integer to a bit array

    >>> int_to_bin_array(11)
    ['1', '0', '1', '1']
    """
    return list(bin(x))[2:]

def split(bits, pos):
    """Convenience method for splitting a list

    >>> split([1, 2, 3, 4, 5, 6], 3)
    [[1, 2, 3], [4, 5, 6]]
    """
    return [bits[0:pos], bits[pos:]]

def c_left_shift(bits, val):
    """Circular left shift by given value

    >>> pprint(c_left_shift(list('0100010011000000011010111101'), 1), 7)
    '1000100 1100000 0011010 1111010'
    >>> pprint(c_left_shift(list('0001001100000001101011110101'), 2), 7)
    '0100110 0000001 1010111 1010100'
    >>> pprint(c_left_shift(list('0010011101100010000111111111'), 2), 7)
    '1001110 1100010 0001111 1111100'
    """
    start = len(bits) - 1
    for i in range(val):
        temp = bits[0]
        for i in range(start):
            bits[i] = bits[i + 1]
        bits[-1] = temp
    return bits

def c_right_shift(bits, val):
    """Circular right shift by given value

    >>> c_right_shift(list('011001'), 1)
    ['1', '0', '1', '1', '0', '0']
    >>> c_right_shift(list('011001'), 2)
    ['0', '1', '0', '1', '1', '0']
    """
    start = len(bits) - 1
    for i in range(val):
        temp = []
        for i in range(start, 0, -1):
            if i > start - 1:
                temp.append(bits[i])
            bits[i] = bits[i - 1]
        for i in range(len(temp)):
            bits[i] = temp[i]
    return bits

def pad(x, num, val='0'):
    """Left-pad bit array to desired number of bits

    >>> pad(list('0110'), 5)
    ['0', '0', '1', '1', '0']
    >>> pad(list('0110'), 5, '1')
    ['1', '0', '1', '1', '0']
    """
    return ([val] * (num - len(x))) + x

def xor(a, b):
    """Takes in two binary arrays and XORs them

    >>> xor(list('0110'), list('1100'))
    ['1', '0', '1', '0']
    """
    xor_result = bin_array_to_int(a) ^ bin_array_to_int(b)
    return int_to_bin_array(xor_result)

def pprint(x, val=8):
    """Returns formatted string from bit array, split by val

    >>> pprint(list('00001111'), 4)
    '0000 1111'
    """
    formatted_string = ''
    for i in range(len(x)):
        if i != 0 and i % val == 0:
            formatted_string += " "
        formatted_string += x[i]
    return formatted_string

def f(right, key, debug=True):
    """f-function for DES

    >>> right = int_to_bin_array(0b11110000101010101111000010101010)
    >>> key = int_to_bin_array(0b010111000000100001001100010101011000111101001111)
    >>> pprint(f(right, key, False))
    '10100000 10111100 11000010 10011101'
    """
    # expansion step
    e = permutate(right, E)
    if debug:
        print "expansion        : ", pprint(e, 6), 'bits', len(e)

    # key mixing
    xor_key_and_e = pad(xor(e, key), 48)
    if debug:
        print "xor_key_and_e    : ", pprint(xor_key_and_e, 6), 'bits', len(xor_key_and_e)

    # round substitution
    sbox_result = s(xor_key_and_e)
    if debug:
        print "sbox_result      : ", pprint(sbox_result, 4), 'bits', len(sbox_result)

    # round permutation
    pbox_result = permutate(sbox_result, P)
    if debug:
        print "pbox_result      : ", hex(bin_array_to_int(pbox_result)), 'bits', len(pbox_result)
    return pbox_result

def key_schedule(c, d, round_index, shift_function, debug=True):
    """Shift according to shift-function and permutate key schedule.

    Returns PC2, C and D in an array

    >>> c = pad(int_to_bin_array(0b0001001100000001101011110101), 28)
    >>> d = pad(int_to_bin_array(0b0010011101100010000111111111), 28)
    >>> round_index = 2
    >>> ary = key_schedule(c, d, round_index, c_left_shift, False)
    >>> pprint(ary[0], 6)
    '110101 001110 010010 000101 110110 001011 010011 101111'
    >>> pprint(ary[1] + ary[2], 7)
    '0100110 0000001 1010111 1010100 1001110 1100010 0001111 1111100'
    """
    shift_val = 0
    rot_one = [0, 1, 8, 15]
    if round_index in rot_one:
        shift_val = 1
    if shift_val == 0:
        shift_val = 2

    c = shift_function(c, shift_val)
    d = shift_function(d, shift_val)
    pc2 = permutate(c + d, PC2)
    if debug:
        print "pc2              : ", pprint(pc2, 6), 'bits', len(pc2)
    return [pc2, c, d]

def subkeys(key, debug=True):
    # permuted choice 1
    pc1 = permutate(key, PC1)
    if debug:
        print "pc1              : ", pprint(pc1, 7), 'bits', len(pc1)
    c, d = split(pc1, 28)
    keys = []
    for i in range(16):
        pc2, c, d = key_schedule(c, d, i, c_left_shift, debug=debug)
        keys.append(pc2)
    return keys

def des(input, key, decrypt=False, debug=True):
    """Encrypts/decrypts an input integer using DES

    >>> plaintext = 0x0123456789abcdef
    >>> key = 0x3b3898371520f75e
    >>> ciphertext = bin_array_to_int(des(plaintext, key, debug=False))
    >>> hex(ciphertext)
    '0xaa39b9777efc3c14L'
    >>> decrypted = des(ciphertext, key, decrypt=True, debug=False)
    >>> hex(bin_array_to_int(decrypted))
    '0x123456789abcdefL'
    """
    # convert input and key to bit arrays
    p = pad(int_to_bin_array(input), 64)
    k = pad(int_to_bin_array(key), 64)
    if debug:
        print "padded p         : ", pprint(p)
        print "padded k         : ", pprint(k)

    # generate round keys
    round_keys = subkeys(k, debug)
    if decrypt:
        round_keys.reverse()

    # initial permutation
    ip = permutate(list(p), IP)
    if debug:
        print "ip               : ", hex(bin_array_to_int(ip))

    # permuted choice 1
    pc1 = permutate(k, PC1)
    if debug:
        print "pc1              : ", pprint(pc1, 7), 'bits', len(pc1)

    c, d = split(pc1, 28)
    l, r = split(ip, 32)

    for i in range(16):
        if debug:
            print "\nround ", i
            print "l                : ", pprint(l), 'bits', len(l)
            print "r                : ", pprint(r), 'bits', len(r)
            print "cd               : ", pprint(c+d, 7), 'bits', len(c+d)
        f_result = f(r, round_keys[i], debug=debug)
        xor_fresult_and_l = pad(xor(f_result, l), 32)
        if debug:
            print "xor_fresult_and_l    : ", hex(bin_array_to_int(xor_fresult_and_l)), 'bits', len(xor_fresult_and_l)
        l = r
        r = xor_fresult_and_l

    # swap l and r
    l, r = r, l

    # final permutation
    inv_ip = permutate(l + r, FP)
    if debug:
        print "inv ip           : ", hex(bin_array_to_int(inv_ip))
    return inv_ip

if __name__ == '__main__':
    import doctest
    doctest.testmod()