/* Simple rainbow tables, experiment / toy example */
/* No cracking, never tested the generator */
/* Implementation with openssl */

/* gcc -O3 -lcrypto ossl.c && time ./a.out */

/*
 * Copyright © 2015 Aleksey Cherepanov <lyosha@openwall.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted.
 */

#pragma GCC diagnostic ignored "-Wdeclaration-after-statement"
#pragma GCC optimize 3

/* Tunable parameters */

#define word_count 3
#define number_of_words 256
/* Not really tunable: total_count == power(number_of_words, word_count) */
/* TODO: what type would 256**5 have? */
/* TODO: total_count can be too big for any type. */
#define total_count (number_of_words * number_of_words * number_of_words)
#define max_length 16

#define chain_length 200
#define color_count chain_length

#define digest_size MD5_DIGEST_LENGTH
#define hash_ctx MD5_CTX
#define hash_init MD5_Init
#define hash_update MD5_Update
#define hash_final MD5_Final

#define tracking_buffer (256 * 256 * 256 / 8 / 8)

/* Code */

#include <stdio.h>
#include <openssl/sha.h>
#include <openssl/md5.h>
#include <time.h>
#include <stdlib.h>

#include <assert.h>

#include <string.h>

#define candidate_number_t unsigned long long
#define position_in_chain_t unsigned long long

/* perl -e 'for (0..255) { print "\"a$_\", " }' */
static const char *words[] = {
    "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a8", "a9", "a10",
    "a11", "a12", "a13", "a14", "a15", "a16", "a17", "a18", "a19", "a20",
    "a21", "a22", "a23", "a24", "a25", "a26", "a27", "a28", "a29", "a30",
    "a31", "a32", "a33", "a34", "a35", "a36", "a37", "a38", "a39", "a40",
    "a41", "a42", "a43", "a44", "a45", "a46", "a47", "a48", "a49", "a50",
    "a51", "a52", "a53", "a54", "a55", "a56", "a57", "a58", "a59", "a60",
    "a61", "a62", "a63", "a64", "a65", "a66", "a67", "a68", "a69", "a70",
    "a71", "a72", "a73", "a74", "a75", "a76", "a77", "a78", "a79", "a80",
    "a81", "a82", "a83", "a84", "a85", "a86", "a87", "a88", "a89", "a90",
    "a91", "a92", "a93", "a94", "a95", "a96", "a97", "a98", "a99", "a100",
    "a101", "a102", "a103", "a104", "a105", "a106", "a107", "a108", "a109",
    "a110", "a111", "a112", "a113", "a114", "a115", "a116", "a117", "a118",
    "a119", "a120", "a121", "a122", "a123", "a124", "a125", "a126", "a127",
    "a128", "a129", "a130", "a131", "a132", "a133", "a134", "a135", "a136",
    "a137", "a138", "a139", "a140", "a141", "a142", "a143", "a144", "a145",
    "a146", "a147", "a148", "a149", "a150", "a151", "a152", "a153", "a154",
    "a155", "a156", "a157", "a158", "a159", "a160", "a161", "a162", "a163",
    "a164", "a165", "a166", "a167", "a168", "a169", "a170", "a171", "a172",
    "a173", "a174", "a175", "a176", "a177", "a178", "a179", "a180", "a181",
    "a182", "a183", "a184", "a185", "a186", "a187", "a188", "a189", "a190",
    "a191", "a192", "a193", "a194", "a195", "a196", "a197", "a198", "a199",
    "a200", "a201", "a202", "a203", "a204", "a205", "a206", "a207", "a208",
    "a209", "a210", "a211", "a212", "a213", "a214", "a215", "a216", "a217",
    "a218", "a219", "a220", "a221", "a222", "a223", "a224", "a225", "a226",
    "a227", "a228", "a229", "a230", "a231", "a232", "a233", "a234", "a235",
    "a236", "a237", "a238", "a239", "a240", "a241", "a242", "a243", "a244",
    "a245", "a246", "a247", "a248", "a249", "a250", "a251", "a252", "a253",
    "a254", "a255"
};

static inline char *number_to_candidate(candidate_number_t number)
{
    static char buf[max_length];
    int i;
    buf[0] = 0;
    for (i = 0; i < word_count; i++) {
        if (i)
            strcat(buf, " ");
        strcat(buf, words[number % number_of_words]);
        number /= number_of_words;
    }
    return buf;
}

static inline candidate_number_t hash_to_number(unsigned char *hash, position_in_chain_t position_in_chain)
{
    return (*((candidate_number_t *)hash) + (position_in_chain % color_count)) % total_count;
}

#define hash_it(buf, out) \
    do { \
        hash_ctx ctx; \
        hash_init(&ctx); \
        hash_update(&ctx, (buf), strlen((buf))); \
        hash_final((out), &ctx); \
    } while(0)

static void die(const char *msg)
{
    fprintf(stderr, "%s", msg);
    exit(1);
}

typedef struct {
    candidate_number_t start;
    candidate_number_t end;
    position_in_chain_t length;
} chain_t;

static chain_t *generate()
{
    unsigned char *tried = malloc(tracking_buffer);
    if (!tried)
        die("malloc() failed\n");

/* #define check_tried(num) (tried[(num) / 8] & (1 << ((num) % 8))) */
/* #define set_tried(num) (tried[(num) / 8] |= (1 << ((num) % 8))) */
/* #define reset_tried(num) (tried[(num) / 8] &= ~(unsigned char)(1 << ((num) % 8))) */

#define C (tracking_buffer * 8)
#define is_in_buffer(num) ((num) >= i && (num) < i + C)
    /* position in buffer */
#define pob(num) ((num) % C)

#define check_tried(num) (is_in_buffer((num)) ? (tried[pob((num)) / 8] & (1 << (pob((num)) % 8))) : 0)
#define set_tried(num)   (is_in_buffer((num)) ? (tried[pob((num)) / 8] |= (1 << (pob((num)) % 8))) : 0)
#define reset_tried(num) (is_in_buffer((num)) ? (tried[pob((num)) / 8] &= ~(unsigned char)(1 << (pob((num)) % 8))) : 0)

    memset(tried, 0, tracking_buffer);
#define max_chains 100000000ULL
    chain_t *chains = malloc(max_chains * sizeof(chain_t));
    if (!chains)
        die("malloc() failed 2\n");
    memset(chains, 0, max_chains * sizeof(chain_t));
    unsigned long long ck = 0;
    candidate_number_t i;
    for (i = 0; i < total_count; i++) {
        if ((i & 0xffff) == 0)
            printf("i = %llu, ck = %llu, reuse ratio = %.5f\n", i, ck, (double)(i - ck) / i);
        if (!check_tried(i)) {
            candidate_number_t ni = i;
            position_in_chain_t n;
            for (n = 0; n < chain_length; n++) {
                unsigned char *w = number_to_candidate(ni);
                set_tried(ni);
                unsigned char h[digest_size];
                hash_it(w, h);
                ni = hash_to_number(h, n);
            }
            if (ck < max_chains - 1) {
                chains[ck].start = i;
                chains[ck].end = ni;
                chains[ck].length = n;
                ck++;
            } else {
                die("too many chains\n");
            }
        }
        reset_tried(i);
    }
    return chains;
}

int main(void)
{
    chain_t *chains = generate();
    return 0;
}