#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <openssl/sha.h>
#include <tlmplib.h>
#include <sys/time.h>
#include <string>

using namespace std;

#include "ddosproof.h"

static unsigned long long tb0[64];
/*
	Initialise the table used to speedup util_nbzero()
*/
void util_filltb0()
{
	unsigned long long v=(unsigned long long)0xffffffffffffffffl;
	for (unsigned i=0; i<64; i++){
		tb0[i] = v;
		v >>=1;
	}
	// for (unsigned i=0; i<16; i++) printf ("tb0[%u] = %016Lx\n",i,tb0[i]);
}

/*
	Validate that a hash starts with a number of zeros (bits)
*/
bool util_nbzero (int nb0, const unsigned char *bufout)
{
	union {
		unsigned char c[8];
		unsigned long long v;
	}u;
	u.c[7] = bufout[0];
	u.c[6] = bufout[1];
	u.c[5] = bufout[2];
	u.c[4] = bufout[3];
	u.c[3] = bufout[4];
	u.c[2] = bufout[5];
	u.c[1] = bufout[6];
	u.c[0] = bufout[7];
	return tb0[nb0] > u.v;
}

bool util_validate (const char *answer, int difficulty)
{
	SHA256_CTX ctx;
	SHA256_Init(&ctx);
	SHA256_Update(&ctx, answer, strlen(answer));
	unsigned char bufout[BUFSIZ];
	SHA256_Final(bufout, &ctx);
	return util_nbzero (difficulty,bufout);
}

	
int util_getwork (int fd, char tmp[100])
{
	int ret = -1;
	unsigned char buf[8];
	if (read(fd,buf,sizeof(buf))!=sizeof(buf)){
		tlmp_error ("Read error on /dev/urandom (%s)\n",strerror(errno));
	}else{
		struct timeval tv;
		gettimeofday (&tv,NULL);
		int len = snprintf (tmp,100,"%ld%ld",tv.tv_sec,tv.tv_usec);
		sprintf (tmp+len,"%02x%02x%02x%02x%02x%02x%02x%02x"
			,buf[0],buf[1],buf[2],buf[3],buf[4],buf[5],buf[6],buf[7]);
		ret = 0;
	}
	return ret;
}
int util_makecookie (int fd, const string &secret, const char *ip, char res[100])
{
	int ret = -1;
	char tmp[100];
	if (util_getwork(fd,tmp)!=-1){
		SHA256_CTX ctx;
		SHA256_Init(&ctx);
		SHA256_Update(&ctx, tmp, strlen(tmp));
		SHA256_Update(&ctx, ip, strlen(ip));
		SHA256_Update(&ctx, secret.c_str(), secret.size());
		unsigned char bufout[BUFSIZ];
		SHA256_Final(bufout, &ctx);
		char *pt = stpcpy(res,tmp);
		*pt++ = ',';
		for (unsigned j = 0; j < SHA256_DIGEST_LENGTH; j++){
			sprintf(pt,"%02x", bufout[j]);
			pt += 2;
		}
		ret = 0;
	}
	return ret;
}

/*
	Check that a cookie is valid.
	A cookie is made of two parts split by a comma:
		A random hex sequence
		the hash of the same sequence combined with a secret
*/
bool util_validcookie(const char *cookie, const char *ip, const string &secret)
{
	bool ret = false;
	const char *pt = strchr(cookie,',');
	if (pt != NULL){
		const char *hash = pt+1;
		SHA256_CTX ctx;
		SHA256_Init(&ctx);
		SHA256_Update(&ctx, cookie, pt-cookie);
		SHA256_Update(&ctx, ip, strlen(ip));
		SHA256_Update(&ctx, secret.c_str(), secret.size());
		unsigned char bufout[BUFSIZ];
		SHA256_Final(bufout, &ctx);
		char out[SHA256_DIGEST_LENGTH*2+1];
		char *ptout = out;
		for (unsigned j = 0; j < SHA256_DIGEST_LENGTH; j++){
			sprintf(ptout,"%02x", bufout[j]);
			ptout += 2;
		}
		// printf ("Compare %s %s %d\n",hash,out,strcmp(hash,out));
		ret = strcmp(hash,out)==0;
	}
	return ret;
}

bool util_validwork (const char *answer, int difficulty)
{
	SHA256_CTX ctx;
	SHA256_Init(&ctx);
	SHA256_Update(&ctx, answer, strlen(answer));
	unsigned char bufout[BUFSIZ];
	SHA256_Final(bufout, &ctx);

	#if 0
		printf ("util_validwork: ");
		for (unsigned j = 0; j < SHA256_DIGEST_LENGTH; j++){
			printf("%02x", bufout[j]);
		}
		printf ("\n");
	#endif
	return util_nbzero (difficulty,bufout);
}