/* 
 * main.cpp 
 * Copyright (c) 2004-2006 Vlad GALU <dudu@dudu.ro> 
 *                    Andrei GAVRILOAIE <gavriloaie_andrei@yahoo.com> 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307, USA.
 */

#include <Class.h>

CFlowTree* pFlowTree=NULL;
CHostTree* pHostTree=NULL;
time_t currentTime;


int main(
	int argc, 
	char* argv[])
{
	pcap_t* pPcap = NULL;
	char  errbuff[PCAP_ERRBUF_SIZE];
	struct bpf_program filter;
	int ipLayerOffset = 0;
	time_t lastHostCleanup;
	time_t lastFlowCleanup;
	int packetsCaptured = 0;
	int i;
	pid_t child;

	pFlowTree = new CFlowTree();
	pHostTree = new CHostTree();
			
	
	if (argc != 3) {
		fprintf(stdout, "\nUsage: %s <interface> <filter>\n",argv[0]);
		return -1;
	};

#ifndef _DEBUG
	if (getppid() == 1)
		return -1;
	if ((child = fork()) < 0)
		exit(1);
	else if (child > 0) 
			exit(0);
	setsid();
	for (i = getdtablesize(); i >= 0; i--)
		close(i);
	signal(SIGCHLD, SIG_IGN); 
	signal(SIGTSTP, SIG_IGN);
	signal(SIGTTOU, SIG_IGN);
	signal(SIGTTIN, SIG_IGN);
#endif	

	dprintf ("Opening interface %s... ", argv[1]);
	if ((pPcap = pcap_open_live(argv[1], 
								SNAPLEN, 1, 10, errbuff)) == NULL) {
		syslog(LOG_LOCAL6|LOG_ALERT, 
			   "%s %d Cannot open pcap device: %s",
				__FILE__, __LINE__, errbuff);
		dprintf("\n%s %d Cannot open pcap device: %s\n",
				__FILE__, __LINE__, errbuff);
		return -1;
	};
	dprintf("done\n");

	
	dprintf("Compiling filter `%s`... ", argv[2]);
	if (pcap_compile(pPcap, &filter, argv[2], 1, 0x00000000) < 0) {
		syslog(LOG_LOCAL6|LOG_ALERT, 
			   "%s %d Cannot compile filter: %s",
				__FILE__, __LINE__, pcap_geterr(pPcap));
		dprintf("\n%s %d Cannot compile filter: %s\n",
				__FILE__, __LINE__, pcap_geterr(pPcap));
		return -1;
	};
	dprintf("done\n");

	
	dprintf("Setting filter... ");
	if (pcap_setfilter(pPcap, &filter) < 0) {
		syslog(LOG_LOCAL6|LOG_ALERT, 
			   "%s %d Cannot set compiled filter: %s",
				__FILE__, __LINE__, pcap_geterr(pPcap));
		dprintf("\n%s %d  Cannot set compiled filter: %s\n",
				__FILE__, __LINE__, pcap_geterr(pPcap));
		return -1;
	};
	dprintf("done\n");

	
	dprintf("Detecting IP layer offset in bytes... ");
	ipLayerOffset = getIPLayerOffset(pPcap);
	if (ipLayerOffset < 0) {
		syslog(LOG_LOCAL6|LOG_ALERT,
			   "%s %d Unknown link type",
				__FILE__, __LINE__);
		dprintf("\n%s %d Unknown link type\n",
				__FILE__, __LINE__);
		return -1;
	}
	dprintf("%d done\n", ipLayerOffset);


	lastFlowCleanup = time(NULL);
	lastHostCleanup = time(NULL);
	
	while (true) {
		if (pcap_dispatch(pPcap, -1, handlePacket,
			 (u_char *)&ipLayerOffset) < 0) {
			syslog(LOG_LOCAL6|LOG_ALERT, 
				   "%s %d pcap_dispatch failed with error: %s",
					__FILE__, __LINE__, pcap_geterr(pPcap));
			dprintf("%s %d pcap_dispatch failed with error: %s",
					__FILE__, __LINE__, pcap_geterr(pPcap));
		};
		
		if ((currentTime - lastFlowCleanup) >= 
					TIME_BETWEEN_FLOW_CLEANUPS) {
			lastFlowCleanup = currentTime;
			pFlowTree->cleanUp(cleanupFlowTest);
		};	
		
		if ((currentTime - lastHostCleanup) >= 
					TIME_BETWEEN_HOST_CLEANUPS) {
			lastHostCleanup = currentTime;
			pHostTree->cleanUp(cleanupHostTest);
		};
		
	};
	
	dprintf("Ending pcap session... ");
	pcap_close(pPcap);
	dprintf("done\n");
}

int 
getIPLayerOffset(pcap_t* pPcap) {
	int ipLayerOffset;

	switch(pcap_datalink(pPcap)) {
	case DLT_NULL:
	case DLT_LOOP:
		ipLayerOffset = 4;
		break;
	case DLT_EN10MB:
		ipLayerOffset = ETHSIZE;
		break;
	case DLT_IEEE802:
		ipLayerOffset = TRSIZE;
		break;
	case DLT_FDDI:
		ipLayerOffset = FDDISIZE;
		break;
	case DLT_SLIP:
		ipLayerOffset = SLIPSIZE;
		break;
	case DLT_PPP:
		ipLayerOffset = PPPSIZE;
		break;
	case DLT_RAW:
		ipLayerOffset = RAWSIZE;
		break;
	case DLT_LINUX_SLL:
		ipLayerOffset = ISDNSIZE;
		break;
	default:
		ipLayerOffset = -1;
		break;
	}

	return ipLayerOffset;
}

void 
handlePacket(
	u_char *pIPLayerOffsetV, 
	const pcap_pkthdr* pPcapHeader, 
	const u_char* pPacketData)
{
	struct ip& ipLayer = *((struct ip*)(pPacketData + *((int*)pIPLayerOffsetV)));
	in_addr& srcAddr = ipLayer.ip_src;
	in_addr& dstAddr = ipLayer.ip_dst;
	unsigned fragmented = ntohs(ipLayer.ip_off) & (IP_MF | IP_OFFMASK);
#if defined(AIX)
#undef ip_hl
	unsigned ip_hl = ipLayer.ip_ff.ip_fhl * 4;
#else
	unsigned ip_hl = ipLayer.ip_hl * 4;
#endif
	u_int packetLen = ntohs(ipLayer.ip_len);;
	u_char zero = 0;
	time_t t_zero = 0;
	currentTime = pPcapHeader->ts.tv_sec;

	if (!fragmented) {
		switch (ipLayer.ip_p) {
		case 6: {
			struct tcphdr& tcp = *((struct tcphdr*)(((char*)&ipLayer) + ip_hl));
			
			if (pFlowTree->insertNode
				(
				(in_addr&)ipLayer.ip_src,
				(in_addr&)ipLayer.ip_dst,
				tcp.th_sport,
				tcp.th_dport,
				tcp.th_flags,
				(u_char&)ipLayer.ip_p,
				(u_char&)ipLayer.ip_tos,
				packetLen,
				(time_t&)pPcapHeader->ts.tv_sec
				))
			{
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)pPcapHeader->ts.tv_sec, 
									  (time_t&)t_zero);
			} else
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)t_zero, 
									  (time_t&)pPcapHeader->ts.tv_sec);
			break; }
		
		case 17: {
			struct udphdr& udp = *((struct udphdr*)(((char*)&ipLayer) + ip_hl));
			
			if (pFlowTree->insertNode
				(
				(in_addr&)ipLayer.ip_src,
				(in_addr&)ipLayer.ip_dst,
				udp.uh_sport,
				udp.uh_dport,
				zero,
				(u_char&)ipLayer.ip_p,
				(u_char&)ipLayer.ip_tos,
				packetLen,
				(time_t&)pPcapHeader->ts.tv_sec
				))
			{
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)pPcapHeader->ts.tv_sec, 
									  (time_t&)t_zero);
			} else
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)t_zero, 
									  (time_t&)pPcapHeader->ts.tv_sec);
			break; }
		
		case 1: {
			struct icmp& icmp = *((struct icmp*)(((char*)&ipLayer) + ip_hl));
			u_short type = (u_short)icmp.icmp_type;
			u_short code = (u_short)icmp.icmp_code;
			
			if (pFlowTree->insertNode (
				(in_addr&)ipLayer.ip_src,
				(in_addr&)ipLayer.ip_dst,
				type, // sport
				code, // dport
				zero,
				(u_char&)ipLayer.ip_p,
				(u_char&)ipLayer.ip_tos,
				packetLen,
				(time_t&)pPcapHeader->ts.tv_sec
				))
			{
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)pPcapHeader->ts.tv_sec, 
									  (time_t&)t_zero);
			} else
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)t_zero, 
									  (time_t&)pPcapHeader->ts.tv_sec);
			break; } 
		
		default: {
			if (pFlowTree->insertNode (
				(in_addr&)ipLayer.ip_src,
				(in_addr&)ipLayer.ip_dst,
				(u_short&)t_zero,
				(u_short&)t_zero,
				zero,
				(u_char&)zero,
				(u_char&)ipLayer.ip_tos,
				packetLen,
				(time_t&)pPcapHeader->ts.tv_sec
				))
			{
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst,
									  (time_t&)pPcapHeader->ts.tv_sec,
									  (time_t&)t_zero);
			} else
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst,
									  (time_t&)t_zero,
									  (time_t&)pPcapHeader->ts.tv_sec);
			break; }
		}
	} else if (pFlowTree->insertNode(
			  (in_addr&)ipLayer.ip_src,
    		  (in_addr&)ipLayer.ip_dst,
			  (u_short&)t_zero,		// set these to 0,
			  (u_short&)t_zero,		// sport = dport = 0 => flow :)
			  zero,					// no tcp flags
			  (u_char&)zero,			// set protocol to IP
			  (u_char&)ipLayer.ip_tos,
			  packetLen,				// this remains uninitialized
			  (time_t&)pPcapHeader->ts.tv_sec))
			{
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)pPcapHeader->ts.tv_sec, 
									  (time_t&)t_zero);
			} else
				pHostTree->insertNode((in_addr&)ipLayer.ip_dst, 
									  (time_t&)t_zero, 
									  (time_t&)pPcapHeader->ts.tv_sec);
};

bool 
cleanupFlowTest(
	void* current, 
	void* udata)
{
	CFlowNode& existingNode = *((CFlowNode*)current);
	if ((currentTime - existingNode.m_flow.records.last) > FLOW_LIFETIME)
		return true;
	return false;
};

bool 
cleanupHostTest(
	void* current, 
	void* udata)
{
	CHostNode& existingNode = *((CHostNode*)current);
	if ((((currentTime - existingNode.m_lastFlow) > HOST_LIFETIME) &&
	    ((currentTime - existingNode.m_lastPacket) > HOST_LIFETIME)) ||
		(existingNode.m_nrFlows <= 0))
		return true;
	return false;
};

// vi:ts=4