#include <stdio.h>

#define assert(x) if(!(x)) fprintf(stderr,"Error in line %d file %s\n", __LINE__,__FILE__)

#define eval_box(b,x) b[(x)>>3][(x) & 0x7] 


char *binString(int n, unsigned int mask, char *buff){
    int i;
    int top;
    for (top=31;!((1<<top) &mask);top--);
    top++;
    for (i=0;i<top;i++){
        buff[top-i-1] =(1<<i)&n?'1':'0';    
    }
    buff[top]=0;
    return buff;

}


unsigned char s1[2][8] = {{5,2,1,6,1,4,7,0},
         {1,4,6,2,0,7,5,3}};

unsigned char s2[2][8] = {{4,0,6,5,7,1,3,2},
        {5,3,0,7,6,2,1,4 }};



unsigned char expand(unsigned char c){
    assert(c<64);
    return ((c & 0x30) << 2)| ((c & 0x4) << 3) | ((c & 0xc)<<1)|((0x9 & c) >> 1 ) | (c & 0x3);
}


unsigned char keyRound(unsigned int K, int round){
    unsigned int mask = (1<<(9-round+1))-1;
    unsigned char ret;
    if (round > 1)
        ret =  ((K & mask) << (round-2) | (K & (~mask))>>(9-round+2)); 
    else
        ret = K>>1;
    //printf("Key %X round %d\n",ret,round);
    return ret;
//    return  (K & mask) << (round-2); //| (K & (~mask))>>(9-round+1)); 
}
unsigned char f(unsigned char R, unsigned char K){
    unsigned char iv;
    iv = expand(R) ^ K;
    return (eval_box(s1,iv>>4) << 3) | eval_box(s2,iv & 0xf);
}

unsigned int encript(unsigned int pt, unsigned int K, int rounds){
    int i;
    unsigned char L,R,tmp;
    L = 0x3f & (pt >> 6);
    R = pt & 0x3f;
    for (i=1;i<=rounds;i++){
//        printf("1:L=%x R=%x key = %x\n",L,R,keyRound(K,i));
        tmp = (L ^ f(R,keyRound(K,i)));
        L = R;
        R=tmp;
//        printf("2:L=%x R=%x\n",L,R);
    }
    return (L<<6) | R;
}

unsigned int decript(unsigned int ct, unsigned int K, int rounds){
    int i;
    unsigned char L,R,tmp;
    // Right is L and Left is Right
    R = 0x3f & (ct >> 6);
    L = ct & 0x3f;
    for (i=rounds;i>=1;i--){
//        printf("1:L=%x R=%x\n",L,R);
        tmp = (L ^ f(R,keyRound(K,i)));
        L = R;
        R=tmp;
//        printf("2:L=%x R=%x\n",L,R);
    }
    return (R<<6) | L;
}
int readInt(){
    char buff[255];
    int i=0;
    char c;
    while(((c = getchar()) != '\n') && i++<255){
        buff[i-1]=c;
    }
    buff[i]=0;
    return atoi(buff);
}

unsigned int readBin(int bits){
    char c;
    int i;
    unsigned int ret = 0;
    for (i=bits-1;i>=0;i--){
        c = getchar();
        ret  = ret | (c == '1'?1<<i:0);
        
    }
    while (getchar()!='\n');
    return ret;
}

int main(int argc, char **argv){
    int weak_flag = 0;
    unsigned int i;
    unsigned int key = 0x99;
    unsigned int pt = 0xbee;
    unsigned int ct = 0x00;
    char buff[32];
    int rounds;
    printf("Key?");
    key = readBin(9);
    printf("Plaintext?");
    pt = readBin(12);
    printf("Key     %s\n",binString(key,0x1ff,buff));
    // part b
    for (i=1;i<=4;i++){
        printf("\nRounds %d\n",i);
        printf("Plaint text %X = %s\n",pt,binString(pt ,0xfff,buff));
        ct = encript(pt, key,i);
        printf("Cipher text %X = %s\n",ct,binString(ct ,0xfff,buff));
        pt = decript(ct, key,i);
        printf("Plaint text %X = %s\n",pt,binString(pt ,0xfff,buff));
    }

    // part c
    weak_flag = 0;
    for (i=0;i<(1<<9);i++){
        int weak = 0;
        for (pt=0;pt<(1<<12);pt++){
            if (pt != encript( encript(pt, i ,4), i ,4))
                break;    
            
            if (pt == ((1<<12)-1))
                weak++;
        }    
    } 
    if (weak_flag)
            printf("Found a weak key");
    else
            printf("No weak key found");
}