/*
 *  zc - zip crack library
 *  Copyright (C) 2012-2018 Marc Ferland
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdlib.h>
#include <stdbool.h>
#include <math.h>

#include "pwstream.h"
#include "libzc_private.h"

static const struct entry null_entry = { -1, -1, -1};

/*
 This algorithm distributes a pool of characters to 'n' password
 streams. These streams can then generate passwords that are mutually
 exclusive. For example:

 pool: a,b,c

 Distribute to 4 streams of 3 characters (last character appears first):
 stream0: [a,a] [a,a] [a,c] --> aaa, baa, caa
 stream1: [a,a] [b,c] [a,c] --> aba, aca, bba, bca, cba, cca
 stream2: [b,b] [a,c] [a,c] --> aab, abb, acb, bab, bbb, bcb, cab, cbb, ccb
 stream3: [c,c] [a,c] [a,c] --> aac, abc, acc, bac, bbc, bcc, cac, cbc, ccc

 The output of the algorithm is a table of streams*characters
 entries. Each entry contains a start and end offset to the character
 pool. Using these offsets we can easily loop over all possible
 'strings' using simple for-loops.

 Pseudo-code (again for 3 characters strings):
 for (i = table[stream0, 2].start; i <= table[stream0, 2].stop; ++i)
    str[0] = pool[i];
    for (j = table[stream0, 1].start; j <= table[stream0, 1].stop; ++j)
       str[1] = pool[j];
       for (k = table[stream0, 0].start; k <= table[stream0, 0].stop; ++k)
          str[2] = pool[k];
          do_stuff(str);

 The entry table is divided in the following way:
 +----------------------------------+
 |s0c0|s1c0|s2c0|s0c1|s1c1|s2c1|....|
 +----------------------------------+
 where s0-->stream0 and c0-->character0.
 */

struct pwstream {
	struct entry *entry;
	size_t initial;
	size_t rows;
	size_t cols;
	size_t real_cols;
	size_t plen;
};

static int compare_entries(const void *a, const void *b)
{
	const struct entry *ea = (const struct entry *)a;
	const struct entry *eb = (const struct entry *)b;

	/* since entries are always mutually exclusive, compare only the
	 * 'start' member */
	return (ea->start > eb->start) - (ea->start < eb->start);
}

static void sort(struct entry *e, size_t streams)
{
	qsort(e, streams, sizeof(struct entry), compare_entries);
}

/**
 * Get entry pointer at [row][col].
 */
static struct entry *get(struct pwstream *t, size_t row, size_t col)
{
	return &t->entry[t->cols * row + col];
}

/**
 * Split the given sequence of length 'len' in streams.
 * i.e.: 1,2,3,4,5 in 3 streams ==> [1,2],[3,4],[5]
 */
static void split_less(size_t plen, size_t streams, struct entry *t)
{
	for (size_t i = 0; i < streams; ++i) {
		int start = (i * plen) / streams;
		t[i].start = start;
		t[i].initial = start;
		t[i].stop = ((i + 1) * plen) / streams - 1;
	}
}

/**
 * 1,2,3 in 4 streams ==> [1],[1],[2],[3]
 */
static void split_more(size_t plen, size_t streams, struct entry *e)
{
	for (size_t i = 0; i < streams; ++i) {
		int tmp = i % plen;
		e[i].start = tmp;
		e[i].stop = tmp;
		e[i].initial = tmp;
	}
	sort(e, streams);
}

/**
 * 1,2,3,4 in 4 streams ==> [1],[2],[3],[4]
 */
static void split_equal(size_t plen, struct entry *e)
{
	for (size_t i = 0; i < plen; ++i) {
		e[i].start = i;
		e[i].stop = i;
		e[i].initial = i;
	}
}

/**
 * Distribute sequence @seq in @streams and store the result at @entry.
 */
static void distribute(size_t plen, size_t streams, struct entry *entry)
{
	if (streams == 1) {
		entry->start = 0;
		entry->stop = plen - 1;
		entry->initial = 0;
	} else if (streams == plen) {
		split_equal(plen, entry);
	} else if (streams > plen) {
		split_more(plen, streams, entry);
	} else
		split_less(plen, streams, entry);
}

/**
 * Compares entries e1 and e2.
 */
static bool is_equal_entries(const struct entry *e1, const struct entry *e2)
{
	return (e1->start == e2->start && e1->stop == e2->stop);
}

/**
 * Initialize entry table with default values.
 */
static void entry_table_init(struct pwstream *t, int start, int stop)
{
	for (size_t i = 0; i < t->rows * t->cols; ++i) {
		t->entry[i].start = start;
		t->entry[i].stop = stop;
		t->entry[i].initial = start;
	}
}

/**
 * Count the number of entries identical to @e on row @row.
 */
static size_t uniq(struct pwstream *t, size_t row, const struct entry *e)
{
	size_t count = 0;
	struct entry *n = get(t, row, 0);

	for (size_t i = 0; i < t->cols; ++i) {
		if (is_equal_entries(&n[i], e))
			++count;
	}
	return count;
}

/**
 * Count the number of entries identical to @e and consider @len
 * entries.
 */
static size_t uniq_from_entry(const struct entry *e, size_t len)
{
	size_t count = 1;           /* first entry is always equal */
	for (size_t i = 1; i < len; ++i) {
		if (is_equal_entries(&e[i], e))
			++count;
	}
	return count;
}

static void recurse(struct pwstream *t, size_t count, struct entry *e)
{
	if (count == 1)
		return;

	distribute(t->plen, count, e);

	size_t u = 0;
	for (size_t i = 0; i < count; i += u) {
		u = uniq_from_entry(&e[i], count - i);
		recurse(t, u, &e[i + t->cols]);
	}
}

static void generate(struct pwstream *pws)
{
	/* do a first distribution for character 0 */
	distribute(pws->plen, pws->cols, pws->entry);

	/* generate the remaining entries */
	for (size_t i = 0, u = 0; i < pws->cols; i += u) {
		u = uniq(pws, 0, &pws->entry[i]);
		if (u > 1)
			recurse(pws, u, get(pws, 1, i));
	}
}

static bool is_before(const struct entry *e, int c)
{
	return (c < e->start);
}

static bool is_after(const struct entry *e, int c)
{
	return (c > e->stop);
}

static bool is_enclosed(const struct entry *e, int c)
{
	return !is_before(e, c) && !is_after(e, c);
}

static void generate_initial_indexes(struct pwstream *pws,
				     const size_t *initial)
{
	for (size_t i = 0; i < pws->rows; ++i) {
		for (size_t j = 0; j < pws->cols; ++j) {
			struct entry *e = get(pws, i, j);
			if (is_enclosed(e, initial[i]))
				e->initial = initial[i];
			else if (is_after(e, initial[i]))
				e->initial = e->stop;
			else
				e->initial = e->start;
		}
	}
}

static size_t ceil_streams(size_t pool_len, size_t pw_len, size_t streams)
{
	long double permut = powl((long double)pool_len, (long double)pw_len);
	if (permut == HUGE_VALL)
		/* assume we won't ever have more than HUGE_VAL streams */
		return streams;
	else if (permut < (long double)streams)
		return (size_t)permut;
	return streams;
}

int pwstream_new(struct pwstream **pws)
{
	struct pwstream *p = calloc(1, sizeof(struct pwstream));

	if (!p)
		return -1;

	p->entry = NULL;
	p->rows = 0;
	p->cols = 0;
	p->plen = 0;

	*pws = p;

	return 0;
}

void pwstream_free(struct pwstream *pws)
{
	if (pws->entry)
		free(pws->entry);
	free(pws);
}

int pwstream_generate(struct pwstream *pws, size_t pool_len, size_t pw_len,
		      size_t streams, const size_t *initial)
{
	if (pws->entry)
		free(pws->entry);

	size_t cstrm = ceil_streams(pool_len, pw_len, streams);

	pws->entry = calloc(cstrm * pw_len, sizeof(struct entry));
	if (!pws->entry) {
		pws->rows = 0;
		pws->cols = 0;
		pws->plen = 0;
		pws->real_cols = 0;
		return -1;
	}

	pws->rows = pw_len;
	pws->cols = cstrm;
	pws->plen = pool_len;
	pws->real_cols = streams;

	entry_table_init(pws, 0, pool_len - 1);
	generate(pws);

	if (initial)
		generate_initial_indexes(pws, initial);

	return 0;
}

const struct entry *pwstream_get_entry(struct pwstream *pws,
				       size_t stream, size_t pos)
{
	if (stream >= pws->cols)
		return &null_entry;
	return get(pws, pos, stream);
}

size_t pwstream_get_pwlen(const struct pwstream *pws)
{
	return pws->rows;
}

size_t pwstream_get_stream_count(const struct pwstream *pws)
{
	return pws->real_cols;
}

bool pwstream_is_empty(struct pwstream *pws, unsigned int stream)
{
	return stream >= pws->cols;
}