#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <memory.h>
#include <string.h>
#include <inttypes.h>

#define ROUNDS 32

//128-bit integers are necessary for full-compliance, but there is not a single
//use-case I know of that needs to support encrypting more than 2^64-1 bytes
#ifdef __GNUC__
typedef unsigned __int128 u128;
#else
typedef uint64_t u128;
#endif

typedef uint8_t byte;

typedef struct
{
    uint64_t elements[2];
} vector;

typedef struct
{
    byte elements[32];
} block;

enum direction
{
    ENCRYPT,
    DECRYPT
};

typedef struct state
{
    vector key_schedule[ROUNDS+6][2];
    vector nonce;
    block mac;

    u128 counter;
    u128 length;

    block (*crypt)(block x, struct state *s);
    enum direction direction;
    char buffer[32];
    uint8_t buffer_len;
} state;

const block null_block = {0};

static const uint64_t pi[8] = {
    0x243F6A8885A308D3, 0x13198A2E03707345, 0xA4093822299F31D1, 0x082EFA98EC4E6C89,
    0x452821E638D01377, 0xBE5466CF34E90C6D, 0xC0AC29B7C97C50DD, 0x3F84D5B5B5470917
};

static vector u128_to_vector(u128 x)
{
    vector t;
    t.elements[0] = 0xFFFFFFFFFFFFFFFF & x;
    t.elements[1] = 0xFFFFFFFFFFFFFFFF & x >> 64;
    return t;
}

static uint64_t rotl(uint64_t x, uint8_t r)
{
    return (x << r) | (x >> (64-r));
}

static vector vector_xor(vector x, vector y)
{
    x.elements[0] ^= y.elements[0];
    x.elements[1] ^= y.elements[1];

    return x;
}

vector F(vector x, vector z)
{
    z.elements[0] += x.elements[0];
    z.elements[1] ^= x.elements[1];
    x.elements[0] += z.elements[1];
    x.elements[1] ^= z.elements[0];

    x.elements[0] = rotl(x.elements[0], 2);
    x.elements[1] = rotl(x.elements[1],17);
    z.elements[0] = rotl(z.elements[0],31);
    z.elements[1] = rotl(z.elements[1],47);

    x.elements[0] = pi[0]*x.elements[0] + pi[1]*x.elements[1]
        + pi[2]*z.elements[0] + pi[3]*z.elements[1];

    x.elements[1] = pi[4]*x.elements[0] + pi[5]*x.elements[1]
        + pi[6]*z.elements[0] + pi[7]*z.elements[1];

    return x;
}

void block_to_vectors(block b, vector x[2])
{
    int i, j, k;
    for (i=0; i<2; i++)
    {
        for (j=0; j<2; j++)
        {
            x[i].elements[j] = 0;
            for (k=0; k<8; k++)
            {
                x[i].elements[j] |= (uint64_t)b.elements[16*i+8*j+k] << 8*k;
            }
        }
    }
}

block vectors_to_block(vector x[2])
{
    block b;
    int i, j, k;
    for (i=0; i<2; i++)
    {
        for (j=0; j<2; j++)
        {
            for (k=0; k<8; k++)
            {
                b.elements[16*i+8*j+k] = (x[i].elements[j] >> 8*k)&0xFF;
            }
        }
    }
    return b;
}
block block_xor(block x, block y)
{
    int i;
    for (i=0; i<32; i++)
    {
        x.elements[i] ^= y.elements[i];
    }
    return x;
}

static block encrypt(block b, state *s)
{
    vector tweak[2];
    vector x[2];
    int i;

    block_to_vectors(b, x);

    //Compute tweak values from the nonce and counter
    //This is an invertible operation to guarantee no two combinations of nonce/counter
    //will result in identical tweak values

    tweak[0] = u128_to_vector(s->counter);
    tweak[1] = s->nonce;
    for (i=0; i<4; i++)
    {
        tweak[i%2] = vector_xor(tweak[i%2],
                F(
                    vector_xor(tweak[(i+1)%2], s->key_schedule[i][0]),
                    s->key_schedule[i][1]));
    }

    //Input whitening
    x[0] = vector_xor(x[0], s->key_schedule[4][0]);
    x[1] = vector_xor(x[1], s->key_schedule[4][1]);

    //Perform the actual encryption rounds
    for (i=5; i<ROUNDS+5; i++)
    {
        x[(i+1)%2] = vector_xor(x[(i+1)%2],
                F(vector_xor(x[i%2], s->key_schedule[i][0]),
                    vector_xor(tweak[i%2], s->key_schedule[i][1])));

        if (i-5 == ROUNDS/2)
        {
            s->mac = block_xor(s->mac, vectors_to_block(x));
        }
    }

    //Output Whitening
    x[0] = vector_xor(x[0], s->key_schedule[ROUNDS+5][0]);
    x[1] = vector_xor(x[1], s->key_schedule[ROUNDS+5][1]);

    return vectors_to_block(x);
}

static block decrypt(block b, state *s)
{
    vector tweak[2];
    vector x[2];
    int i;

    block_to_vectors(b, x);

    //Compute tweak values from the nonce and counter
    //This is an invertible operation to guarantee no two combinations of nonce/counter
    //will result in identical tweak values

    tweak[0] = u128_to_vector(s->counter);
    tweak[1] = s->nonce;
    for (i=0; i<4; i++)
    {
        tweak[i%2] = vector_xor(tweak[i%2],
                F(
                    vector_xor(tweak[(i+1)%2], s->key_schedule[i][0]),
                    s->key_schedule[i][1]));
    }

    //Remove Output Whitening
    x[0] = vector_xor(x[0], s->key_schedule[ROUNDS+5][0]);
    x[1] = vector_xor(x[1], s->key_schedule[ROUNDS+5][1]);

    //Perform the actual encryption rounds
    for (i=ROUNDS+4; i>= 5; i--)
    {
        if (i-5 == ROUNDS/2)
        {
            s->mac = block_xor(s->mac, vectors_to_block(x));
        }

        x[(i+1)%2] = vector_xor(x[(i+1)%2],
                F(vector_xor(x[i%2], s->key_schedule[i][0]),
                    vector_xor(tweak[i%2], s->key_schedule[i][1])));
    }

    //Remove Input whitening
    x[0] = vector_xor(x[0], s->key_schedule[4][0]);
    x[1] = vector_xor(x[1], s->key_schedule[4][1]);

    return vectors_to_block(x);
}

/* API Functions */

void reinit(state *s, vector nonce, enum direction direction)
{
    s->nonce = nonce;
    s->mac = null_block;
    s->counter = 0;
    s->length = 0;
    s->buffer_len = 0;
    s->direction = direction;
    if (direction == ENCRYPT)
    {
        s->crypt = &encrypt;
    }
    else
    {
        s->crypt = &decrypt;
    }
}

state *init(vector *key, size_t key_length, vector nonce, enum direction direction)
{
    unsigned int i, j, k;
    vector x[2] = {0};
    u128 counter = 0;
    state *s = calloc(1, sizeof(state));

    static const vector e = {.elements = {0xB7E151628AED2A6A,0xBF7158809CF4F3C7}};

    for(i=0; i<key_length; i++)
    {
        x[0] = vector_xor(x[0], F(x[1], key[i]));
        x[1] = vector_xor(x[1], F(x[0], e));
    }
    for (i=0; i<ROUNDS+6; i++)
    {
        for (j=0; j<2; j++)
        {
            x[0] = vector_xor(x[0], u128_to_vector(counter++));
            for(k=0; k<key_length; k++)
            {
                x[0] = vector_xor(x[0], F(x[1], key[k]));
                x[1] = vector_xor(x[1], F(x[0], e));
            }
            s->key_schedule[i][j] = x[0];
        }
    }
    reinit(s, nonce, direction);

    return s;
}

void dest(state **s)
{
    //You should really zero memory here, but memset will be optimized, GCC does not
    //support memset_s, and I'm lazy
    free(*s);
    *s = NULL;
}

size_t append(char *in, size_t in_len, char *out, size_t out_len, state *s,
        size_t *bytes_written)
{
    size_t in_pos = 0;
    size_t out_pos = 0;
    size_t copy_length;
    block b;

    if (s->direction == DECRYPT)
    {
        if (in_len < 32) //We always need to process at least 32 bytes when decrypting
        {
            *bytes_written = 0;
            return 0;
        }
        in_len -= 32; //Exactly 32 bytes will always be placed in the buffer
    }

    if (s->buffer_len > 0 && s->buffer_len < 32)
    {
        if (in_len+s->buffer_len >= 32)
        {
            copy_length = 32-s->buffer_len;
        }
        else
        {
            copy_length = in_len;
        }

        memcpy(s->buffer+s->buffer_len, in, copy_length);
        in_pos += copy_length;
        s->buffer_len += copy_length;
    }

    if (s->buffer_len == 32)
    {
        memcpy(b.elements, s->buffer, 32);
        b = s->crypt(b, s);
        s->counter++;
        memcpy(out+out_pos, b.elements, 32);

        out_pos += 32;
        s->buffer_len = 0;
    }

    while (in_len-in_pos >= 32 && out_len-in_pos >= 32)
    {
        memcpy(b.elements, in+in_pos, 32);
        in_pos += 32;
        b = s->crypt(b, s);
        s->counter++;
        memcpy(out+out_pos, b.elements, 32);
        out_pos += 32;
    }

    if (s->direction == ENCRYPT)
    {
        copy_length = in_len-in_pos;
        if (copy_length > 31)
        {
            copy_length = 31;
        }
    }
    else
    {
        copy_length = 32;
    }

    if (copy_length > 0)
    {
        memcpy(s->buffer, in+in_pos, copy_length);
        s->buffer_len = copy_length;
        in_pos += copy_length;
    }
    s->length += out_pos;
    *bytes_written = out_pos;

    return in_pos;
}

//Result is freed on call to dest(state), you
//should not free result or access after calling dest
char * finalize(size_t *out_length, state *s)
{
    block b;
    //Something to note, I'm using a non-standard padding scheme
    //When setting the padding bytes, each byte is set to the number of plaintext bytes in
    //the final block. When determining the number of plaintext bytes when decrypting,
    //the value of the last byte mod [block_length] is used so there is no invalid padding
    if (s->direction == ENCRYPT)
    {
        //pad and encrypt
        memcpy(b.elements, s->buffer, s->buffer_len);
        memset(b.elements+s->buffer_len, s->buffer_len, 32-s->buffer_len);
        s->buffer_len = 32;
        b = s->crypt(b, s);
        memcpy(s->buffer, b.elements, s->buffer_len);
    }
    else
    {
        //decrypt and remove padding
        if (s->buffer_len == 32)
        {
            memcpy(b.elements, s->buffer, s->buffer_len);
            b = s->crypt(b, s);
            s->buffer_len = b.elements[31] % 32;
            memcpy(s->buffer, b.elements, s->buffer_len);
        }
        else
        {
            return NULL;
        }
    }

    *out_length = s->buffer_len;
    s->length += s->buffer_len;

    //Finalize MAC computation by encrypting using the length instead of the counter for the
    //tweak value
    s->counter = s->length;
    //s->mac = encrypt(s->mac, s);

    //return a reference to the buffer
    return s->buffer;
}

char * get_mac(size_t *out_length, state *s)
{
    if (s->length != s->counter) //Not finalized
    {
        *out_length = 0;
        return NULL;
    }
    else
    {
        *out_length = 32;
        return (char *)s->mac.elements;
    }
}

int check_mac(char *mac, size_t mac_length, state *s)
{
    int i;
    if (mac_length > 32)
    {
        return 0; //False
    }
    else
    {
        uint8_t check = 0;
        for (i=0; i<mac_length; i++)
        {
            //if any bit differed, check will be non-zero
            check |= s->mac.elements[i] ^ (uint8_t)mac[i];
        }
        return !check;
    }
}

//Encrypt or Decrypt all bytes except for the last [trailing_bytes_len] bytes
size_t process_io(FILE *input, FILE *output,
        char *in_buffer, size_t in_buffer_len, size_t trailing_bytes_len,
        char *out_buffer, size_t out_buffer_len,
        state *s)
{
    size_t buffer_len, bytes_read, bytes_written;
    size_t buffer_pos = 0;
    char *tail;

    if (trailing_bytes_len > 0)
    {
        buffer_pos = fread(in_buffer, 1, trailing_bytes_len, stdin);
        if (buffer_pos != trailing_bytes_len)
        {
            fprintf(stderr, "Not enough trailing bytes\n");
            exit(1);
        }
    }
    while (trailing_bytes_len < (buffer_len = buffer_pos + fread(in_buffer+buffer_pos, 1,
                    in_buffer_len-buffer_pos, stdin)))
    {

        bytes_read = append(in_buffer, buffer_len-trailing_bytes_len, out_buffer,
                out_buffer_len, s, &bytes_written);

        if (bytes_written != fwrite(out_buffer, 1, bytes_written, stdout))
        {
            fprintf(stderr, "Error writing to output\n");
            exit(1);
        }

        memmove(in_buffer, in_buffer+bytes_read, buffer_len-bytes_read);
        buffer_pos = buffer_len-bytes_read;
    }

    tail = finalize(&bytes_written, s);

    if (bytes_written != fwrite(tail, 1, bytes_written, stdout))
    {
        fprintf(stderr, "Error writing to output\n");
        exit(1);
    }
    return buffer_pos;
}

int main(int argc, char **argv)
{
    if (argc == 2)
    {
        char direction = argv[1][0];
        vector key = {0};
        vector nonce;
        state *s;
        char *mac;
        size_t bytes_read, bytes_written;
        char out_buffer[8192];

        block x = {0};

        if (direction == 'e')
        {
            char in_buffer[8192];
            FILE *random = fopen("/dev/random", "rb");

            if (random == NULL)
            {
                fprintf(stderr, "Error reading from /dev/random\n");
                exit(1);
            }

            bytes_read = fread(nonce.elements, 1, sizeof(nonce.elements), random);

            fclose(random);
            random = NULL;
            if (bytes_read != 16)
            {
                fprintf(stderr, "Error reading from /dev/random\n");
                exit(1);
            }

            if (16 != fwrite(nonce.elements, 1, 16, stdout))
            {
                fprintf(stderr, "Error writing to output\n");
                exit(1);
            }

            s = init(&key, 1, nonce, ENCRYPT);

            process_io(stdin, stdout,
                    in_buffer, sizeof(in_buffer), 0, out_buffer, sizeof(out_buffer), s);

            mac = get_mac(&bytes_written, s);

            if (bytes_written != fwrite(mac, 1, bytes_written, stdout))
            {
                fprintf(stderr, "Error writing to output\n");
                exit(1);
            }
            dest(&s);
        }
        else
        {
            char in_buffer[8192+32];
            bytes_read = fread(nonce.elements, 1, sizeof(nonce.elements), stdin);

            if (bytes_read != 16)
            {
                fprintf(stderr, "File did not start with nonce\n");
                exit(1);
            }
            s = init(&key, 1, nonce, DECRYPT);

            bytes_read = process_io(stdin, stdout,
                    in_buffer, sizeof(in_buffer), 32, out_buffer, sizeof(out_buffer), s);

            if (!check_mac(in_buffer, bytes_read, s))
            {
                fprintf(stderr, "MAC check failed\n");
            }
            dest(&s);
        }
    }
    return 0;
}