/* Public Domain */

#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netdb.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>
#include <openssl/sha.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/ssl.h>
#include <openssl/err.h>

#define MOT_KEY_SHA1_HASH "\x1d\x3f\xb6\x62\x79\x4d\x8c\x70\xfb\x57\xb4\xcb\x49\x2e\x27\xf6\x6f\x15\x2e\x4f"
#define MOT_KEY_EXPONENT 0x10001
#define MOT_KEY_SIZE 2048

#define SERVER "droid-developers.org"
#define PORT 448

#ifdef USE_BITS
static int max_bit_count = 0;
#endif

char diff[20];

struct sha1_hash {
	unsigned char* hex;
	short bits[160];
	int bit_count;
};

struct keys_pool {
	RSA* key;
	void* next;
};

/* Function for init SHA1 hash fast-access structure
 * ------------------------------------------------------------------
 *  Take empty structure with fullfilled only string of hexadecimal 
 *  chars and fill other fields, e.g. bits[] array and bits_count
 *  accumulator
 */
void sha1_hash_init(struct sha1_hash *hash)
{
	int i, pos = 0;

	unsigned int *input= (unsigned int *)hash->hex;
	hash->bit_count = 0;

	for(i = 31; i >= 0; i--)
	{
		if (((*input >> i) & 1)) 
		{
			hash->bits[pos] = 1;
			hash->bit_count++;
		}
		else hash->bits[pos] = 0;
		pos++;
	}
}

/* Function for compare two sha1 hashes.
 * -------------------------------------------------------------------
 * Return 1 if hashes identical
 * Return 0 if hashes differs 
 */
int sha1_hash_compare_hex(struct sha1_hash *hash1, struct sha1_hash *hash2)
{
	int i;

	for (i = 0; hash1->hex[i]; i++ )
		if ( (hash1->hex[i] & hash2->hex[i]) != hash1->hex[i] ) {
			diff[i] = hash1->hex[i] ^ hash2->hex[i];
			printf("out %d: %x %x\n", i, hash1->hex[i], hash2->hex[i]);
			printf("differs in %d byte: %s\n", i + 1, diff);
			return 0;
		}
	return 1;
}

unsigned char *DER_encode_RSA_public(RSA *rsa, int *len) 
{
	unsigned char *next;
	
	next = NULL;
	*len = i2d_RSAPublicKey(rsa, &next);
	return next;
}

unsigned char *DER_encode_RSA_private(RSA *rsa, int *len) 
{
	unsigned char *next;
	
	next = NULL;
	*len = i2d_RSAPrivateKey(rsa, &next);
	return next;
}


// TODO: Make keys generation as different thread, 
// which can generate pool of the RSA keys.
void generate_keys(struct keys_pool *pool)
{
	pool->key = RSA_generate_key(MOT_KEY_EXPONENT, MOT_KEY_SIZE, NULL, NULL);
}

// Simple structure to keep track of the handle, and
// of what needs to be freed later.
typedef struct {
	int socket;
	SSL *sslHandle;
	SSL_CTX *sslContext;
} connection;

// Establish a regular tcp connection
int tcpConnect ()
{
	int error, handle;
	struct hostent *host;
	struct sockaddr_in server;

	host = gethostbyname (SERVER);
	handle = socket (AF_INET, SOCK_STREAM, 0);
	if (handle == -1)
	{
		perror ("Socket");
		handle = 0;
	}
	else
	{
		server.sin_family = AF_INET;
		server.sin_port = htons (PORT);
		server.sin_addr = *((struct in_addr *) host->h_addr);
		bzero (&(server.sin_zero), 8);
		
		error = connect (handle, (struct sockaddr *) &server, sizeof (struct sockaddr));

		if (error == -1)
		{
			perror ("Connect");
			handle = 0;
		}
	}

	return handle;
}

// Establish a connection using an SSL layer
connection *sslConnect (void)
{
	connection *c;

	c = malloc (sizeof (connection));
	c->sslHandle = NULL;
	c->sslContext = NULL;

	c->socket = tcpConnect ();
	if (c->socket)
	{
		// Register the error strings for libcrypto & libssl
		SSL_load_error_strings ();
		// Register the available ciphers and digests
		SSL_library_init ();
		// New context saying we are a client, and using SSL 2 or 3
		c->sslContext = SSL_CTX_new (SSLv23_client_method ());
		if (c->sslContext == NULL)
			ERR_print_errors_fp (stderr);

		// Create an SSL struct for the connection
		c->sslHandle = SSL_new (c->sslContext);
		if (c->sslHandle == NULL)
			ERR_print_errors_fp (stderr);

		// Connect the SSL struct to our connection
		if (!SSL_set_fd (c->sslHandle, c->socket))
			ERR_print_errors_fp (stderr);

		// Initiate SSL handshake
		if (SSL_connect (c->sslHandle) != 1)
			ERR_print_errors_fp (stderr);
	}
	else
	{
		perror ("Connect failed");
	}
	return c;
}

// Disconnect & free connection struct
void sslDisconnect (connection *c)
{
	if (c->socket)
		close (c->socket);
	
	if (c->sslHandle)
	{
		SSL_shutdown (c->sslHandle);
		SSL_free (c->sslHandle);
	}
	
	if (c->sslContext)
		SSL_CTX_free (c->sslContext);
	
	free (c);
}

// Read all available text from the connection
char *sslRead (connection *c)
{
	const int readSize = 1024;
	char *rc = NULL;
	int received, count = 0;
	char buffer[1024];

	if (c)
	{
		while (1)
		{
			if (!rc)
				rc = malloc (readSize * sizeof (char) + 1);
			else
				rc = realloc (rc, (count + 1) * readSize * sizeof (char) + 1);
			received = SSL_read (c->sslHandle, buffer, readSize);
			buffer[received] = '\0';
			
			if (received > 0)
				strcat (rc, buffer);
				
			if (received < readSize)
				break;
				count++;
			}
		}
	return rc;
}

// Write text to the connection
void sslWrite (connection *c, char *text)
{
	if (c)
	SSL_write (c->sslHandle, text, strlen (text));
}

#ifdef STAT
void client_register(void)
{
	connection *c;
	char *response =  NULL;
	char *ssl_data =  NULL;

	c = sslConnect();
	strcat(ssl_data, "CLIENT_REGISTER");
	sslWrite(c, ssl_data);
	response = sslRead (c);
	printf("%s\n", response);
	sslDisconnect (c);
	free (response);
}
#endif

int main(void)
{
	RSA* key;
	const unsigned char* key_public = NULL;
	const unsigned char* key_private = NULL;
	char* ssl_data = NULL;
	struct sha1_hash *origin;
	struct sha1_hash *hash;
	struct keys_pool *pool;
	struct diff_buffer *diff;
	connection *c;
	char *response =  NULL;

	int key_size = 0;
	int n = 0;

#ifdef STAT
	client_register();
#endif

	hash = malloc(sizeof(*hash));
	origin = malloc(sizeof(*origin));
	pool = malloc(sizeof(*pool));
	diff =  malloc(sizeof(diff));
	origin->hex = MOT_KEY_SHA1_HASH;
	sha1_hash_init(origin);
	
	while (1) {
		
		//generate_keys(pool);
		pool->key =  RSA_generate_key(MOT_KEY_SIZE,MOT_KEY_EXPONENT, NULL, NULL);
		key_size =  RSA_size(pool->key);
		key_public = DER_encode_RSA_public(pool->key, &key_size);
		hash->hex = SHA1(key_public, key_size, NULL);

		if (sha1_hash_compare_hex(origin, hash) == 1)
		{
			printf("key found!\n");
			RSA_print_fp(stdout, key, 3);
			key_private = DER_encode_RSA_private(pool->key, &key_size);

			c = sslConnect ();
			strcat(ssl_data, "KEY_FOUND");
			strcat(ssl_data, key_public);
			strcat(ssl_data, " | ");
			strcat(ssl_data, key_private);
			sslWrite (c, ssl_data);
			response = sslRead (c);
			printf ("%s\n", response);
			sslDisconnect (c);
			free (response);
			return EXIT_SUCCESS;
		}
		RSA_free(pool->key);
		n++;
		if ( !(n % 10) ) {
			printf("%d keys done!\n", n);
		}
	}
	return EXIT_SUCCESS;
}