#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <pthread.h>
#include <time.h>

#define prefetch(x) //
#define prefetchw(x) //
#define _C(fcolor) "\e[1;" #fcolor "m"
#define _DC "\e[1;%um"
#define ROTATE_LEFT(x, n) (((x) << (n)) | ((x) >> (32-(n))))
#define F(x,y,z) (((x) & (y)) | ((~(x)) & (z)))
#define G(x,y,z) (((x) & (z)) | ((y) & (~(z))))
#define H(x,y,z) ((x) ^ (y) ^ (z))
#define I(x,y,z) ((y) ^ ((x) | (~(z))))
#define md5_s1_0 7
#define md5_s1_1 12
#define md5_s1_2 17
#define md5_s1_3 22
#define md5_s2_0 5
#define md5_s2_1 9 
#define md5_s2_2 14
#define md5_s2_3 20
#define md5_s3_0 4
#define md5_s3_1 11
#define md5_s3_2 16
#define md5_s3_3 23
#define md5_s4_0 6
#define md5_s4_1 10
#define md5_s4_2 15
#define md5_s4_3 21

#define ROTATE_LEFT(x, n) (((x) << (n)) | ((x) >> (32-(n))))
#define F(x,y,z) (((x) & (y)) | ((~(x)) & (z)))
#define G(x,y,z) (((x) & (z)) | ((y) & (~(z))))
#define H(x,y,z) ((x) ^ (y) ^ (z))
#define I(x,y,z) ((y) ^ ((x) | (~(z))))

uint32_t md5_hash(uint32_t *X, uint32_t *digest)
{
	uint32_t A, B, C, D;
	A = 0x67452301; 
	B = 0xefcdab89; 
	C = 0x98badcfe; 
	D = 0x10325476;
	/// round one (unrolled)
	A = B + ROTATE_LEFT((A + F(B, C, D) + X[0] + 0xd76aa478), md5_s1_0);
	D = A + ROTATE_LEFT((D + F(A, B, C) + X[1] + 0xe8c7b756), md5_s1_1);
	C = D + ROTATE_LEFT((C + F(D, A, B) + X[2] + 0x242070db), md5_s1_2);
	B = C + ROTATE_LEFT((B + F(C, D, A) + X[3] + 0xc1bdceee), md5_s1_3);
	A = B + ROTATE_LEFT((A + F(B, C, D) + X[4] + 0xf57c0faf), md5_s1_0);
	D = A + ROTATE_LEFT((D + F(A, B, C) + X[5] + 0x4787c62a), md5_s1_1);
	C = D + ROTATE_LEFT((C + F(D, A, B) + X[6] + 0xa8304613), md5_s1_2);
	B = C + ROTATE_LEFT((B + F(C, D, A) + X[7] + 0xfd469501), md5_s1_3);
	A = B + ROTATE_LEFT((A + F(B, C, D) + X[8] + 0x698098d8), md5_s1_0);
	D = A + ROTATE_LEFT((D + F(A, B, C) + X[9] + 0x8b44f7af), md5_s1_1);
	C = D + ROTATE_LEFT((C + F(D, A, B) + X[10] + 0xffff5bb1), md5_s1_2);
	B = C + ROTATE_LEFT((B + F(C, D, A) + X[11] + 0x895cd7be), md5_s1_3);
	A = B + ROTATE_LEFT((A + F(B, C, D) + X[12] + 0x6b901122), md5_s1_0);
	D = A + ROTATE_LEFT((D + F(A, B, C) + X[13] + 0xfd987193), md5_s1_1);
	C = D + ROTATE_LEFT((C + F(D, A, B) + X[14] + 0xa679438e), md5_s1_2);
	B = C + ROTATE_LEFT((B + F(C, D, A) + X[15] + 0x49b40821), md5_s1_3);
	/// round two (unrolled)
	A = B + ROTATE_LEFT((A + G(B, C, D) + X[1] + 0xf61e2562), md5_s2_0);
	D = A + ROTATE_LEFT((D + G(A, B, C) + X[6] + 0xc040b340), md5_s2_1);
	C = D + ROTATE_LEFT((C + G(D, A, B) + X[11] + 0x265e5a51), md5_s2_2);
	B = C + ROTATE_LEFT((B + G(C, D, A) + X[0] + 0xe9b6c7aa), md5_s2_3);
	A = B + ROTATE_LEFT((A + G(B, C, D) + X[5] + 0xd62f105d), md5_s2_0);
	D = A + ROTATE_LEFT((D + G(A, B, C) + X[10] + 0x02441453), md5_s2_1);
	C = D + ROTATE_LEFT((C + G(D, A, B) + X[15] + 0xd8a1e681), md5_s2_2);
	B = C + ROTATE_LEFT((B + G(C, D, A) + X[4] + 0xe7d3fbc8), md5_s2_3);
	A = B + ROTATE_LEFT((A + G(B, C, D) + X[9] + 0x21e1cde6), md5_s2_0);
	D = A + ROTATE_LEFT((D + G(A, B, C) + X[14] + 0xc33707d6), md5_s2_1);
	C = D + ROTATE_LEFT((C + G(D, A, B) + X[3] + 0xf4d50d87), md5_s2_2);
	B = C + ROTATE_LEFT((B + G(C, D, A) + X[8] + 0x455a14ed), md5_s2_3);
	A = B + ROTATE_LEFT((A + G(B, C, D) + X[13] + 0xa9e3e905), md5_s2_0);
	D = A + ROTATE_LEFT((D + G(A, B, C) + X[2] + 0xfcefa3f8), md5_s2_1);
	C = D + ROTATE_LEFT((C + G(D, A, B) + X[7] + 0x676f02d9), md5_s2_2);
	B = C + ROTATE_LEFT((B + G(C, D, A) + X[12] + 0x8d2a4c8a), md5_s2_3);
	/// round three (unrolled)
	A = B + ROTATE_LEFT((A + H(B, C, D) + X[5] + 0xfffa3942), md5_s3_0);
	D = A + ROTATE_LEFT((D + H(A, B, C) + X[8] + 0x8771f681), md5_s3_1);
	C = D + ROTATE_LEFT((C + H(D, A, B) + X[11] + 0x6d9d6122), md5_s3_2);
	B = C + ROTATE_LEFT((B + H(C, D, A) + X[14] + 0xfde5380c), md5_s3_3);
	A = B + ROTATE_LEFT((A + H(B, C, D) + X[1] + 0xa4beea44), md5_s3_0);
	D = A + ROTATE_LEFT((D + H(A, B, C) + X[4] + 0x4bdecfa9), md5_s3_1);
	C = D + ROTATE_LEFT((C + H(D, A, B) + X[7] + 0xf6bb4b60), md5_s3_2);
	B = C + ROTATE_LEFT((B + H(C, D, A) + X[10] + 0xbebfbc70), md5_s3_3);
	A = B + ROTATE_LEFT((A + H(B, C, D) + X[13] + 0x289b7ec6), md5_s3_0);
	D = A + ROTATE_LEFT((D + H(A, B, C) + X[0] + 0xeaa127fa), md5_s3_1);
	C = D + ROTATE_LEFT((C + H(D, A, B) + X[3] + 0xd4ef3085), md5_s3_2);
	B = C + ROTATE_LEFT((B + H(C, D, A) + X[6] + 0x04881d05), md5_s3_3);
	A = B + ROTATE_LEFT((A + H(B, C, D) + X[9] + 0xd9d4d039), md5_s3_0);
	D = A + ROTATE_LEFT((D + H(A, B, C) + X[12] + 0xe6db99e5), md5_s3_1);
	C = D + ROTATE_LEFT((C + H(D, A, B) + X[15] + 0x1fa27cf8), md5_s3_2);
	B = C + ROTATE_LEFT((B + H(C, D, A) + X[2] + 0xc4ac5665), md5_s3_3);
	/// round four (unrolled)
	A = B + ROTATE_LEFT((A + I(B, C, D) + X[0] + 0xf4292244), md5_s4_0);
	D = A + ROTATE_LEFT((D + I(A, B, C) + X[7] + 0x432aff97), md5_s4_1);
	C = D + ROTATE_LEFT((C + I(D, A, B) + X[14] + 0xab9423a7), md5_s4_2);
	B = C + ROTATE_LEFT((B + I(C, D, A) + X[5] + 0xfc93a039), md5_s4_3);
	A = B + ROTATE_LEFT((A + I(B, C, D) + X[12] + 0x655b59c3), md5_s4_0);
	D = A + ROTATE_LEFT((D + I(A, B, C) + X[3] + 0x8f0ccc92), md5_s4_1);
	C = D + ROTATE_LEFT((C + I(D, A, B) + X[10] + 0xffeff47d), md5_s4_2);
	B = C + ROTATE_LEFT((B + I(C, D, A) + X[1] + 0x85845dd1), md5_s4_3);
	A = B + ROTATE_LEFT((A + I(B, C, D) + X[8] + 0x6fa87e4f), md5_s4_0);
	D = A + ROTATE_LEFT((D + I(A, B, C) + X[15] + 0xfe2ce6e0), md5_s4_1);
	C = D + ROTATE_LEFT((C + I(D, A, B) + X[6] + 0xa3014314), md5_s4_2);
	B = C + ROTATE_LEFT((B + I(C, D, A) + X[13] + 0x4e0811a1), md5_s4_3);
	A = B + ROTATE_LEFT((A + I(B, C, D) + X[4] + 0xf7537e82), md5_s4_0);
	D = A + ROTATE_LEFT((D + I(A, B, C) + X[11] + 0xbd3af235), md5_s4_1);
	C = D + ROTATE_LEFT((C + I(D, A, B) + X[2] + 0x2ad7d2bb), md5_s4_2);
	B = C + ROTATE_LEFT((B + I(C, D, A) + X[9] + 0xeb86d391), md5_s4_3);
	A += 0x67452301; 
	B += 0xefcdab89; 
	C += 0x98badcfe; 
	D += 0x10325476;
	if( (digest[0] == A) && (digest[1] == B) && (digest[2] == C) && (digest[3] == D) )
	{
		return 1;
	}
	return 0;
}

/// generate next sequence from current sequence.
// returns zero if no more sequences exist, or one for success.
int32_t sequence_next(uint8_t *sequence, uint32_t charset_length, uint32_t length)
{
   uint32_t x;
   for(x = 0; sequence[x] == (charset_length-1); ++x)
   {
      if(x == (length-1))
      {
         return 0;
      }
      sequence[x] = 0;
   }
   ++sequence[x];
   return 1;
}

int32_t sequence_nextx(uint8_t *sequence, uint32_t charset_length, uint32_t length, uint32_t m)
{
   uint32_t x;
   for(x = 0; sequence[x] >= (charset_length-1); ++x)
   {
      if(x == (length-1))
      {
         return 0;
      }
      sequence[x] = sequence[x] - (charset_length - 1);
   }
   sequence[x] += m;
   return 1;
}

struct tThreadInfo
{
	uint32_t	id;
	uint32_t	charset_len;
	uint32_t	charset_offs;
	uint32_t	seqlen_min;
	uint32_t	seqlen_max;
	uint32_t	jump;
	uint32_t	ejump;
};

uint8_t *g_charset_map;
uint32_t g_charset_len;

uint32_t g_seqlen_max;

uint32_t g_threads_stop;
volatile uint32_t g_thread_active_cnt;

uint32_t g_digest_target[4];

time_t g_start;

uint32_t lock(uint32_t state, uint32_t *lck)
{
	asm volatile("xchg %1, %0" : "=r" (state), "=m" (*lck) : "0" (state));
	return state;
}

void* threaded_worker(struct tThreadInfo *tinfo)
{
	/// create md5 buffer large enough for one block.
	uint8_t sequence[g_seqlen_max];
	uint8_t *l_charset_map;
	uint8_t md5buffer[64];
	uint32_t *md5buffer32 = (uint32_t*)&md5buffer[0];

	uint32_t x, l_charset_len, sequence_wlen;

	uint32_t digest[4];
	uint32_t loop = 0;

	time_t start;
	float diff;

	l_charset_len = tinfo->charset_len;
	l_charset_map = (uint8_t*)malloc(tinfo->charset_len);
	printf("cm:");
	for(x = tinfo->charset_offs; x < (tinfo->charset_offs + tinfo->charset_len); ++x)
	{
		l_charset_map[x - tinfo->charset_offs] = g_charset_map[x];
		printf("%c", g_charset_map[x]);
	}
	printf("\n");

	for(x = 0; x < 16; ++x)
	{
		md5buffer32[x] = 0;
	}

	start = clock();
	g_start = start;

	/// prepare sequence
	for(x = 0; x < tinfo->seqlen_max; ++x)
	{
		sequence[x] = 0;
	}

	printf("running\n");

	sequence_wlen = tinfo->seqlen_min;
	md5buffer32[15] = 0;
	while(sequence_wlen <= tinfo->seqlen_max)
	{
		sequence[sequence_wlen-1] = tinfo->jump;
		md5buffer[sequence_wlen] = 0x80;
		md5buffer32[14] = sequence_wlen * 8;
		/// write message, hash message, advance sequence
		while(1)
		{
			//printf("(%u)seq:", tinfo->id);
			for(x = 0; x < sequence_wlen; ++x)
			{
				md5buffer[x] = l_charset_map[sequence[x]];
				//printf("%c", md5buffer[x]);
			}
			//printf("\n");
			if(md5_hash((uint32_t*)&md5buffer[0], &g_digest_target[0]) == 1)
			{
				diff = clock() - g_start;
				printf("time: %f\n", (double)(clock() - g_start) / (double)CLOCKS_PER_SEC);
				printf("h/s:" "%f\n", loop / diff);
				printf("(%x) I found a collision!\n", tinfo->id);
				printf("message:");
				for(x = 0; x < sequence_wlen; ++x)
				{
					printf(_DC "%c", 30+(x%7), g_charset_map[sequence[x]]);
				}
				printf("\n");
				asm("lock decl %0" : "=m" (g_thread_active_cnt));
				printf("g_thread_active_cnt:%x\n", g_thread_active_cnt);
				exit(1);
				return 0;
			}
			++loop;
/*		
			if(loop & 16)
			{
				if(((clock() - start) / CLOCKS_PER_SEC) > 3)
				{
					start = clock();
					diff = (clock() - g_start) / CLOCKS_PER_SEC;
					printf("h/s:%f loop:%u diff:%f\n", (loop / diff) / 1000000.0, loop, diff);
				}
			}
*/
			if(sequence_next(&sequence[0], l_charset_len, sequence_wlen) == 0)
			{
				break;
			}
			if(sequence[sequence_wlen-1] == tinfo->ejump)
			{
				break;
			}
			
			if(g_threads_stop == 1)
			{
				printf("stopping\n");
				asm("lock decl %0" : "=m" (g_thread_active_cnt));
				printf("g_thread_active_cnt:%x\n", g_thread_active_cnt);
				return 0;
			}
			
		}
		++sequence_wlen;
	}

	asm("lock decl %0" : "=m" (g_thread_active_cnt));
	printf("g_thread_active_cnt:%x\n", g_thread_active_cnt);
	return 0;
}



int32_t md5_search_core(uint32_t CHARSET_LENGTH, uint32_t seqlen_min, uint32_t seqlen_max, uint32_t THREAD_COUNT)
{	
	uint32_t x;
	pthread_t thread_workers[THREAD_COUNT];
	struct tThreadInfo tinfo[THREAD_COUNT];
	g_thread_active_cnt = THREAD_COUNT;
	g_seqlen_max = seqlen_max;
	for(x = 1; x < THREAD_COUNT; ++x)
	{
		tinfo[x].jump = (CHARSET_LENGTH/THREAD_COUNT) * x;
		tinfo[x].ejump = (CHARSET_LENGTH/THREAD_COUNT) * (x+1);
		tinfo[x].id = x + 1;
		tinfo[x].charset_len = CHARSET_LENGTH;
		tinfo[x].charset_offs = 0;
		tinfo[x].seqlen_min = seqlen_min;
		tinfo[x].seqlen_max = seqlen_max;
		pthread_create(&thread_workers[x], NULL, (void*(*)(void*))&threaded_worker, (void*)(&tinfo[x]));
	}
	tinfo[0].jump = 0;
	tinfo[0].ejump = (CHARSET_LENGTH/THREAD_COUNT) * 1;
	tinfo[0].id = 1;
	tinfo[0].charset_len = CHARSET_LENGTH;
	tinfo[0].charset_offs = 0;
	tinfo[0].seqlen_min = seqlen_min;
	tinfo[0].seqlen_max = seqlen_max;
	threaded_worker(&tinfo[0]);
	// keep waiting until all have returned.
	while(g_thread_active_cnt != 0);
	return 1;
} 

int main(int argc, char *argv[])
{
	uint32_t t;
	char __charset[] = {'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z'}; 
	if( argc < 5 )
	{
		printf("[hash] [min-search-length] [max-search-length] [thread-count]\n");
		return -1;
	}
	printf("initializing..\n");
	
	g_charset_map = (uint8_t*)malloc(sizeof(__charset));
	memcpy(g_charset_map, __charset, sizeof(__charset));

	for(t = 0; t < 16; ++t)
	{
		sscanf(&argv[1][t*2], "%2x", &((uint8_t*)(&g_digest_target[0]))[t]);
	}

	

	printf("digest_target is " _C(4));
	for(t = 0; t < 4; ++t)
	{
		printf(_DC "%08x", 30+(t%7), g_digest_target[t]);
	}
	printf(_C(0) "\n");

	g_start = clock();
	/// <size of char map>, <maximum message size we are searching for>, <number of threads to use>
	md5_search_core(sizeof(__charset), atoi( argv[2] ), atoi( argv[3] ), atoi(argv[4]));
	return 1;
}