/*
 * Copyright (c) 2000-2008
 * Author: Weiming Zhou
 *
 * Permission to use, copy, modify, distribute and sell this software
 * and its documentation for any purpose is hereby granted without fee,
 * provided that the above copyright notice appear in all copies and
 * that both that copyright notice and this permission notice appear
 * in supporting documentation.  
 */

#include <windows.h>
#include <time.h>
#include "CapiGlobal.h"
#include "BinTree.h"
#include "AVLTree.h"
#include "HashAVLTree.h"

#define     TCPHEAD_LENGTH  20
#define     DEFAULT_CONNECT_TIME    5000
#define     MAX_CONNECT_TIMES       5

typedef struct CLIENTDATA_st {
    DWORD   dwIPAddr;
    clock_t StartClock;
    clock_t TotalClock;
    DWORD   dwConnectTimes;
    INT     nTcpData[TCPHEAD_LENGTH];
} CLIENTDATA;

/**	判斷是否DoS攻擊的回調函數

	@param	CLIENTDATA * pData - CLIENTDATA指針	
	@return	BOOL - TRUE表示是DoS攻擊，FALSE表示不是	
*/
typedef  BOOL   (*ISDOSATTACKFUNC)(CLIENTDATA *pData);

/**	DoSUnattack創建函數，創建一個哈希AVL樹。

	@param	UINT uBucketCount - 哈希AVL樹的BUCKET數組大小	
	@return	HASHAVLTREE * - 哈希AVL樹指針	
*/
HASHAVLTREE * DoSUnattack_Create(UINT uBucketCount)
{
    HASHAVLTREE *pTree;
    pTree = HashAVLTree_Create(uBucketCount);

    return pTree;
}

/**	哈希AVL樹的比較回調函數，比較兩個IP地址是否相等

	@param	void *p1 - CLIENTDATA指針類型	
	@param	void *p2 - DWORD類型	
	@return	INT - 0表示相等，-1表示p1小于p2，1表示p1大于p2。	
*/
INT IPCompare(void *p1, void *p2)
{
    CLIENTDATA  *pData;
    DWORD dwIPAddr;

    pData = (CLIENTDATA *)p1;
    dwIPAddr = (DWORD)p2;

    if ( pData->dwIPAddr == dwIPAddr )
    {
        return 0;
    }
    else if ( pData->dwIPAddr < dwIPAddr )
    {
        return -1;
    }
    else
    {
        return 1;
    }
}

/**	哈希AVL樹的計算哈希值的回調函數

	@param	void *p - 要計算哈希值的對應關鍵詞指針	
	@param	UINT uBucketCount - 哈希AVL樹中的BUCKET數組大小	
	@return	INT - BUCKET數組下標	
*/
INT HashClientData(void *p, UINT uBucketCount)
{
    CLIENTDATA  *pData = (CLIENTDATA *)p;

    return pData->dwIPAddr % uBucketCount;
}

/**	接收客戶端的連接

	@param	HASHAVLTREE *pTree - 哈希AVL樹指針	
	@param	DWORD dwIPAddr - 對方IP地址	
	@param	INT *pTcpIpData - 連接包數據	
	@param	ISDOSATTACKFUNC IsDosAttack - 判斷是否DoS攻擊的回調函數	
	@return	void - 無
*/
void DoSUnattack_RecvConnection(HASHAVLTREE *pTree, DWORD dwIPAddr,
                    INT *pTcpIpData, ISDOSATTACKFUNC IsDosAttack)
{
    CLIENTDATA *pData;
    clock_t     ConnectClock;
    
    pData = (CLIENTDATA *)HashAVLTree_Find(pTree, (void *)dwIPAddr, 
                                HashInt, IPCompare);

    ConnectClock = clock();
    if ( pData != NULL )
    {
        pData->dwConnectTimes += 1;
        if ( (*IsDosAttack)(pData) )
        {
            /* 發現攻擊，拒絕連接，將對方IP加入黑名單中 */
        }
    }
    else
    {
        pData = (CLIENTDATA *)malloc(sizeof(CLIENTDATA));
        if ( pData != NULL )
        {
            pData->dwConnectTimes = 1;
            pData->dwIPAddr = dwIPAddr;
            pData->StartClock = clock();
            pData->TotalClock = 0;
            memcpy(pData->nTcpData, pTcpIpData, TCPHEAD_LENGTH * sizeof(INT)); 

            HashAVLTree_Insert(pTree, pData, HashClientData, IPCompare);
        }
    }

    return;
}

/**	DosUnattack的連接任務, 負責和客戶端的連接

	@param	HASHAVLTREE *pTree - 哈希AVL樹指針	
	@param	ISDOSATTACKFUNC IsDosAttack - 判斷是否DoS攻擊的回調函數	
	@return	void - 無	
*/
void DosUattack_ConnectTask(HASHAVLTREE *pTree, 
                           ISDOSATTACKFUNC IsDosAttack)
{
    CLIENTDATA  *pData;

    HashAVLTree_EnumBegin(pTree);

    while( (pData = HashAVLTree_EnumNext(pTree) ) != NULL )
    {
        if ( (*IsDosAttack)(pData) )
        {
            /* 發現攻擊，拒絕連接，將對方IP計入黑名單*/
            /* ..... */
        }
        else
        {
            /* 和客戶端建立連接 */
        }
    }
}