#include <stdio.h>
#include <stdint.h>
#include "speck.h"


// This function is only used for the "x86" Speck compilation and as reference
void FuncER16(u16 *x, u16 *y, u16 k)
{
    u16 tmp_x = *x;
    u16 tmp_y = *y;

    *x = (((tmp_x)>>(7)) | ((tmp_x)<<(16-(7))));
    *x += *y;

    *x = *x ^ k;

    *y = (((tmp_y)<<(2)) | (tmp_y>>(16-(2))));
    *y = *y ^ *x;

}


#ifdef ARM
// This function is used when running on the CW
void FuncER16_ASM(u16 *x, u16 *y, u16 k)
{

    asm volatile (
        "nop\n\t"
        "push	{r4, r5, lr}\n\t"
        "ldrh	r5, [r0, #0]\n\t"
        "ldrh	r4, [r1, #0]\n\t"
        "lsls	r3, r5, #9\n\t"
        "orr.w	r3, r3, r5, lsr #7\n\t"
        "uxth	r3, r3\n\t"
        "strh	r3, [r0, #0]\n\t"
        "ldrh	r5, [r1, #0]\n\t"
        "add	r3, r5\n\t"
        "eors	r2, r3\n\t"
        "lsls	r3, r4, #2\n\t"
        "orr.w	r3, r3, r4, lsr #14\n\t"
        "uxth	r3, r3\n\t"
        "strh	r2, [r0, #0]\n\t"
        "strh	r3, [r1, #0]\n\t"
        "ldrh	r2, [r0, #0]\n\t"
        "eors	r3, r2\n\t"
        "strh	r3, [r1, #0]\n\t"
        "pop	{r4, r5, pc}\n\t"
    );

}
#endif


void Words16ToBytes(u16 words[],u8 bytes[],int numwords)
{
    int i,j=0;
    for(i=0;i<numwords;i++){
        bytes[j]=(u8)words[i];
        bytes[j+1]=(u8)(words[i]>>8);
        j+=2;
    }
}

void BytesToWords16(u8 bytes[],u16 words[],int numbytes)
{
    int i,j=0; for(i=0;i<numbytes/2;i++){
        words[i]=(u16)bytes[j] | ((u16)bytes[j+1]<<8);
        j+=2;
    }
}

void Speck3264KeySchedule(u16 K[],u16 rk[])
{
    u16 i,D=K[3],C=K[2],B=K[1],A=K[0];
#ifdef ARM
    for(i=0;i<22;){
        rk[i]=A;
        ER16(B,A,i++);
        rk[i]=A;
        ER16(C,A,i++);
        rk[i]=A;
        ER16(D,A,i++);
    }
#endif

#ifndef ARM
    for(i=0;i<22;){

        printf("A = 0x%x ; B = 0x%x ; C = 0x%x ; D = 0x%x\n", A, B, C, D);

        rk[i]=A;
        ER16(B,A,i++);
        printf("rk[%d] =  0x%x\n", i-1, A);

        printf("A = 0x%x ; B = 0x%x ; C = 0x%x ; D = 0x%x\n", A, B, C, D);
        rk[i]=A;
        ER16(C,A,i++);
        printf("rk[%d] =  0x%x\n", i-1, A);

        printf("A = 0x%x ; B = 0x%x ; C = 0x%x ; D = 0x%x\n", A, B, C, D);
        rk[i]=A;
        ER16(D,A,i++);
        printf("rk[%d] =  0x%x\n  <- D = 0x%x", i-1, A, D);
        printf("----------------------\n");
    }
#endif
}


void Speck3264Encrypt(u16 Pt[],u16 Ct[],u16 rk[])
{
    u16 i;
    Ct[0]=Pt[0]; Ct[1]=Pt[1];

    // full 22  rounds
    for(i=0;i<22;) {
        //ER16(Ct[1],Ct[0],rk[i++]);
#ifdef ARM
        FuncER16_ASM(&Ct[1], &Ct[0],rk[i++]);
        //FuncER16(&Ct[1], &Ct[0], rk[i++]);
#else
        ER16(Ct[1],Ct[0],rk[i++]);
#endif

    }
}


void Speck3264Decrypt(u16 Pt[],u16 Ct[],u16 rk[])
{
    int i;
    Pt[0]=Ct[0]; Pt[1]=Ct[1];

    for(i=21;i>=0;) DR16(Pt[1],Pt[0],rk[i--]);
}


void Speck3264_EncryptBlock(u8 pt[], u8 k[], u8 ct[]) {

    u16 Pt[2] = {0};
    u16 K[4] = {0};
    u16 rk[34] = {0};
    u16 Ct[2] = {0};

    BytesToWords16(pt,Pt,8);
    BytesToWords16(k,K,16);


    Speck3264KeySchedule(K,rk);

#ifndef ARM
    // DEBUG Purposes
    for (int i=0; i < 16; i++)
    {
        printf("Key: 0x%x\n", rk[i]);
    }
#endif
    Speck3264Encrypt(Pt,Ct,rk);
    Words16ToBytes(Ct,ct,2);
}



#ifndef ARM
int main() {


    // test are from https://github.com/inmcm/Simon_Speck_Ciphers
    //u8 key[8] = {0x00, 0x01, 0x08, 0x09, 0x10, 0x11, 0x18, 0x19};
    u8 pt[4] = {0x4c, 0x69, 0x74, 0x65};

    u8 key[8] = {0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88};


    u8 ct[4] = {0x0};

    Speck3264_EncryptBlock(pt, key, ct);
    printf("[[ Speck 32/64  ]]\n");

    printf("The output: \n");
    for (int i = 0; i < 4; i++) {
        printf("- %08x\n", ct[i]);
    }
    return 0;
}
#endif