/* libnebula.c
 *
 * Copyright (C) 2009 Tillmann Werner <tillmann.werner@gmx.de>
 *
 * This file is free software; as a special exception the author gives
 * unlimited permission to copy and/or distribute it, with or without
 * modifications, as long as this notice is preserved.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY, to the extent permitted by law; without even the
 * implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 *
 * Description:
 *   client library for data submissions to a nebula server
 *
 *
 *  $Id$
 */

#define _GNU_SOURCE 1

#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <zlib.h>

#include <nebula.h>

#include "hmac.h"
#include "md5.h"
#include "sha512.h"
#include "util.h"

#include <ctype.h>


const char *nebula_strerr(nebula *n) {
	if (!n) return NULL;

	if (n->nerrno == NERR_PROXYERRNO)
		return strerror(errno);
	else
		return nebula_errstr[n->nerrno];
}


nebula *nebula_new(void) {
	nebula	*n = NULL;

	if ((n = calloc(1, sizeof(nebula))) == NULL) {
		return NULL;
	}

	return n;
}


int nebula_init(nebula *n) {
	// initialize HMAC setup
	size_t keylen = 0;

	if (n->secret) {
		if (strlen(n->secret) <= HMAC_BLOCK_SIZE) {
			n->hmac.key = n->secret;
			n->secret = NULL;
		} else {
			if ((n->hmac.key = mem_sha512sum((u_char *) n->secret, strlen(n->secret))) == NULL) {
				n->nerrno = NERR_HMACINIT;
				return 0;
			}
		}
		keylen = strlen(n->hmac.key);
	}
	memset(n->hmac.k_ipad, 0x36, HMAC_BLOCK_SIZE);
	memset(n->hmac.k_opad, 0x5c, HMAC_BLOCK_SIZE);

	int i;
        if (n->hmac.key) for (i=0; i<strlen(n->hmac.key); i++) {
                n->hmac.k_ipad[i] ^= n->hmac.key[i];
                n->hmac.k_opad[i] ^= n->hmac.key[i];
        }

	return 1;
}


int nebula_authenticate(nebula *n) {
	u_int32_t	nonce;
	ssize_t		bytes_read, total_bytes;
	u_char		*cbuf;
	char		*sha512sum;
	size_t		keylen;

	keylen = n->hmac.key ? strlen(n->hmac.key) : 0;

	for (bytes_read = 1, total_bytes = 0; bytes_read && total_bytes < 4; total_bytes += bytes_read)
		bytes_read = recv(n->sockfd, &nonce + total_bytes, 4, 0);

	if (bytes_read < 0) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	nonce = ntohl(nonce);

	// hash secret with nonce
	if ((cbuf = malloc(keylen + 4)) == NULL) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	memcpy(cbuf, n->hmac.key, keylen);
	memcpy(cbuf + keylen, &nonce, 4);

	if ((sha512sum = mem_sha512sum(cbuf, keylen + 4)) == NULL) {
		n->nerrno = NERR_HMACHASH;
		return 0;
	}
	free(cbuf);

	// send hashed secret or nonce
	if (send(n->sockfd, sha512sum, 128, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	free(sha512sum);


	return 1;
}


int nebula_togglepersist(nebula *n, int pswitch) {
	u_int32_t msg_type;

	if (!nebula_authenticate(n)) return 0;

	msg_type = pswitch > 0 ? htonl(1) : htonl(2);
	if (send(n->sockfd, &msg_type, 4, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	return 1;
}


int nebula_connect(nebula *n, struct in_addr host, u_int16_t port, int flags) {
	switch (flags) {
	case 0:
		break;
	case NSCK_PERSIST:
		break;
	case NSCK_NONBLCK:
		break;
	default:
		n->nerrno = NERR_INVALIDFLAGS;
		return 0;
	}

	if ((n->sockfd = socket(PF_INET, SOCK_STREAM, 0)) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	n->server.sin_family	= PF_INET;
	n->server.sin_addr	= host;
	n->server.sin_port	= htons(port);

	if (connect(n->sockfd, &n->server, sizeof(n->server)) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	if (flags & NSCK_PERSIST) {
		// request persistent mode from server
		if (nebula_togglepersist(n, 1) == 0) return -1;
	}

	return n->sockfd;
}


int nebula_send(nebula *n, u_char protocol, u_int16_t port, u_char *data, size_t len) {
	u_char		*cbuf;
	char		*md5sum, *sha512sum, response[9];

	// authenticate
	if (!nebula_authenticate(n)) return -1;


	// send message type, data submission has type 0
	u_int32_t msg_type = 0;
	if (send(n->sockfd, &msg_type, 4, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// send md5 hash
	if ((md5sum = mem_md5sum(data, len)) == NULL) {
		n->nerrno = NERR_HMACHASH;
		return 0;
	}
	if (send(n->sockfd, md5sum, 32, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	free(md5sum);


	if (!nebula_read_line(n->sockfd, response, 9)) {
		n->nerrno = NERR_RECV;
		return 0;
	}
	if (strncmp((char *) response, "KNOWN", 5) == 0) {
		return 1;
	}
	if (strncmp((char *) response, "UNKNOWN", 7)) {
		return 1;
		n->nerrno = NERR_INVALID;
		return 0;
	}

	// send protocol
	if (send(n->sockfd, &protocol, 1, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// send port
	u_int16_t hmac_port = htons(port);
	if (send(n->sockfd, &hmac_port, 2, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// send length of uncompressed attack
	u_int32_t len_uncompressed = htonl(len);
	if (send(n->sockfd, &len_uncompressed, 4, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// compress attack
	unsigned long cbuf_len = (len * 1.1) + 12;
	if ((cbuf = calloc(1, cbuf_len)) == NULL) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	switch (compress(cbuf, &cbuf_len, data, len)) {
	case Z_OK:
		break;
	case Z_MEM_ERROR:
		n->nerrno = NERR_ZMEM;
		return 0;
	case Z_BUF_ERROR:
		n->nerrno = NERR_ZBUF;
		return 0;
	default:
		n->nerrno = NERR_ZUNKNOWN;
		return 0;
	}

	// send length of compressed attack
	u_int32_t len_compressed = htonl(cbuf_len);
	if (send(n->sockfd, &len_compressed, 4, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// send compressed attack
	int sent, bytes;
	for (sent = 0, bytes = 0; sent < cbuf_len; sent += bytes) {
		if ((bytes = send(n->sockfd, cbuf+sent, sent + BUFSIZ < cbuf_len ? BUFSIZ : cbuf_len - sent, 0)) == -1) {
			n->nerrno = NERR_PROXYERRNO;
			return 0;
		}
	}

	// append protocol and port to cattack for HMAC calculation
	if ((cbuf = realloc(cbuf, cbuf_len+3)) == NULL) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	memcpy(cbuf+cbuf_len, &protocol, 1);
	memcpy(cbuf+cbuf_len+1, &hmac_port, 2);

	if ((sha512sum = hmac(n->hmac.k_ipad, n->hmac.k_opad, &cbuf, cbuf_len+3)) == NULL) {
		n->nerrno = NERR_HMAC;
		return 0;
	}
	free(cbuf);

	// send length of HMAC
	int hmac_len = htons(strlen(sha512sum));
	if (send(n->sockfd, &hmac_len, 2, 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}

	// send HMAC
	if (send(n->sockfd, sha512sum, strlen(sha512sum), 0) == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	free(sha512sum);

	// wait for OK
	if (!nebula_read_line(n->sockfd, (char *) response, 9)) {
		n->nerrno = NERR_RECV;
		return 0;
	} else if (strlen((char *) response) != 2 || strncmp((char *) response, "OK", 2) != 0) {
		n->nerrno = NERR_INVALID;
		return 0;
	}

	return 1;
}


int nebula_disconnect(nebula *n) {
	int rv = shutdown(n->sockfd, SHUT_RDWR);

	close(n->sockfd);
	if (rv == -1) {
		n->nerrno = NERR_PROXYERRNO;
		return 0;
	}
	return 1;
}


int nebula_socket(nebula *n) {
	return n->sockfd;
}


void nebula_cleanup(nebula *n) {
	if (n) {
		if (n->secret) free(n->secret);
		free(n->hmac.key);
		free(n);
	}

	return;
}