DES algorithm that works with bit arrays and uses PKCS7 padding

DES.py

import bitarray
import itertools
from collections import deque


class DES(object):

    _initial_permutation = [
        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 = [
        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_function = [
        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 = [
        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
    ]

    _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
    ]

    _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
    ]

    _left_rotations = [
        1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1
    ]

    _sbox = [
        # 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]
        ]
    ]

    def __init__(self, key):
        self.key = key

    def encrypt(self, message):
        padded = self.pkcs7_padding(message, pad=True)
        result = []

        for block in padded:
            result += self.encrypt_64bit(''.join(str(i) for i in block))

        return result

    def decrypt(self, message, msg_in_bits=False):
        bits_array_msg = []
        if msg_in_bits:
            bits_array_msg = message
        else:
            bits_array_msg = self._string_to_bitsarray(message)

        if len(bits_array_msg) % 64 != 0:
            raise ValueError('Ciphered code must be a multiple of 64')

        blocks_lst = [
            bits_array_msg[i:i + 64] for i in range(0, len(bits_array_msg), 64)
        ]

        result = []

        for block in blocks_lst:
            decrypted = self.decrypt_64bit(block, msg_in_bits=True)
            bl = list(
                ''.join(chr(int(
                    ''.join(
                        map(str, decrypted[i:i + 8])),
                    2)) for i in range(0, len(decrypted), 8)))
            bl = self._unpad(bl)
            result += bl

        return ''.join(result)

    def encrypt_64bit(self, message):
        return self.crypt(message, encrypt=True)

    def decrypt_64bit(self, message, msg_in_bits=False):
        return self.crypt(message, encrypt=False, msg_in_bits=msg_in_bits)

    def crypt(self, message, encrypt=True, msg_in_bits=False):
        bits_array_msg = []
        if msg_in_bits:
            bits_array_msg = message
        else:
            bits_array_msg = self._string_to_bitsarray(message)

        bits_array_key = self._string_to_bitsarray(self.key)

        if len(bits_array_msg) != 64:
            raise ValueError('Message must be 64 bit!')

        if len(bits_array_key) != 64:
            raise ValueError('Key must be 64 bit!')

        # Compute 16 48-bit subkeys
        subkeys = self.get_subkeys(bits_array_key)

        # Convert the message using the initial permutation block
        msg = [bits_array_msg[i - 1] for i in self._initial_permutation]

        L, R = msg[:32], msg[32:]

        if encrypt:
            for i in range(16):
                prev_r = R
                r_feistel = self.feistel_function(R, subkeys[i])
                R = [L[i] ^ r_feistel[i] for i in range(32)]
                L = prev_r
        else:
            for i in reversed(range(16)):
                prev_l = L
                l_feistel = self.feistel_function(L, subkeys[i])
                L = [R[i] ^ l_feistel[i] for i in range(32)]
                R = prev_l

        before_final_permute = L + R
        return [before_final_permute[i - 1] for i in self._final_permutation]

    def pkcs7_padding(self, message, block_size=8, pad=True):
        msg = list(message)
        blocks_lst = [
            msg[i:i + block_size] for i in range(0, len(msg), block_size)
        ]

        s = block_size
        return [
            self._pad(b, s) if len(b) < block_size else b for b in blocks_lst
        ] if pad else blocks_lst

    def feistel_function(self, r_32bit, subkey_48bit):
        r_48bit = [r_32bit[i - 1] for i in self._expansion_function]
        subkey_xor_r = [r_48bit[i] ^ subkey_48bit[i] for i in range(48)]

        # Divide subkey_xor_r into 8 6-bit blocks for computing s-boxes
        b_6_bit_blocks = [subkey_xor_r[i:i + 6] for i in range(0, 48, 6)]

        # Compute 8 s-boxes and concatenate them into 32-bit vector
        after_sboxes_32bit = list(itertools.chain(*[
            self.compute_s_box(
                self._sbox[i], b_6_bit_blocks[i]) for i in range(8)
        ]))
        # Compute the permutation and return the 32-bit block
        return [int(after_sboxes_32bit[i - 1]) for i in self._permutation]

    def compute_s_box(self, sbox, b_6_bit):
        row = int(''.join(str(x) for x in [b_6_bit[0], b_6_bit[5]]), 2)
        col = int(''.join(str(x) for x in b_6_bit[1:5]), 2)
        s_box_res = sbox[row][col]
        return list('{0:04b}'.format(s_box_res))

    def get_subkeys(self, bits_array_key):
        # Extract 8 parity bits from the key (8, 16, 24, 32, 40, 48, 56, 64)
        # key_56bit = bits_array_key
        # del key_56bit[7::8]

        # Compute Permuted Choice 1 on the key
        key_56bit = [bits_array_key[i - 1] for i in self._pc1]

        # Split the key into two 28-bit subkeys
        key_56_left, key_56_right = [
            key_56bit[i:i + 28] for i in range(0, 56, 28)
        ]

        # Compute 16 48-bit keys using left rotations and permuted choice 2
        subkeys_48bit = []
        C, D = key_56_left, key_56_right
        for i in range(16):
            C_deque, D_deque = deque(C), deque(D)
            C_deque.rotate(-self._left_rotations[i])
            D_deque.rotate(-self._left_rotations[i])

            C, D = list(C_deque), list(D_deque)
            CD = C + D
            subkeys_48bit.append([CD[i - 1] for i in self._pc2])

        return subkeys_48bit

    def _string_to_bitsarray(self, string):
        ba = bitarray.bitarray()
        ba.fromstring(string)
        return [1 if i else 0 for i in ba.tolist()]

    def _pad(self, arr, block_size):
        z = block_size - len(arr)
        return arr + [z] * z

    def _unpad(self, arr):
        if str(arr[-1]).isdigit():
            arr_str = ''.join(str(i) for i in arr)
            i = j = int(arr[-1])

            while arr_str[-1] == str(j) and i > 0:
                arr_str = arr_str[:-1]
                i -= 1

            return list(arr_str)
        else:
            return arr



# ==============================================================

d = DES('qwertyui')

cipher = d.encrypt('hello world!')

print(cipher)

deciphered = d.decrypt(cipher, msg_in_bits=True)

print(deciphered)