#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#ifdef WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <unistd.h>
#include <netinet/in.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <netdb.h>
#include <fcntl.h>
#include <arpa/inet.h>
#endif

#include "pscan.h"

typedef struct _net_instance
{
	struct _net_instance *next;
	unsigned long ip;
	unsigned short port;
	unsigned long time_start;
	int sock;
} net_instance;

typedef struct _net_data
{
	rescb cb;
	net_instance *instances;
	unsigned long wait_time;
} net_data;

static net_data *nd;
void add_next_port();

static void instance_remove(net_instance *ni)
{
	if (nd->instances == ni)
		nd->instances = ni->next;
	else
	{
		net_instance *t = nd->instances;
		while (t->next != ni) t = t->next;
		t->next = ni->next;
	}
#ifdef WIN32
	closesocket(ni->sock);
#else
	close(ni->sock);
#endif
	free(ni);
}

static int check_connection_status(int sock, int mode)
{
#ifdef WIN32
	fd_set fds;
	fd_set *fds_ptr_1 = NULL, *fds_ptr_2 = NULL;
	struct timeval	tv;

	memset(&tv, 0, sizeof(struct timeval));
	FD_ZERO(&fds);
	FD_SET(sock, &fds);
	if (mode == 1) 
		fds_ptr_1 = &fds;
	else 
		fds_ptr_2 = &fds;
	if (select(0, NULL, fds_ptr_1, fds_ptr_2, &tv) == 1) 
		return 1;
	else return 0;
#else
	struct pollfd pfd;
	short mevent;

	pfd.fd = sock;
	pfd.revents = 0;
	if (mode == 1) 
		mevent = POLLOUT;
	else // mode == 0
		mevent = POLLHUP;
	pfd.events = mevent;
	poll(&pfd, 1, 0);
	if (pfd.revents & mevent)
		return 1;
	else 
		return 0;
#endif
}

static net_instance *check_instance(net_instance *ni, unsigned long time)
{
	net_instance *n = ni->next;

	if (check_connection_status(ni->sock, 1) == 1)
	{
		// port open
		nd->cb(OPEN, ni->port);	// here we call on_response()
		add_next_port();
		instance_remove(ni);
	}
	else if (check_connection_status(ni->sock, 0) == 1)
	{
		// port closed
		nd->cb(CLOSED, ni->port);	
		add_next_port();
		instance_remove(ni);
	}
	else
	{
		// not known yet
		// check for timeout
		if (nd->wait_time != 0 && (nd->wait_time + ni->time_start) < time)
		{
			// timed-out
			nd->cb(TIMEDOUT, ni->port);
			add_next_port();
			instance_remove(ni);
		}
	}

	return n;
}

void net_init()
{
#ifdef WIN32
	WSADATA wsadata;
	WSAStartup(MAKEWORD(2, 2), &wsadata);
#endif
}

void net_end()
{
#ifdef WIN32
	WSACleanup();  
#endif
}

void struct_init(rescb cb, unsigned long wait_time)
{
	nd = (net_data *) malloc(sizeof(net_data));
	memset(nd, 0, sizeof(net_data));
	nd->cb = cb;
	nd->wait_time = wait_time;
}

void net_uninit()
{
	while (nd->instances) 
		instance_remove(nd->instances);
	free(nd);
}

int net_check(unsigned long time)
{
	net_instance *ni = nd->instances;

	if (!ni) 
		return 0;
	while (ni) 
		ni = check_instance(ni, time);
	return 1;
}

int net_add(unsigned long ip, unsigned short port, unsigned long time)
{
	struct sockaddr_in addr;
#ifdef WIN32
	unsigned long mode = 1;
#else
	int flags;
#endif

	net_instance *ni = (net_instance *) malloc(sizeof(net_instance));
	ni->ip = ip;
	ni->port = port;
	ni->sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (ni->sock < 1) 
	{
		free(ni);
		return 0;
	}
#ifdef WIN32
	if (ioctlsocket(ni->sock, FIONBIO, &mode) == SOCKET_ERROR)
#else
	flags = fcntl(ni->sock, F_GETFL, 0);
	if (flags < 0)
	{
		free(ni);
		return 0;
	}
	if (fcntl(ni->sock, F_SETFL, flags | O_NONBLOCK) == -1)
#endif
	{
		free(ni);
		return 0;
	}

	addr.sin_addr.s_addr = ip;
	addr.sin_port = htons(port);
	addr.sin_family = AF_INET;
	connect(ni->sock, (struct sockaddr *) &addr, sizeof(struct sockaddr_in));
	ni->time_start = time;

	ni->next = nd->instances;
	nd->instances = ni;

	return 1;
}


char *net_get_aip(unsigned long ip)
{
// some compilers complain about this line of code
//	return inet_ntoa(*((struct in_addr *)&ip));	
// so, this is a workaroud...

struct in_addr *ia = (struct in_addr *) &ip;
static struct in_addr ib;
ib = *ia;
return inet_ntoa(ib);
}


unsigned long net_get_iip(char *aip)
{
	return inet_addr(aip);
}


unsigned long ResolveHostname(char *hostname)
{
	struct addrinfo *addressInfo;
	unsigned long ResolvedIp;

	getaddrinfo (hostname, NULL, NULL, &addressInfo);
	ResolvedIp = ((struct sockaddr_in *)(addressInfo->ai_addr))->sin_addr.s_addr;
			
	freeaddrinfo(addressInfo);
	
	return ResolvedIp;
}
