#include <xprobe.h>
#include "ttl_calc.h"
#include "util.h"
#define _XPROBE_MODULE
#include "xplib.h"
#include "xprobe_module.h"
#include "interface.h"

extern Interface *ui;

int TTL_Mod::init (void) {
	
	xprobe_mdebug(XPROBE_DEBUG_MODULES, "[TTL_Mod]: Initialized\n");
return OK;

}

int TTL_Mod::exec (Target *Tgt, OS_Matrix *osmtx) {
	int ttldistance;

	if ((ttldistance = get_ttl_distance (Tgt)) > 0) {
		Tgt->set_distance(ttldistance);
		osmtx->add_result (get_id(), 1, XPROBE_MATCH_YES);
		xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] TTL distance to target is: %d\n", get_name(), Tgt->get_distance());
	} else 
		xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] TTL distance calculation failed\n", get_name());
	
	return OK;

}

int TTL_Mod::fini (void) {
	xprobe_debug(XPROBE_DEBUG_MODULES, "%s module has been deinitilized\n", get_name());
return OK;
}


int TTL_Mod::parse_keyword(int os_id, const char *keyword, const char *value){
    xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] Parsing for %i : %s  = %s\n", 
				get_name(), os_id, keyword, value);
return OK;
}



int ttl_mod_init(Xprobe_Module_Hdlr *pt) {

	TTL_Mod *ttl_mod = new TTL_Mod;

	xprobe_mdebug(XPROBE_DEBUG_MODULES, "Initializing the TTL precalculation module\n");
	/* register module and keywords */
	pt->register_module(ttl_mod);
return(OK);
}

int TTL_Mod::get_ttl_distance(Target *tgt) {

	struct in_addr remote = tgt->get_addr();
	struct in_addr local = tgt->get_interface_addr();
	struct icmp_hdr *icmph;
	struct tcphdr *tcph;
	struct ip *iph;
	struct timeval timeo;
    int retval, ttl_distance = 1, failures=0, done = 0;
	u_char *buf, *udppayload;
	bool tcptest=true;
	time_t timestart;

	if ( (buf = (u_char *) malloc (BUFSIZE)) == NULL ) {
		ui->error("get_ttl_distance: failed to malloc()\n");
		return (-1);
	}
	if ( (udppayload = (u_char *) malloc (sizeof (struct DNSHEADER))) == NULL ) {
		ui->error("get_ttl_distance: failed to malloc()\n");
		return (-1);
	}
	memset (buf, 0, BUFSIZE);
	memset (udppayload, 0, sizeof(struct DNSHEADER));
	build_DNS_reply(udppayload);
    TCP tcp(inet_ntoa(remote));
    UDP udp(inet_ntoa(remote));
    IP sn(inet_ntoa(local), 123);
	sn.init_device (tgt->get_interface(), 0, 1024);
	tcp.set_src(inet_ntoa(local));
	tcp.set_id(getrandom(U_INTMAX));
	tcp.set_fragoff(IP_DF);

    tcp.set_dstport (tgt->get_port(IPPROTO_TCP, XPROBE_TARGETP_OPEN));
	tcp.set_srcport (getrandom(U_SHORTMAX));
	tcp.set_flags (TCPPACKETFLAGS);
	tcp.set_seq (getrandom(U_INTMAX));
	tcp.set_ack (0);

	udp.set_src (inet_ntoa(local));
	udp.set_id (getrandom(U_INTMAX));
	udp.set_fragoff(IP_DF);
	udp.set_srcport (53);
	udp.set_dstport (tgt->get_port(IPPROTO_UDP, XPROBE_TARGETP_CLOSED));
	timeo.tv_usec = 0;
	timeo.tv_sec = TIMEOUT;
	tcp.timeout(timeo);
	udp.timeout(timeo);
	sn.timeout(timeo);

	while (1) {
		if (tcptest) {
			tcp.set_ttl(ttl_distance);
			tcp.sendpack("");
		} else {
			udp.set_ttl(ttl_distance);
			udp.sendpack((char *)udppayload, sizeof(struct DNSHEADER));	
		}
		timestart = time(NULL);
		done = 0;
		while (!done) {
			retval = sn.sniffpack(buf, BUFSIZE);	
			if (retval > 0 ) {
				if (sn.get_proto() == IPPROTO_ICMP && 
					retval > (int) (sizeof(struct ip) + sizeof (struct icmp_hdr))) {
						icmph = (struct icmp_hdr *) buf;
						iph = (struct ip *) (buf + sizeof (struct icmp_hdr));
						if (icmph->type == ICMP_TIMXCEED &&
							icmph->code == ICMP_TIMXCEED_INTRANS &&
							iph->ip_src.s_addr == local.s_addr &&
							iph->ip_dst.s_addr == remote.s_addr) {
								ttl_distance++;
								break;
						}else
						if (tcptest == false &&
							icmph->type == ICMP_UNREACH &&
							icmph->code == ICMP_UNREACH_PORT &&
                        	iph->ip_src.s_addr == local.s_addr &&
                        	iph->ip_dst.s_addr == remote.s_addr) {
								return ttl_distance;
						}
				}else
				if (tcptest == true && sn.get_proto() == IPPROTO_TCP &&
					retval >= (int) sizeof (struct tcphdr) ) {
						tcph = (struct tcphdr *) buf;
						if (ntohl (tcph->th_ack) == tcp.get_seq()+1) {
							return ttl_distance;
						}
				}
			}
			if ( (time(NULL) - timestart) > TIMEOUT ) {
				xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] Timed out\n", get_name());
				retval = 0;
			}
			/* this is ugly need something more efficient and nice */
			if (retval < 1) {
				if (failures == 2) {
					xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] UDP test failed also\n", get_name());
					return (0);
				}
				if (tcptest == true && failures == 1) {
					xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] Switched to UDP test\n", get_name());
					tcptest = false;
					failures++;
					done = 1;
				}
				if (tcptest == true) {
					xprobe_debug(XPROBE_DEBUG_MODULES, "[%s] Switching TCP test, now sending to closed port\n", get_name());
					tcp.set_dstport(tgt->get_port(IPPROTO_TCP, XPROBE_TARGETP_CLOSED));
					failures++;
					done = 1;
				}
			}
		} /* inside while() */
	} /* outside while() */

return 0;
}
