/* 
 * FlowTree.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::CFlowTree()
{
	m_pRoot = NULL;
	m_numberOfNodes = 0;
};

CFlowTree::~CFlowTree()
{
#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 
CFlowTree::insertNode(
	in_addr& saddr,
	in_addr& daddr,
	u_short& sport,
	u_short& dport,
	u_char& tcp_flags,
	u_char& proto,
	u_char& tos,
	u_int& bytes,
	time_t& now)
{
	CFlowNode* pNode = new CFlowNode(saddr,
									 daddr,
									 sport,
									 dport,
									 tcp_flags,
									 proto,
									 tos,
									 bytes,
									 now);
					
	return insertNode(pNode);
};

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

	CFlowNode* pTemp = m_pRoot;

	while (true) {
		if (*pTemp == *pNode) {
			pTemp->m_flow.records.last = pNode->m_flow.records.last;
			pTemp->m_flow.records.packets++;
			pTemp->m_flow.records.bytes += pNode->m_flow.records.bytes;

            int deltaTime = pTemp->m_flow.records.last - pTemp->m_flow.stats.lastcheck;
            if (deltaTime > 5 ) {
                double pps = (double)pTemp->m_flow.records.packets / (double)deltaTime;
                if (pps > MAX_PACKETS_PER_SEC_IN_FLOW) {
					struct protoent* proto = getprotobynumber(pTemp->m_flow.records.proto);
					syslog(LOG_LOCAL6|LOG_ALERT, 
						   "FLOOD: %d.%d.%d.%d - %.2f %s packets/s from %d.%d.%d.%d [DSCP %02X]",
						   ((unsigned char*)&pTemp->m_flow.records.daddr)[0],
						   ((unsigned char*)&pTemp->m_flow.records.daddr)[1],
						   ((unsigned char*)&pTemp->m_flow.records.daddr)[2],
						   ((unsigned char*)&pTemp->m_flow.records.daddr)[3],
						   pps,
						   proto->p_name,
						   ((unsigned char*)&pTemp->m_flow.records.saddr)[0],
						   ((unsigned char*)&pTemp->m_flow.records.saddr)[1],
						   ((unsigned char*)&pTemp->m_flow.records.saddr)[2],
						   ((unsigned char*)&pTemp->m_flow.records.saddr)[3],
						   pTemp->m_flow.records.tos >> 2);
							
					dprintf("FLOOD: %d.%d.%d.%d - %.2f %s packets/s from %d.%d.%d.%d [DSCP %02X]\n",
							((unsigned char*)&pTemp->m_flow.records.daddr)[0],
							((unsigned char*)&pTemp->m_flow.records.daddr)[1],
							((unsigned char*)&pTemp->m_flow.records.daddr)[2],
							((unsigned char*)&pTemp->m_flow.records.daddr)[3],
							pps,
							proto->p_name,
							((unsigned char*)&pTemp->m_flow.records.saddr)[0],
							((unsigned char*)&pTemp->m_flow.records.saddr)[1],
							((unsigned char*)&pTemp->m_flow.records.saddr)[2],
							((unsigned char*)&pTemp->m_flow.records.saddr)[3],
							pTemp->m_flow.records.tos >> 2);
                }
                pTemp->m_flow.records.packets = 0;
                pTemp->m_flow.stats.lastcheck = pTemp->m_flow.records.last;
            }
	
			delete pNode;
			pNode = NULL;
			
			return false;
		}
		
		if (*pTemp < *pNode) {
			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 > *pNode) {
			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 
CFlowTree::deleteNode(CFlowNode* 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 {
				CFlowNode *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 {
				CFlowNode* 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 
CFlowTree::cleanUp(CLEANFUNC cleanfunc)
{
	stack<CFlowNode*> tempStack;
	CFlowNode* pCursor = m_pRoot;			//set cursor to root of binary tree
	CFlowNode* pTemp;
	stack<CFlowNode *> 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, "Flow cleanup took %2.5f seconds (%u -> %u)\n",
		(double)(end - start)/(double)CLOCKS_PER_SEC, before, after);
#endif
	
};

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

#ifdef _DEBUG
int 
CFlowTree::show(void)
{
	if (m_pRoot == NULL)
		return 0;
	int nrNodes = 0;
	m_pRoot->show(nrNodes);
	return nrNodes;
};
#endif //_DEBUG

// vi:ts=4