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

#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

void __attribute__((__fastcall__)) __attribute__((__always_inline__))
	md5_hash(uint8_t *nm8, uint32_t nmlength, uint32_t *digest)
{
	uint32_t 	*nm32;
	uint32_t 	i, j;
	uint32_t 	AA, BB, CC, DD;
	register uint32_t A, B, C, D;
	uint32_t	mlength;
	uint32_t 	X[16];
	prefetchw(nm8);
	nm32 = (uint32_t*)nm8;
	AA = 0x67452301; BB = 0xefcdab89; CC = 0x98badcfe; DD = 0x10325476;
	prefetch(&X[0]);
	prefetchw(&X[0]);
	for(i = 0; i < (nmlength/64); ++i)
	{
		// move block from arbitrary memory location to stack. (could gain speed changing this)
		for(j = 0, mlength = i*64; j < 16; ++j)
		{
			X[j] = nm32[mlength+j];
		}
		A = AA;
		B = BB;
		C = CC;
		D = DD;
		/// 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);

		AA = AA + A;
		BB = BB + B;
		CC = CC + C;
		DD = DD + D;
	}

	// release allocated memory.
	//free(nm8);

	digest[0] = AA;
	digest[1] = BB;
	digest[2] = CC;
	digest[3] = DD;
	return;
}

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


uint8_t  g_sequence_wmsb;				/// current sequence working most significant byte
uint32_t g_sequence_wlen;				/// current sequence working length
uint32_t g_sequence_lock;
uint8_t *g_charset_map;
uint32_t g_charset_len;

uint32_t g_threads_stop;
uint32_t g_thread_active_cnt;

uint32_t g_digest_target[4];
uint32_t g_digest_target_min;
uint32_t g_digest_target_max;
clock_t g_digest_stats_time;
uint32_t g_digest_stats_count;
uint8_t *g_digest_stats_message;

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(uint32_t arg)
{
	/// create md5 buffer large enough for one block.
	uint8_t md5buffer[64];
	uint32_t *md5buffer32 = (uint32_t*)&md5buffer[0];

	uint8_t *sequence;
	uint32_t  sequence_wlen, x;
	uint8_t sequence_wmsb;

	uint32_t digest[4];
	uint32_t loop = 0;

	time_t start;

	start = clock();

	for(x = 0; x < 64; ++x)
	{
		md5buffer[x] = 0;
	}

	sequence = (uint8_t*)malloc(g_digest_target_max);
	while((g_sequence_wlen) <= g_digest_target_max)
	{
		if(lock(arg, &g_sequence_lock) == 0)
		{
			/// negotiation between threads
			if((g_sequence_wmsb+1) >= g_charset_len)
			{
				sequence_wmsb = g_sequence_wmsb = 0;
				sequence_wlen = ++g_sequence_wlen;
			}else{
				sequence_wmsb = ++g_sequence_wmsb;
				sequence_wlen = g_sequence_wlen;
			}
			g_sequence_lock = 0;

			/// prepare sequence
			for(x = 0; x < sequence_wlen; ++x)
			{
				sequence[x] = 0;
			}
			sequence[sequence_wlen-1] = sequence_wmsb;

			/// prepare md5buffer
			for(x = 0; x < g_digest_target_max; ++x)
			{
				md5buffer[x] = 0;
			}

			md5buffer[sequence_wlen-1] = sequence_wmsb;
			md5buffer[sequence_wlen] = 0x80;
			md5buffer32[15] = 0;
			md5buffer32[14] = sequence_wlen * 8;
			/// write message, hash message, advance sequence
			while(1)
			{
				for(x = 0; x < (sequence_wlen); ++x)
				{
					md5buffer[x] = g_charset_map[sequence[x]];
				}
				//printf("cs:");
				//for(x = 0; x < sequence_wlen; ++x)
				//{
				//	printf("%02x", sequence[x]);
				//}
				//printf("\n");
				md5_hash(&md5buffer[0], 64, &digest[0]);
				if(	(digest[0] == g_digest_target[0]) && (digest[1] == g_digest_target[1]) &&
					(digest[2] == g_digest_target[2]) && (digest[3] == g_digest_target[3]) )
				{
					g_threads_stop = 1;
					printf(_C(0) "h/s:" _C(36) "%f\n", (float)loop / (float)((clock() - start) / CLOCKS_PER_SEC));
					printf(_C(0) "(" _C(36) "%x" _C(0) ") I found a collision!\n", arg);
					printf(_C(0) "message:");
					for(x = 0; x < sequence_wlen; ++x)
					{
						printf(_DC "%c", 30+(x%7), g_charset_map[sequence[x]]);
					}
					printf("\n" _C(0));
				}
				++loop;
				if(sequence_next(sequence, g_charset_len, sequence_wlen - 1) == 0)
				{
					break;
				}
				if(g_threads_stop == 1)
				{
					asm("lock decl %0" : "=m" (g_thread_active_cnt));
					printf("g_thread_active_cnt:%x\n", g_thread_active_cnt);
					return 0;
				}
				
			}
			/// advance len or sequence most significant, or both.
		}
	}
	free(sequence);

	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 SEQUENCE_LENGTH, uint32_t THREAD_COUNT)
{	
	uint32_t x;
	pthread_t thread_workers[THREAD_COUNT];
	g_sequence_wlen = 2;
	g_sequence_wmsb = 0;
	g_charset_len = CHARSET_LENGTH;
	g_thread_active_cnt = THREAD_COUNT;
	for(x = 0; x < THREAD_COUNT; ++x)
	{
		pthread_create(&thread_workers[x], NULL, (void*(*)(void*))&threaded_worker, (void*)(x+1));
	}
	// keep waiting until all have returned.
	while(g_thread_active_cnt != 0)
	{
		sleep(1);
	}
	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 < 3 )
	{
		printf("[hash] [min-search-length] [max-search-length]\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_digest_target_min = atoi( argv[2] );
	g_digest_target_max = atoi( argv[3] );
	/// <size of char map>, <maximum message size we are searching for>, <number of threads to use>
	md5_search_core(sizeof(__charset), g_digest_target_max, 1);
	return 1;
}