/* Simple rainbow tables with distinguished points, experiment / toy example */
/* 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 64
/* 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 color_count 2000000

#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)

/* (candidate_number & mask) == 0  ends chains */
#define dp_mask 0xff

/* used to allocate memory to remember the table, specific for the prototype */
#define max_chains 100000000ULL

/* Postponing parameters */

/* skip in the first pass if chain is longer than  max_x_dp_mask * (dp_mask + 1) */
#define max_x_dp_mask 3.6
/* skip in the first pass if chain is shorter than postpone_min_len */
#define postpone_min_len 0

/* 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;

#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 reset_tried(num)

/* #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) */

static chain_t *generate()
{
    unsigned long long hash_k = 0;
    unsigned char *tried = malloc(tracking_buffer);
    if (!tried)
        die("malloc() failed\n");
    memset(tried, 0, tracking_buffer);
    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;
    position_in_chain_t max_len = 0;
#define max_intermediates ((unsigned long long)((dp_mask + 1) * max_x_dp_mask))
    candidate_number_t chain_intermediates[max_intermediates];
    unsigned long long skipped = 0;
    unsigned int oi;
    for (oi = 0; oi < 2; oi++) {
        int skip_bad_chains = !oi;
        /* int skip_bad_chains = 0; */
        for (i = 0; i < total_count; i++) {
            if ((i & 0xffff) == 0)
#define print_stats() printf("i %llu, ck %llu, rr %.3f, ml %llu, hk %dks, s %llu %.3f %.3f\n", i, ck, (skip_bad_chains ? (double)(i - ck) / i : 0), max_len, hash_k / total_count, skipped, (double)skipped / total_count, (double)skipped / i)
                print_stats();
            if (!check_tried(i)) {
                candidate_number_t ni = i;
                position_in_chain_t n = 0;
                do {
                    unsigned char *w = number_to_candidate(ni);
                    if (skip_bad_chains) {
                        if (n < max_intermediates)
                            chain_intermediates[n] = ni;
                    } else {
                        set_tried(ni);
                    }
                    unsigned char h[digest_size];
                    hash_it(w, h);
                    hash_k++;
                    ni = hash_to_number(h, n);
                    n++;
                } while ((ni & dp_mask) && (n < max_intermediates || !skip_bad_chains));
                if ((n >= postpone_min_len && n < max_intermediates) || !skip_bad_chains) {
                    if (skip_bad_chains) {
                        position_in_chain_t ti;
                        for (ti = 0; ti < n; ti++) {
                            set_tried(chain_intermediates[ti]);
                        }
                    }
                    if (ck < max_chains - 1) {
                        chains[ck].start = i;
                        chains[ck].end = ni;
                        chains[ck].length = n;
                        ck++;
                    } else {
                        die("too many chains\n");
                    }
                    if (n > max_len) {
                        max_len = n;
                    }
                } else {
                    skipped++;
                }
            }
            reset_tried(i);
        }
        print_stats();
        candidate_number_t not_covered = 0;
        for (i = 0; i < total_count; i++) {
            if (!check_tried(i))
                not_covered++;
        }
        printf("not covered: %llu %f\n", not_covered, (double)not_covered / total_count);
    }
    free(tried);
    return chains;
}

int cmp_length(const void *a, const void *b)
{
    /* %% undefined behaviour when the delta is in not range of int type? */
    return ((const chain_t *)a)->length - ((const chain_t *)b)->length;
}

int cmp_end(const void *a, const void *b)
{
    candidate_number_t ae = ((const chain_t *)a)->end;
    candidate_number_t be = ((const chain_t *)b)->end;
    return ae < be ? -1 : ae > be;
}

static inline chain_t *search_linear(chain_t *chains, unsigned long long k, candidate_number_t n)
{
    unsigned long long p;
    for (p = 0; p < k && chains[p].end != n; p++)
        ;
    return p < k ? &chains[p] : 0;
}

static inline chain_t *search_binary(chain_t *chains, unsigned long long k, candidate_number_t n)
{
    chain_t key;
    key.end = n;
    chain_t *r;
    r = bsearch(&key, chains, k, sizeof(key), cmp_end);
    if (!r)
        return 0;
    while (r > chains && r->end == n) {
        r--;
    }
    if (r->end != n)
        r++;
    return r;
}

#define search search_binary
/* #define search search_linear */

candidate_number_t crack_hash_no_dp(chain_t *chains, unsigned long long ml, unsigned long long k, const unsigned char *orig_h)
{
    ml += 1;
    unsigned char h[digest_size], ch[digest_size];
    unsigned char *w, *cw;
    candidate_number_t cn, ni, i, offset, end, t;
    for (offset = 0; offset < ml; offset++) {
        memcpy(h, orig_h, digest_size);
        for (i = offset; i < ml; i++) {
            ni = hash_to_number(h, i);
            /* search возвращает указатель на первую цепочку */
            chain_t *p = search(chains, k, ni);
            if (p) {
                for (end = p->end; p->end == end && p <= &chains[k - 1]; p++) {
                    /* memcpy(ch, h, digest_size); */
                    cn = p->start;
                    for (t = 0; t < p->length; t++) {
                        cw = number_to_candidate(cn);
                        hash_it(cw, ch);
                        if (!memcmp(ch, orig_h, digest_size)) {
                            return cn;
                        }
                        cn = hash_to_number(ch, t);
                    }
                }
            }
            /* We end chain on Distinguished Point. */
            if ((ni & dp_mask) == 0)
                break;
            w = number_to_candidate(ni);
            hash_it(w, h);
        }
    }
    /* %% Bad value, нужно специальное значения для фейла, возможно, через дополнительный параметр */
    return 0;
}

candidate_number_t crack_hash_dp(chain_t *chains, unsigned long long ml, unsigned long long k, const unsigned char *orig_h)
{
    ml += 1;
    /* %% длина цепочки; промежуточные точки для неполного вычисления */
    unsigned char h[digest_size], ch[digest_size];
    unsigned char *w, *cw;
    candidate_number_t cn, ni, i, offset, end, t;
    for (offset = 0; offset < ml; offset++) {
        memcpy(h, orig_h, digest_size);
        i = offset;
        do {
            ni = hash_to_number(h, i);
            w = number_to_candidate(ni);
            hash_it(w, h);
            i++;
        } while (i < ml && (ni & dp_mask));
        if ((ni & dp_mask) == 0) {
            /* If we found a Distinguished Point, then search it among chains. */
            chain_t *p = search(chains, k, ni);
            /* printf("%p\n", p); */
            if (p) {
                /* %% тут мы можем выехать за границу массива, хорошо бы иметь ноль там, или проверять длину */
                for (end = p->end; p->end == end; p++) {
                    /* memcpy(ch, h, digest_size); */
                    cn = p->start;
                    for (t = 0; t < p->length; t++) {
                        cw = number_to_candidate(cn);
                        hash_it(cw, ch);
                        if (!memcmp(ch, orig_h, digest_size)) {
                            return cn;
                        }
                        cn = hash_to_number(ch, t);
                    }
                }
            }
        }
    }
    /* %% So we can't distinguish 0 and 'Not found', make an arg for that. */
    return 0;
}

/* #define crack_hash crack_hash_no_dp */
#define crack_hash crack_hash_dp

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

    unsigned long long i, k = 0;
    for (i = 1; chains[i].length; i++)
        ;
    k = i;
    printf("k = %llu\n", k);

    double ave = 0;
    for (i = 0; i < k; i++) {
        ave += chains[i].length;
    }
    ave /= k;
    printf("ave: %f\n", ave);

    unsigned long long short_k = 0;
    for (i = 0; i < k; i++) {
        if (chains[i].length < ave) {
            short_k++;
        }
    }
    printf("short_k = %llu  %f\n", short_k, (double)short_k / k);

    qsort(chains, k, sizeof(chain_t), cmp_length);
    double avep = 0;
    double part = 0.95;
    unsigned long long kk = k * part;
    for (i = 0; i < kk; i++) {
        avep += chains[i].length;
    }
    avep /= kk;
    printf("ave %f for part %f\n", avep, part);
    unsigned long long ml = chains[k - 1].length;
#define pmp2(p, len) printf("ml %llu for part %f; ml x %.3f, ave x %.3f\n", (len), (p), (double)(len) / ml, (double)(len) / ave)
#define pmp(p) pmp2((p), chains[(unsigned long long)(k * (p) - 1)].length)
    pmp(0.5);
    pmp(0.8);
    pmp(0.9);
    pmp(0.95);
    pmp(0.98);
    pmp(0.99);
    pmp(0.999);
    pmp(1.0);

    qsort(chains, k, sizeof(chain_t), cmp_end);
    const char *test_good = "a0 a1 a3";
    const char *test_bad = "bad";
    unsigned char h[digest_size];
    unsigned long long r;
    hash_it(test_good, h);
    r = crack_hash(chains, ml, k, h);
    printf("good: %llu\n", r);
    printf("good: %s\n", number_to_candidate(r));
    hash_it(test_bad, h);

    /* r = crack_hash(chains, ml, k, h); */
    /* printf("bad: %llu\n", r); */

#if 0
    const char *test;
    test = 1 ? test_good : test_bad;
    hash_it(test, h);
    for (i = 0; i < 1000; i++) {
        r = crack_hash(chains, ml, k, h);
        if ((i & 0xf) == 0) {
            printf("i = %llu, r = %llu\n", i, r);
            if ((test == test_good) == !r) {
                printf("failure\n");
                return 1;
            }
        }
    }
#endif

#if 0
    printf("before full test\n");
    char *w1, *w2, *w3, *wr;
    size_t i1, i2, i3, l1, l2, l3, gk = 0;
#define forw(w, i, l) for (w = words[i = 0], l = strlen(w); i < number_of_words; w = words[++i], l = strlen(w))
    forw(w1, i1, l1) forw(w2, i2, l2) forw(w3, i3, l3) {
        hash_ctx c;
        unsigned char o[digest_size];
        hash_init(&c);
        hash_update(&c, w1, l1);
        hash_update(&c, " ", 1);
        hash_update(&c, w2, l2);
        hash_update(&c, " ", 1);
        hash_update(&c, w3, l3);
        hash_final(o, &c);
        r = crack_hash(chains, ml, k, o);
        wr = number_to_candidate(r);
        if (wr[l1] == ' ' && wr[l2 + l1 + 1] == ' ' && wr[l1 + l2 + l3 + 2] == '\0' && !strncmp(wr, w1, l1) && !strncmp(&wr[l1 + 1], w2, l2) && !strncmp(&wr[l1 + l2 + 2], w3, l3)) {
            /* ok */
            gk++;
            if ((gk & 0xff) == 0) {
                printf("ok: %llu\n", gk);
            }
        } else {
            printf("fail: wr = %s, right = '%s %s %s', r = %llu\n", wr, w1, w2, w3, r);
        }
    }
#endif

    free(chains);
    return 0;
}