#include "spn.h"

unsigned char do_subst(unsigned char w, unsigned char subst[16]) {
    return subst[w];
}

unsigned char do_perm(unsigned char w, unsigned char perm[16]) {
    unsigned char result = 0;
    for (int i = 0; i < 8; i++)
        if ((w >> i) & 1)
            result |= (1 << perm[i]);
    return result;
}

unsigned short encrypt(unsigned short w, unsigned int   key, unsigned char  perm[16], unsigned char  subst[16]) {
    for (int i = 0; i < 2; i++) {
        unsigned char rk = (key >> (8 * (3 - i))) & 0xFF;
        w ^= (rk << 8) | rk;

        unsigned short tmp = 0;
        for (int j = 0; j < 4; ++j) {
            unsigned char nib = (w >> (j * 4)) & 0xF;
            tmp |= (subst[nib] << (j * 4));
        }
        w = tmp;

        unsigned short permute = 0;
        for (int j = 0; j < 16; ++j)
            if ((w >> j) & 1) {
                permute |= (1 << perm[j]);
            }
        w = permute;
    }
    return w;     
}

unsigned short decrypt(unsigned short w, unsigned int   key, unsigned char  perm[16], unsigned char  subst[16]) {
    unsigned char inverseSubst[16];
    unsigned char inversePerm[16];
    for (int i = 0; i < 16; ++i) {
        inverseSubst[subst[i]] = i;
        inversePerm[perm [i]] = i;
    }

    for (int i = 1; i >= 0; --i) {
        unsigned short permInverse = 0;
        for (int j = 0; j < 16; ++j)
            if ((w >> j) & 1) {
                permInverse |= (1 << inversePerm[j]);
            }
        w = permInverse;

        unsigned short tmp = 0;
        for (int j = 0; j < 4; ++j) {
            unsigned char sousBloc = (w >> (j * 4)) & 0xF;
            tmp |= (inverseSubst[sousBloc] << (j * 4));
        }
        w = tmp;

        unsigned char rk = (key >> (8 * (3 - i))) & 0xFF;
        w ^= (rk << 8) | rk;
    }
    return w; 
}