/* Copyright 2006 Neil Edelman, distributed under the terms of the
 GNU General Public License, see copying.txt */

/* this is a socket programme for Windows that goes with the KillSpam */

#include <winsock2.h>
#include <stdlib.h> /* malloc(), free() */
#include <stdio.h> /* printf(), fprintf(), perror() */
#include <string.h> /* strncmp() */
#include "net.h"

#define WSA_VERSION 2

struct tagNet {
	enum {
		WSA_DLL = 1,
		FOUND_HOST = 2,
		CONNECTED = 4
	} flags;
	int sd;
	unsigned long nlAddr;
	unsigned short nsPort;
	unsigned int bufSize;
	char *buffer;
};

Net *create_net(unsigned int bufSize) {
	WSADATA wsaData;
	Net *net;
	int ret;

	/* require a buffer with at least a single character in it plus a null */
	if(bufSize < 2) return 0;
	/* allocate and initialize the new structure */
	if(!(net = malloc(sizeof(Net) + bufSize))) return 0;
	net->flags = 0;
	net->sd = INVALID_SOCKET;
	net->nlAddr = INADDR_NONE;
	net->bufSize = bufSize;
	net->buffer = (char *)(net + 1);
	/* load Winsock 2.0 DLL */
	if((ret = WSAStartup(WSA_VERSION, &wsaData)) != 0) {
		printf("Error initializing windows socket library (%d.)\n", ret);
		destroy_net(net);
		return 0;
	} else {
		printf("Windows socket library loaded.\n");
		net->flags |= WSA_DLL;
	}
	return net;
}

void destroy_net(Net *net) {
	if(!net) return;
	if(net->flags & CONNECTED) disconnect_net(net);
	if(net->flags & WSA_DLL) {
		printf("Releasing Windows sockets library.\n");
		/* unload Winsock DLL */
		if(WSACleanup() != 0) {
			printf("Failed shutting down socket library (%d.)\n", WSAGetLastError());
		}
	}
	free(net);
}

int set_net_host(Net *net, char *hostName, unsigned short port) {
	unsigned long nlAddr;

	if(!net || !net->flags & WSA_DLL || !hostName) return 0;
	/* first check if this is a valid dotted IP address */
	if((nlAddr = inet_addr(hostName)) == INADDR_NONE) {
		/* make sure it wasn't specifically INADDR_NONE */
		if(strncmp(hostName, "255.255.255.255", 15) != 0) {
			struct hostent *ent;

			printf("Looking up %s.\n", hostName);
			/* lastly try to look up the host address with DNS */
			if(!(ent = gethostbyname(hostName))) {
				/* herror("Looking up host address"); <- not in WSA */
				printf("Couldn't find address %s (%d.)\n", hostName, WSAGetLastError());
			} else {
				struct in_addr inadr;

				nlAddr = inadr.s_addr = *((unsigned long *)ent->h_addr_list[0]);
				printf("Found I.P. %s.\n", inet_ntoa(inadr));
			}
		}
	}
	if(nlAddr == INADDR_NONE) return 0;
	net->nlAddr = nlAddr;
	/* set the port using network byte order */
	net->nsPort = htons(port);
	net->flags |= FOUND_HOST;
	return 1;
}

int connect_net(Net *net) {
	struct sockaddr_in inaddr;

	if(!net
	|| !(net->flags & FOUND_HOST)
	|| (net->flags & CONNECTED)) return 0;
	/* create the socket */
	if((net->sd = socket(PF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) {
		perror("Unable to create socket");
		/* printf("Unable to create socket (%u.)\n", WSAGetLastError()); */
		return 0;
	} else {
		printf("Opened socket.\n");
	}
	/* connect to the host */
	memset(&inaddr, 0, sizeof(inaddr));
	inaddr.sin_family = AF_INET;
	inaddr.sin_port = net->nsPort;
	inaddr.sin_addr.s_addr = net->nlAddr;
	if(connect(net->sd, (const struct sockaddr *)&inaddr, sizeof(inaddr)) != 0) {
		perror("Failed connecting to host");
		/* printf("Failed connecting to host (%u.)\n", WSAGetLastError()); */
		disconnect_net(net);
		return 0;
	}
	else {
		net->flags |= CONNECTED;
		printf("Connected on port %u.\n", ntohs(inaddr.sin_port));
	}
	return 1;
}

void disconnect_net(Net *net) {
	if(!net) return;
	if((net->flags & CONNECTED)) {
		int len;

		printf("Disconnecting.\n");
		/* initiate sending shutdown */
		if(shutdown(net->sd, SD_SEND) != 0) {
			perror("Error shutting down outgoing connection");
			/* printf("Error shutting down outgoing connection (%u.)\n", WSAGetLastError()); */
		}
		/* loop on recv until all data is received */
		while((len = recv(net->sd, net->buffer, net->bufSize, 0))) {
			if(len == SOCKET_ERROR) {
				perror("Error receiving left-over data");
				/* printf("Error receiving left-over data (%u.)\n", WSAGetLastError()); */
				break;
			}
		}
		/* shut down reading on this connection */
		if(shutdown(net->sd, SD_RECEIVE) != 0) {
			perror("Error shutting down incoming connection");
			/*printf("Error shutting down incoming connection (%u.)\n", WSAGetLastError()); */
		}
		net->flags &= ~CONNECTED;
	}
	if(net->sd != INVALID_SOCKET) {
		printf("Closing socket.\n");
		/* release the socket descriptor */
		if(closesocket(net->sd) != 0) {
			perror("Unable to close socket");
			/* printf("Unable to close socket (%u.)\n", WSAGetLastError()); */
		} else {
			net->sd = INVALID_SOCKET;
		}
	}
}

int send_net_string(Net *net, char *str, int len) {
	int ret;

	if(!net || !(net->flags & CONNECTED)) return 0;
	for( ; ; ) {
		if((ret = send(net->sd, str, len, 0)) == SOCKET_ERROR) {
			perror("Error sending data");
			/* printf("Error sending data (%u.)\n", WSAGetLastError()); */
			return 0;
		}
		str += ret;
		if((len -= ret) <= 0) break;
	}
	return 1;
}

int read_net_string(Net *net, char **bufAddr, unsigned int *lenAddr, long usec) {
	struct timeval tv;
	struct fd_set readfds;
	int len;

	if(!net
	|| !(net->flags & CONNECTED)
	|| usec < 0) return 0;
	tv.tv_sec = 0;
	tv.tv_usec = usec;
	/* initialize to an empty buffer */
	*net->buffer = 0;
	*bufAddr = net->buffer;
	*lenAddr = 0;
	/* since this is Windows, it doesn't quite work like it should */
	/* fdStdin = fileno(stdin); */
	FD_ZERO(&readfds);
	/* FD_SET(fdStdin, &readfds); */
	FD_SET(net->sd, &readfds);
	/* fdMax = (net->sd > fdStdin) ? net->sd : fdStdin; */
	if((len = select(net->sd, &readfds, 0, 0, &tv)) == SOCKET_ERROR) {
		perror("Error checking input status");
		return 0;
	}
	/* ioctlsocket(net->sd, FIONREAD, &amount);
	if(!amount) return 0; */
	if(FD_ISSET(net->sd, &readfds)) {
		if((len = recv(net->sd, net->buffer, net->bufSize - 1, 0)) == SOCKET_ERROR) {
			perror("Error receiving data");
			/* printf("Error receiving data (%u.)\n", WSAGetLastError()); */
			return 0;
		} else if(!len) {
			printf("Connection terminated by host.\n");
			/* successfully reading nothing indicates disconnection */
			net->flags &= ~CONNECTED;
			return 0;
		}
		/* this will never happen */
		if(len > net->bufSize - 1) len = net->bufSize - 1;
		/* terminate the buffer with a null character */
		net->buffer[len] = 0;
		*lenAddr = len;
	}
	return 1;
}