Enhance netstat

// ENetStat.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "ENetStat.h"
#include "TCPTable.h"
#include "UDPClass.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

/////////////////////////////////////////////////////////////////////////////
// The one and only application object

CWinApp theApp;
int		m_nOption = 1; //help
u_long	m_ulRemoteIP = 0;
u_short m_usLocalPort = 0;
using namespace std;

/////////////////////////////////////////////////////////////////////////////
// Routines
void ShowBanner()
{
	cerr << _T("\n\tENetStat - enhanced netstat") << endl;
	cerr << _T("\tVersion 0.1 - by y0d4") << endl;
}
//////////////////////////////////////////////////////////////////////
//
void ShowUsage()
{
	cerr << _T("Usage:  ENetStat [options] ") << endl << endl;
	cerr << _T("\t[options] are any of:") << endl;

	cerr << _T("\t/a \tshow all connections") << endl;
	cerr << _T("\t/? \tshow this help page") << endl;
	cerr << _T("\t/t \tshow TCP connections") << endl;
	cerr << _T("\t/e \tshow established connections") << endl;
	cerr << _T("\t/k \tkill specified 'established' connection (specify remote IP and local port)") << endl;
	cerr << _T("\t/u \tshow UDP connections") << endl;
	cerr << _T("\t/v \tshow version") << endl;

	cerr << endl;
}
//////////////////////////////////////////////////////////////////////
// ParseOptions() looks for options on the command line and sets global
// flags so the rest of the program knows about them.  
BOOL ParseOptions(int argc, char* argv[], char *envp[])
{
	CString str;

	int nIndex;
	for (nIndex = 1; nIndex < argc; nIndex++)
	{
		if (*argv[nIndex] == '-' || *argv[nIndex] == '/')
		{
			if (argv[nIndex][1] == 'a' || argv[nIndex][1] == 'A')
				m_nOption = 0;
			else if (argv[nIndex][1] == '?')
				m_nOption = 1;
			else if (argv[nIndex][1] == 't' || argv[nIndex][1] == 'T')
				m_nOption = 2;
			else if (argv[nIndex][1] == 'u' || argv[nIndex][1] == 'U')
				m_nOption = 3;
			else if (argv[nIndex][1] == 'v' || argv[nIndex][1] == 'V')
				m_nOption = 4;
			else if (argv[nIndex][1] == 'e' || argv[nIndex][1] == 'E')
				m_nOption = 5;
			else if (argv[nIndex][1] == 'k' || argv[nIndex][1] == 'K')
			{
				int len = sizeof(argv);
				
				//get remote ip
				int i = 2;				
				do {
					if (argv[nIndex][i] != 0)
						str += argv[nIndex][i];
					i++;
				} while (argv[nIndex][i]);
				m_ulRemoteIP = inet_addr(str.GetBuffer(str.GetLength()));
				if (INADDR_NONE == m_ulRemoteIP)
					cerr << _T("\t \tInvalid argument in case of remote IP\n") << endl;


				//get local port
				i = str.GetLength();
				if(i <= 4)
				{
					cerr << _T("\t \tInvalid argument in case of local port\n") << endl;
					return FALSE;
				}
				str.Empty();
				nIndex ++;
				do {
					if (argv[nIndex][i] != 0)
						str += argv[nIndex][i];
					i++;
				} while (argv[nIndex][i]);
				m_usLocalPort = htons((u_short)atoi(str.GetBuffer(str.GetLength())));

				//set the selected option
				m_nOption = 6;
			}
			else
			{
				cerr << _T("Error: unrecognized option: ") << argv[nIndex] << endl;
				return FALSE;
			}
		}
	}

	return TRUE;
}

int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
{
	int nRetCode = 0;

	if (!AfxWinInit(::GetModuleHandle(NULL), NULL, ::GetCommandLine(), 0))
	{
		cerr << _T("MFC Failed to initialize.\n");
		return 1;
	}

	if (argc < 1 || !ParseOptions(argc, argv, envp))
		ShowUsage();

	CTCPTable *pTcp = new CTCPTable();
	CUDPClass *pUdp = new CUDPClass();
	DWORD dwRez = 0;

	switch(m_nOption)
	{
		case 0:		//all
			dwRez = pTcp->PrintTableEx();
			if (dwRez != 1) dwRez = pTcp->PrintTable();
			printf("\n");

			dwRez = pUdp->PrintTableEx();
			if (dwRez != 1) dwRez = pUdp->PrintTable();
			printf("\n");
		break;

		case 1:		//help
			ShowUsage();
			printf("\n");
		break;

		case 2:		//tcp
			dwRez = pTcp->PrintTableEx();
			if (dwRez != 1) dwRez = pTcp->PrintTable();
			printf("\n");
		break;

		case 3:		//udp
			dwRez = pUdp->PrintTableEx();
			if (dwRez != 1) dwRez = pUdp->PrintTable();
			printf("\n");
		break;

		case 4:		//ver
			ShowBanner();
			printf("\n");
		break;

		case 5:		//tcp
			pTcp->m_bEstb = true;
			dwRez = pTcp->PrintTableEx();
			if (dwRez != 1) dwRez = pTcp->PrintTable();
			printf("\n");
		break;
				
		case 6:
			pTcp->KillConnection(m_ulRemoteIP, m_usLocalPort);
			pTcp->m_bEstb = true;
			dwRez = pTcp->PrintTableEx();
			if (dwRez != 1) dwRez = pTcp->PrintTable();
			printf("\n");
		break;

		default:	//help
			ShowUsage();
			printf("\n");
		break;
	}

	delete pTcp;
	pTcp = NULL;
	delete pUdp;
	pUdp = NULL;

	return nRetCode;
}