saes.py

# pip3 install pyfinite
from pyfinite import ffield

class SAES:
    def __init__(self, key):
        # generator za F_16 = F_2[x]/x^4+x+1
        X_generator = 0b10011
        
        # generator za F_2[y]/y^4+1
        Y_generator = 0b10001
        
        # generator za F_16[z]/z^2+1
        Z_generator = 0b101

        # polinom a(y) = y^3+y^2+1 za funkciju S
        self.a = 0b1101

        # polinom b(y) = y^3+1 za funkciju S
        self.b = 0b1001
        
        # Polja sa odgovarajucim generatorima
        self.X_field = ffield.FField(4, gen=X_generator, useLUT=0)
        self.Y_field = ffield.FField(4, gen=Y_generator, useLUT=0)
        self.Z_field = ffield.FField(4, gen=Z_generator, useLUT=0)
        
        # inicijalizacija S tabele
        self._init_S()
        
        # prosirivanje kljuca
        self._extend_key(key)

    def _init_S(self):
        self.S_box = {}
        self.S_box_inv = {}
        
        for i in range(16):
            # Prvo odredjujemo inverz broja i u polju X
            if i == 0:
                N = 0b0
            else:
                N = self.X_field.Inverse(i)

            # zapisujemo inverz kao element polja Y
            Ny = self.Y_field.Multiply(N, 1)
            
            # vrsimo transformaciju Ny u polju Y: Ny*a(y) + b(y)
            s = self.Y_field.Add(self.Y_field.Multiply(self.a, Ny), self.b)
            
            # upisujemo vrednost u tabelu
            self.S_box[i] = s
            self.S_box_inv[s] = i
    
    def S(self, x):
        return self.S_box[x]

    def S_inv(self, x):
        return self.S_box_inv[x]

    # funkcija SubNib
    def _sub_nib(self, nibble_pair):
        n1 = (0b11110000 & nibble_pair) >> 4
        n2 = (0b1111 & nibble_pair)
        
        n1_sub = self.S(n1)
        n2_sub = self.S(n2)
        
        return n1_sub * 16 + n2_sub
    
    # funkcija RotNib
    def _rot_nib(self, nibble_pair):  
        n1 = (0b11110000 & nibble_pair) >> 4
        n2 = (0b1111 & nibble_pair)
        
        return n2 * 16 + n1

    # pretvaranje bitova u matricu koja opisuje jedno stanje
    def _bytes_to_matrix(self, byte_array):
        return [
            [(byte_array[0] & 0b11110000) >> 4, (byte_array[1] & 0b11110000) >> 4],
            [(byte_array[0] & 0b1111), (byte_array[1] & 0b1111)]
        ]

    # funkcija za prosirivanje kljuca
    def _extend_key(self, key):
        if len(key) != 2:
            raise Exception('Invalid key length')
        
        W = [ord(key[0]), ord(key[1]), 0, 0, 0, 0]
        
        RC = []
        x2 = self.X_field.ConvertListToElement([0,1,0,0])

        for i in range(1,4):
            x_i = [0,0,0,0]
            x_i[-i-1] = 1
            
            xi = self.X_field.ConvertListToElement(x_i)
            RC.append(self.X_field.Multiply(xi, x2))

        RCON = [0] + [rc * 16 for rc in RC]

        for i in range(2,6):
            if i % 2 == 0:
                k = self._sub_nib(self._rot_nib(W[i-1]))
                W[i] = RCON[i//2] ^ k ^ W[i-2]
            else:
                W[i] = W[i-2] ^ W[i-1]

        self.extended_key = W

    # Ak funkcija, dodavanje kljuca bit po bit (involucija)
    def _add_key(self, i, state):
        k_i = self.extended_key[i * 2 : i * 2 + 2]
        key_mat = self._bytes_to_matrix(k_i)
        
        return [
            [key_mat[0][0] ^ state[0][0], key_mat[0][1] ^ state[0][1]],
            [key_mat[1][0] ^ state[1][0], key_mat[1][1] ^ state[1][1]]
        ]

    # NS funkcija, primena funkcije S na svaki od niblova trenutnog stanja
    def _nibble_substitution(self, state):
        return [
            [self.S(state[0][0]), self.S(state[0][1])],
            [self.S(state[1][0]), self.S(state[1][1])]
        ]

    # NS inverzna funkcija, primena funkcije S^-1 na svaki od niblova trenutnog stanja
    def _nibble_substitution_inv(self, state):
        return [
            [self.S_inv(state[0][0]), self.S_inv(state[0][1])],
            [self.S_inv(state[1][0]), self.S_inv(state[1][1])]
        ]

    # SR funkcija, pomeranje drugog reda za jedno mesto (involucija)
    def _shift_row(self, state):
        return [
            [state[0][0], state[0][1]],
            [state[1][1], state[1][0]]
        ]

    # MC funkcija, primena formule (N_i*z + N_j)*c(z) = z(N_i + N_j*x^2) + (N_i*x^2 + N_j)
    # na obe kolone trenutnog stanja
    def _mix_column(self, state):
        Ni1 = state[0][0]
        Ni2 = state[0][1]
        Nj1 = state[1][0]
        Nj2 = state[1][1]
        return [
            [
            self.X_field.Add(Ni1, self.X_field.Multiply(Nj1, 0b100)),
            self.X_field.Add(Ni2, self.X_field.Multiply(Nj2, 0b100))
            ],
            [
            self.X_field.Add(Nj1, self.X_field.Multiply(Ni1, 0b100)),
            self.X_field.Add(Nj2, self.X_field.Multiply(Ni2, 0b100))
            ]
        ]

    # MC inverzna funkcija, primena formule (N_i*z + N_j)*c^{-1}(z) = z(N_i*x^3 + N_i + N_j*x) + (N_i*x + N_jx^3 + N_j)
    # na obe kolone trenutnog stanja
    def _mix_column_inv(self, state):
        Ni1 = state[0][0]
        Ni2 = state[0][1]
        Nj1 = state[1][0]
        Nj2 = state[1][1]
        return [
            [
            self.X_field.Add(self.X_field.Multiply(Ni1, 0b1001), self.X_field.Multiply(Nj1, 0b10)),
            self.X_field.Add(self.X_field.Multiply(Ni2, 0b1001), self.X_field.Multiply(Nj2, 0b10))
            ],
            [
            self.X_field.Add(self.X_field.Multiply(Ni1, 0b10), self.X_field.Multiply(Nj1, 0b1001)),
            self.X_field.Add(self.X_field.Multiply(Ni2, 0b10), self.X_field.Multiply(Nj2, 0b1001))
            ]
        ]
    
    # sifrovanje 16bitnog teksta
    def _encrypt_bytes(self, data):
        state = self._bytes_to_matrix(data)
        
        state = self._add_key(0, state)
        state = self._nibble_substitution(state)
        state = self._shift_row(state)
        state = self._mix_column(state)
        state = self._add_key(1, state)
        state = self._nibble_substitution(state)
        state = self._shift_row(state)
        state = self._add_key(2, state)
        
        return [state[0][0] * 16 + state[1][0], state[0][1] * 16 + state[1][1]]
    
    # desifrovanje 16bitnog sifrata
    def _decrypt_bytes(self, data):
        state = self._bytes_to_matrix(data)
        
        state = self._add_key(2, state)
        state = self._shift_row(state)
        state = self._nibble_substitution_inv(state)
        state = self._add_key(1, state)
        state = self._mix_column_inv(state)
        state = self._shift_row(state)
        state = self._nibble_substitution_inv(state)
        state = self._add_key(0, state)
        
        return [state[0][0] * 16 + state[1][0], state[0][1] * 16 + state[1][1]]

    # funkcija koja dati tekst deli na grupe od 2 karaktera i svaku grupu pojedinacno sifruje
    def encrypt(self, string_data):
        data_bytes = [ord(x) for x in string_data]
        n = len(data_bytes)
        
        if n % 2 == 1:
            data_bytes.append(ord(' '))
        
        encrypted_bytes = []

        for i in range(0, n, 2):
            data_bytes_slice = data_bytes[i:i+2]
            encrypted_bytes += self._encrypt_bytes(data_bytes_slice)

        return ''.join([chr(x) for x in encrypted_bytes])

    # funkcija koja dati sifrat deli na grupe od 2 karaktera i svaku grupu pojedinacno desifruje
    def decrypt(self, encoded_data):
        data_bytes = [ord(x) for x in encoded_data]
        n = len(data_bytes)
        decrypted_bytes = []
        
        for i in range(0, n, 2):
            data_bytes_slice = data_bytes[i:i+2]
            decrypted_bytes += self._decrypt_bytes(data_bytes_slice)
        
        return ''.join([chr(x) for x in decrypted_bytes])

def main():
    message="0111110110010100"
    key="0101000101100101"
    
    A=SAES(key)
    B=SAES(key)

    encrypted_data=A.encrypt(str(message))
    decrypted_data=B.decrypt(str(encrypted_data))
    print("A sends this: ", message)
    print("A encrypts message: ", encrypted_data)
    print("B decrypts message: ", decrypted_data)
    
if __name__=="__main__":
   main()