/* 
 * HostTree.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>

CHostTree::CHostTree()
{
	m_pRoot = NULL;
	m_numberOfNodes = 0;
}

CHostTree::~CHostTree()
{
#ifdef _DEBUG
	clock_t startTime;
	clock_t endTime;

	startTime = clock();
#endif
	if (m_pRoot != NULL) {
		delete m_pRoot;
		m_pRoot = NULL;
	}
	m_numberOfNodes = 0;
#ifdef _DEBUG
	endTime = clock();
	printf("Tree deleted in %.4f seconds\n",
		((double)(endTime-startTime)) / CLOCKS_PER_SEC);
#endif
}

bool 
CHostTree::insertNode(
	in_addr& ip, 
	time_t& lastflow, 
	time_t& lastpacket)
{
	CHostNode* pNode = new CHostNode(ip, lastflow, lastpacket);
	return insertNode(pNode);
};

bool 
CHostTree::insertNode(CHostNode* pNode)
{
	if (m_pRoot == NULL) {
		m_pRoot = pNode;
		m_pRoot->m_pParent = NULL;
		m_numberOfNodes++;
		return true;
	}

	CHostNode* pTemp = m_pRoot;

	while (true) {
		if (pTemp->m_ip.s_addr == pNode->m_ip.s_addr) {
			if (pNode->m_lastFlow) { // new flow
				pTemp->m_lastFlow = pNode->m_lastFlow;
				pTemp->m_nrFlows++;
			
				int deltaTime = pTemp->m_lastFlow - pTemp->m_lastCheck;
				if (deltaTime > 5) {
					double fps = (double)pTemp->m_nrFlows / (double)deltaTime;
					if (fps > MAX_FLOWS_PER_SEC) {
						syslog(LOG_LOCAL6|LOG_ALERT, "FLOOD: %d.%d.%d.%d - %.2f flows/s",
							((unsigned char*)&pTemp->m_ip)[0],
							((unsigned char*)&pTemp->m_ip)[1],
							((unsigned char*)&pTemp->m_ip)[2],
							((unsigned char*)&pTemp->m_ip)[3],
							fps);
						dprintf("FLOOD: %d.%d.%d.%d - %.2f flows/s\n",
							((unsigned char*)&pTemp->m_ip)[0],
							((unsigned char*)&pTemp->m_ip)[1],
							((unsigned char*)&pTemp->m_ip)[2],
							((unsigned char*)&pTemp->m_ip)[3],
							fps);
					}
					pTemp->m_nrFlows = 0;
					pTemp->m_lastCheck = pTemp->m_lastFlow;
				}
			}

			if (pNode->m_lastPacket) { // flow update
				pTemp->m_lastPacket = pNode->m_lastPacket;
				pTemp->m_nrPackets++;
		
				int deltaTime = pTemp->m_lastPacket - pTemp->m_lastCheck;
				if (deltaTime > 5) {
					double pps = (double)pTemp->m_nrPackets / (double)deltaTime;
					if (pps > MAX_PPS_PER_HOST) {
						syslog(LOG_LOCAL6|LOG_ALERT, "FLOOD: %d.%d.%d.%d - %.2f packets/s",
							((unsigned char*)&pTemp->m_ip)[0],
							((unsigned char*)&pTemp->m_ip)[1],
							((unsigned char*)&pTemp->m_ip)[2],
							((unsigned char*)&pTemp->m_ip)[3],
							pps);
						dprintf("FLOOD: %d.%d.%d.%d - %.2f packets/s\n",
							((unsigned char*)&pTemp->m_ip)[0],
							((unsigned char*)&pTemp->m_ip)[1],
							((unsigned char*)&pTemp->m_ip)[2],
							((unsigned char*)&pTemp->m_ip)[3],
							pps);
					}
					pTemp->m_nrPackets = 0;
					pTemp->m_lastCheck = pTemp->m_lastPacket;
				}
			}

			delete pNode;
			pNode = NULL;
			return false;
		}
		
		if (pTemp->m_ip.s_addr < pNode->m_ip.s_addr) {
			if (pTemp->m_pRight == NULL) {
				pTemp->m_pRight = pNode;
				pNode->m_pParent = pTemp;
				m_numberOfNodes++;
				return true;
			} else {
				pTemp = pTemp->m_pRight;
				continue;
			}
		}
		
		if (pTemp->m_ip.s_addr > pNode->m_ip.s_addr) {
			if (pTemp->m_pLeft == NULL) {
				pTemp->m_pLeft = pNode;
				pNode->m_pParent = pTemp;
				m_numberOfNodes++;
				return true;
			} else {
				pTemp = pTemp->m_pLeft;
				continue;
			}
		}
	}

	return false;
}

void CHostTree::deleteNode(CHostNode* pNode)
{
	if (pNode == NULL)
		return;

	if (pNode == m_pRoot) {
		if ((pNode->m_pRight == NULL) && 
			(pNode->m_pLeft == NULL)) { //no children
			delete m_pRoot;
			m_pRoot = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight != NULL) && 
			 (pNode->m_pLeft == NULL)) { //only right child
			m_pRoot = m_pRoot->m_pRight;
			m_pRoot->m_pParent = NULL;
			pNode->m_pRight = NULL;
			delete pNode;
			pNode = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight == NULL) && 
			(pNode->m_pLeft != NULL)) { //only left child
			m_pRoot = m_pRoot->m_pLeft;
			m_pRoot->m_pParent = NULL;
			pNode->m_pLeft = NULL;
			delete pNode;
			pNode = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight != NULL) && 
			(pNode->m_pLeft != NULL)) { //both children
			if (pNode->m_pLeft->m_pRight == NULL) {
				pNode->m_pLeft->m_pRight = m_pRoot->m_pRight;
				m_pRoot->m_pRight->m_pParent = pNode->m_pLeft;
				pNode->m_pLeft->m_pParent = NULL;
				m_pRoot = pNode->m_pLeft;

				pNode->m_pLeft = NULL;
				pNode->m_pRight = NULL;
				delete pNode;
				pNode = NULL;
				m_numberOfNodes--;
				return;
			} else {
				CHostNode* pPredecesor = pNode->m_pLeft;
				while (pPredecesor->m_pRight != NULL)
					pPredecesor = pPredecesor->m_pRight;

				if (pPredecesor->m_pLeft == NULL) {
					pPredecesor->m_pParent->m_pRight = NULL;
					pPredecesor->m_pParent = NULL;

					pPredecesor->m_pLeft = m_pRoot->m_pLeft;
					pPredecesor->m_pRight = m_pRoot->m_pRight;

					pPredecesor->m_pLeft->m_pParent = pPredecesor;
					pPredecesor->m_pRight->m_pParent = pPredecesor;

					m_pRoot = pPredecesor;

					pNode->m_pLeft = NULL;
					pNode->m_pRight = NULL;
					delete pNode;
					pNode = NULL;
					m_numberOfNodes--;
					return;
				} else {
					pPredecesor->m_pParent->m_pRight = pPredecesor->m_pLeft;
					pPredecesor->m_pLeft->m_pParent = pPredecesor->m_pParent;

					pPredecesor->m_pParent = NULL;

					pPredecesor->m_pLeft = m_pRoot->m_pLeft;
					pPredecesor->m_pRight = m_pRoot->m_pRight;

					pPredecesor->m_pLeft->m_pParent = pPredecesor;
					pPredecesor->m_pRight->m_pParent = pPredecesor;

					m_pRoot = pPredecesor;

					pNode->m_pLeft = NULL;
					pNode->m_pRight = NULL;
					delete pNode;
					pNode = NULL;
					m_numberOfNodes--;
					return;
				}
			}
		}
	} else {
		if ((pNode->m_pRight == NULL) && 
			(pNode->m_pLeft == NULL)) { //no children
			if (pNode->m_pParent->m_pLeft == pNode)
				pNode->m_pParent->m_pLeft = NULL;
			else
				pNode->m_pParent->m_pRight = NULL;
			delete pNode;
			pNode = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight != NULL) && 
			(pNode->m_pLeft == NULL)) { //only right child
			if (pNode->m_pParent->m_pLeft == pNode)
				pNode->m_pParent->m_pLeft = pNode->m_pRight;
			else
				pNode->m_pParent->m_pRight = pNode->m_pRight;

			pNode->m_pRight->m_pParent = pNode->m_pParent;
			pNode->m_pRight = NULL;
			delete pNode;
			pNode = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight == NULL) && 
			(pNode->m_pLeft != NULL)) { //only left child
			if (pNode->m_pParent->m_pLeft == pNode)
				pNode->m_pParent->m_pLeft = pNode->m_pLeft;
			else
				pNode->m_pParent->m_pRight = pNode->m_pLeft;
				
			pNode->m_pLeft->m_pParent = pNode->m_pParent;
			pNode->m_pLeft = NULL;
			delete pNode;
			pNode = NULL;
			m_numberOfNodes--;
			return;
		}

		if ((pNode->m_pRight != NULL) && 
			(pNode->m_pLeft != NULL)) { //both children
			if (pNode->m_pLeft->m_pRight == NULL) {
				if (pNode->m_pParent->m_pLeft == pNode)
					pNode->m_pParent->m_pLeft = pNode->m_pLeft;
				else
					pNode->m_pParent->m_pRight = pNode->m_pLeft;
					
				pNode->m_pLeft->m_pParent = pNode->m_pParent;

				pNode->m_pLeft->m_pRight = pNode->m_pRight;
				pNode->m_pRight->m_pParent = pNode->m_pLeft;

				pNode->m_pLeft = NULL;
				pNode->m_pRight = NULL;
				delete pNode;
				pNode = NULL;
				m_numberOfNodes--;

				return;
			} else {
				CHostNode* pPredecesor = pNode->m_pLeft;
				while (pPredecesor->m_pRight != NULL)
					pPredecesor = pPredecesor->m_pRight;

				if (pPredecesor->m_pLeft == NULL) {
					if (pNode->m_pParent->m_pLeft == pNode)
						pNode->m_pParent->m_pLeft = pPredecesor;
					else
						pNode->m_pParent->m_pRight = pPredecesor;
					
					pPredecesor->m_pParent->m_pRight = NULL;
					pPredecesor->m_pParent = pNode->m_pParent;
					
					pPredecesor->m_pLeft = pNode->m_pLeft;
					pPredecesor->m_pRight = pNode->m_pRight;
					
					pPredecesor->m_pLeft->m_pParent = pPredecesor;
					pPredecesor->m_pRight->m_pParent = pPredecesor;
					
					pNode->m_pLeft = NULL;
					pNode->m_pRight = NULL;
					delete pNode;
					pNode = NULL;
					m_numberOfNodes--;
					
					return;
				} else {
					if (pNode->m_pParent->m_pLeft == pNode)
						pNode->m_pParent->m_pLeft = pPredecesor;
					else
						pNode->m_pParent->m_pRight = pPredecesor;
					
					pPredecesor->m_pParent->m_pRight = pPredecesor->m_pLeft;
					pPredecesor->m_pLeft->m_pParent = pPredecesor->m_pParent;

					pPredecesor->m_pParent = pNode->m_pParent;
					
					pPredecesor->m_pLeft = pNode->m_pLeft;
					pPredecesor->m_pRight = pNode->m_pRight;
					
					pPredecesor->m_pLeft->m_pParent = pPredecesor;
					pPredecesor->m_pRight->m_pParent = pPredecesor;
					
					pNode->m_pLeft = NULL;
					pNode->m_pRight = NULL;
					delete pNode;
					pNode = NULL;
					m_numberOfNodes--;
					
					return;
				}
			}
		}
	}
};

void 
CHostTree::cleanUp(CLEANFUNC cleanfunc)
{
	stack<CHostNode*> tempStack;
	CHostNode* pCursor = m_pRoot;		//set cursor to root of binary tree
	CHostNode* pTemp;
	stack<CHostNode*> toBeDeleted;
	bool done = false;
#ifdef _DEBUG
	time_t start, end;
	u_int before, after;
	before = m_numberOfNodes;
	start = clock();
#endif
	
	
	while (!done) {
		if (pCursor != NULL) {
			tempStack.push(pCursor);		//place pointer to node on the stack
											//before traversing the node's left subtree
			pCursor = pCursor->m_pLeft;		//traverse the left subtree
		} else {
			//backtrack from the empty subtree and
			//visit the node at the top of the stack;
			//however, if the stack is empty, you are
			//done
			if (!tempStack.empty()) {
				pCursor = tempStack.top();
				tempStack.pop();
				if (cleanfunc(pCursor, NULL))
				    toBeDeleted.push(pCursor);
				pCursor = pCursor->m_pRight;
			} else
				done = true;
		}
	}
	
	while (!toBeDeleted.empty()) {
	    pCursor = toBeDeleted.top();
	    toBeDeleted.pop();
	    deleteNode(pCursor);
	}
#ifdef _DEBUG
	end = clock();
	after = m_numberOfNodes;
	fprintf(stdout, "Host cleanup took %2.5f seconds (%u -> %u)\n",
		(double)(end - start)/(double)CLOCKS_PER_SEC, before, after);
#endif
	
};

u_int 
CHostTree::getNumberOfNodes(void)
{
	return m_numberOfNodes;
};
	

#ifdef _DEBUG
void 
CHostTree::show(FILE* file)
{
	if (m_pRoot == NULL)
		return;
	m_pRoot->show(file);
};
#endif //_DEBUG

// vi:ts=4