/*
=========================================================================================================================
|  Implementation of SIMON 32/64 in 32 Rounds. Please checkout the following link for more information on SIMON cipher. |
|  https://eprint.iacr.org/2013/404.pdf                                                                                 |
=========================================================================================================================
Copyright (c) 2015 Milad Khandouzy (m.khandouzy@smartarts.ir) and Nima Azizzadeh (n.azizzadeh@smartarts.ir)

Permission to use, copy, modify, and/or distribute this software for any
purpose with or without fee is hereby granted, provided that the above
copyright notice and this permission notice appear in all copies.

*/

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

typedef uint8_t u8;
typedef uint16_t u16;
typedef uint64_t u64;
typedef uint32_t u32;

static u8 z[62] =
{1,1,1,1,1,0,1,0,0,0,1,0,0,1,0,1,0,1,1,0,0,0,0,1,1,1,0,0,1,1,0,1,1,1,1,1,0,1,0,0,0,1,0,0,1,0,1,0,1,1,0,0,0,0,1,1,1,0,0,1,1,0};

void KeyExpansion ( u16 k[] )
{
    u8 i;
    u16 tmp;
    for ( i=4 ; i<32 ; i++ )
    {
        tmp = ROTATE_RIGHT_16(k[i-1],3);
        tmp = tmp ^ k[i-3];
        tmp = tmp ^ ROTATE_RIGHT_16(tmp,1);
        k[i] = ~k[i-4] ^ tmp ^ z[i-4] ^ 3;
    }
}

void Encrypt ( u16 text[], u16 crypt[], u16 key[] )
{
    u8 i;
    u16 tmp;
    crypt[0] = text[0];
    crypt[1] = text[1];

    for ( i=0 ; i<32 ; i++ )
    {
        tmp = crypt[0];
        crypt[0] = crypt[1] ^ ((ROTATE_LEFT_16(crypt[0],1)) & (ROTATE_LEFT_16(crypt[0],8))) ^ (ROTATE_LEFT_16(crypt[0],2)) ^ key[i];
        crypt[1] = tmp;
    }
}

void Decrypt ( u16 text[], u16 crypt[], u16 key[] )
{
    u8 i;
    u32 tmp;
    crypt[0] = text[0];
    crypt[1] = text[1];

    for ( i=0 ; i<32 ; i++ )
    {
        tmp = crypt[1];
        crypt[1] = crypt[0] ^ ((ROTATE_LEFT_16(crypt[1],1)) & (ROTATE_LEFT_16(crypt[1],8))) ^ (ROTATE_LEFT_16(crypt[1],2)) ^ key[31-i];
        crypt[0] = tmp;
    }
}

#include <fstream>
#include <sstream>
#include <string>
#include <set>
#include <algorithm> 

#include <cstring>

#include <linux/random.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

// hur durr
int random_fd = open("/dev/urandom", O_RDONLY);
int getrandom(void *buf, size_t buflen, unsigned int flags) {
    return read(random_fd, buf, buflen);
}


using blk = u32;
std::set<blk> read_corpus(const char *filename) {
    std::ifstream t(filename);
    std::stringstream buffer;
    buffer << t.rdbuf();
    auto s = buffer.str();
    s.erase(std::remove_if(s.begin(), s.end(),
            [](int c) { return c == '\r'; }
        ), s.end());

    for(size_t i=0; i<s.size(); i++)
        if(s[i] == '\n') s[i] = ' ';

    std::set<blk> seen;

    for(size_t i=0; i<s.size()-3; i++) {
        blk b = (s[i+2] << 24) | (s[i+3] << 16) | (s[i+0] << 8) | s[i+1];
        seen.emplace(b);
    }

#if 0
    printf("%zd\n", seen.size());
    for(auto &i : seen) {
        printf("%08x - %c%c%c%c\n", i, (i >> 24), (i >> 16) & 0xff, (i >> 8) & 0xff, i & 0xff);
    }
#endif
    return seen;
}

void worker(int ignored, auto &corpus, auto cb) {
    u16 crypt[2];
    u16 k[32] = {}, k0[4];
    
    getrandom((void*)k0, 8, GRND_NONBLOCK);
    memcpy(k, k0, sizeof(k0));

    KeyExpansion ( k );
    auto c1 = blk{}, c2 = blk{}, p1 = blk{}, p2 = blk{};

    auto it = corpus.begin(), end = corpus.end();
    for(; it != end; it++) {
        Encrypt((u16*)&*it, crypt, k);
        if((crypt[0] | crypt[1]) & 0x8080) continue;
        auto jt = corpus.find(*(blk*)&crypt[0]);
        if(jt != end)
        {
            p1 = *it;
            c1 = *jt;
            break;
        }
    }
#if 0
    p2 = p1; c2 = c1;
#else
    for(it++; it != end; it++) {
        Encrypt((u16*)&*it, crypt, k);
        if((crypt[0] | crypt[1]) & 0x8080) continue;
        auto jt = corpus.find(*(blk*)&crypt[0]);
        if(jt != end)
        {
            p2 = *it;
            c2 = *jt;
            break;
        }
    }
#endif

    if(c1 && c2) {
        cb(p1, p2, c1, c2, k0);
    }
}

int main ()
{
    auto corpus = read_corpus("pg10.txt");
    #pragma omp parallel for
    for(auto i = 0; i<10000; i++) {
        worker(i, corpus, [&](auto p1, auto p2, auto c1, auto c2, auto k0) {
            #pragma omp critical
            printf("(%04x%04x%04x%04x) -> (%04x%04x%04x%04x) K=%04x%04x%04x%04x ",
                p1 & 0xffff, p1 >> 16, p2 & 0xffff, p2 >> 16,
                c1 & 0xffff, c1 >> 16, c2 & 0xffff, c2 >> 16,
                k0[3], k0[2], k0[1], k0[0]);
            printf("'%c%c%c%c%c%c%c%c' -> '%c%c%c%c%c%c%c%c'\n",
                (p1 >> 8) & 0xff, (p1 >> 0) & 0xff, (p1 >> 24) & 0xff, (p1 >> 16) & 0xff,
                (p2 >> 8) & 0xff, (p2 >> 0) & 0xff, (p2 >> 24) & 0xff, (p2 >> 16) & 0xff,
                (c1 >> 8) & 0xff, (c1 >> 0) & 0xff, (c1 >> 24) & 0xff, (c1 >> 16) & 0xff,
                (c2 >> 8) & 0xff, (c2 >> 0) & 0xff, (c2 >> 24) & 0xff, (c2 >> 16) & 0xff);
        });
    }
    return 0;
}