/*---------------------------------------------------------------------------
    Copyright (c) 1998 - 2000  Microsoft Corporation
    Module Name: IpStat.exe    
    File       : IpStat.cpp
    Description: This file demonstrates the use of IP Helper APIs to
                 get IP statistics.
    Author:
    Frank Li            April 21, 1998

    
    Revision History:
    Who         When        What
    --------    --------    ----------------------------------
    Frank Li    04-21-98    created   
---------------------------------------------------------------------------*/

#include "IpStat.h"

void Usage(char * pszProgramName)
{
    printf("Manipulates IP Statistics.\n\n");


    printf("%s -p proto      Shows connections for the protocol specified\n", pszProgramName);
    printf("                     by proto, proto may be tcp or udp.\n");  
    printf("%s -s [-p proto] Displays per-protocol statistics.\n", pszProgramName);
    printf("                     By default, statistics are shown for\n");
    printf("                     IP, ICMP, TCP and UDP; the -p option\n");
    printf("                     may be used to specify a subset of the default.\n");

    printf("Examples:\n\n");

    printf("> IpStat -p tcp\n");
    printf("> IpStat  -s\n");

    WSACleanup();
    exit(1);
}


void _cdecl main(int argc, char **argv)
{
    WORD wVersionRequested = MAKEWORD(1,1);
    WSADATA wsaData;
    int nRet;

    nRet = WSAStartup(wVersionRequested, &wsaData);
    if (wsaData.wVersion != wVersionRequested)
    {    
        fprintf(stderr,"\n Wrong version\n");
        return;
    }

    if ((argc < 2) || (argv[1][0] != '-'))        
        Usage("IpStat");
    if (strlen(argv[1]) > 2)     
        Usage("IpStat");  
    
    switch(argv[1][1]) 
    {    
    case 'p':
        // Print connection table
        if (argc == 3)
            DoGetConnTable(argv[2]);
        else
            Usage("IpStat");
        break;
    case 's':   
        // show statistics
        if (argc == 4 && argv[2][1] == 'p')
            DoGetStat(argv[3]); // Get stat for a specific protocol
        else if (argc == 2)
            DoGetStat(); // Get stat for all protocols
        else
            Usage("IpStat");
        break;    
         
        
    default:
        // help
        Usage("IpStat");        
        break;    
    }
    WSACleanup();
}

void DoGetConnTable(char* pszProto)
{
    DWORD dwStatus;
    if (_strnicmp(pszProto, "tcp", 3) == 0)
    {
        //Print Tcp Connnection Table
        PMIB_TCPTABLE pTcpTable = NULL;
        dwStatus = MyGetTcpTable(pTcpTable, TRUE);
        if (dwStatus != NO_ERROR)
        {
            printf("Ipstat: Couldn't get tcp connection table.\n");
            if (pTcpTable)
                free(pTcpTable);
            return;
        }
        else
        {
            DumpTcpTable(pTcpTable);
            free(pTcpTable);
        }
    }
    else if (_strnicmp(pszProto, "udp", 3) == 0)
    {
        //Print Udp Table
        PMIB_UDPTABLE pUdpTable = NULL;
        dwStatus = MyGetUdpTable(pUdpTable, TRUE);
        if (dwStatus != NO_ERROR)
        {
            printf("Ipstat: Couldn't get udp table.\n");
            if (pUdpTable)
                free(pUdpTable);
            return;
        }
        else
        {
            DumpUdpTable(pUdpTable);
            free(pUdpTable);
        }
    }
    else
        Usage("IpStat");
}

void DoGetStat(char* pszProto /*NULL*/)
{
    if (pszProto == NULL)
    {
        // by default, display all statistics
        {
        PMIB_IPSTATS pIpStats = NULL;
        if (MyGetIpStatistics(pIpStats) != NO_ERROR)
            printf("IpStat: error in getting ip statistics.\n");
        else
            PrintIpStats(pIpStats);

        if (pIpStats)
            free(pIpStats);
        }
        {
        PMIB_ICMP pIcmpStats = NULL;
        if (MyGetIcmpStatistics(pIcmpStats) != NO_ERROR)
            printf("IpStat: error in getting icmp statistics.\n");
        else
            PrintIcmpStats(&(pIcmpStats->stats));

        if (pIcmpStats)
            free(pIcmpStats);

        }
        {
        PMIB_TCPSTATS pTcpStats = NULL;
        if (MyGetTcpStatistics(pTcpStats) != NO_ERROR)
            printf("IpStat: error in getting tcp statistics.\n");
        else
            PrintTcpStats(pTcpStats);

        if (pTcpStats)
            free(pTcpStats);
        }
        {
        PMIB_UDPSTATS pUdpStats = NULL;
        if (MyGetUdpStatistics(pUdpStats) != NO_ERROR)
            printf("IpStat: error in getting udp statistics.\n");
        else
            PrintUdpStats(pUdpStats);

        if (pUdpStats)
            free(pUdpStats);

        }
    }
    else if (_strnicmp(pszProto, "ip", 2) == 0)
    {
        PMIB_IPSTATS pIpStats = NULL;
        if (MyGetIpStatistics(pIpStats) != NO_ERROR)
            printf("IpStat: error in getting ip statistics.\n");
        else
            PrintIpStats(pIpStats);

        if (pIpStats)
            free(pIpStats);
    }

    else if (_strnicmp(pszProto, "icmp", 4) == 0)
    {
        PMIB_ICMP pIcmpStats = NULL;
        if (MyGetIcmpStatistics(pIcmpStats) != NO_ERROR)
            printf("IpStat: error in getting icmp statistics.\n");
        else
            PrintIcmpStats(&(pIcmpStats->stats));

        if (pIcmpStats)
            free(pIcmpStats);

    }

    else if (_strnicmp(pszProto, "tcp", 3) == 0)
    {
        PMIB_TCPSTATS pTcpStats = NULL;
        if (MyGetTcpStatistics(pTcpStats) != NO_ERROR)
            printf("IpStat: error in getting tcp statistics.\n");
        else
            PrintTcpStats(pTcpStats);

        if (pTcpStats)
            free(pTcpStats);

    }

    else if (_strnicmp(pszProto, "udp", 3) == 0)
    {
        PMIB_UDPSTATS pUdpStats = NULL;
        if (MyGetUdpStatistics(pUdpStats) != NO_ERROR)
            printf("IpStat: error in getting udp statistics.\n");
        else
            PrintUdpStats(pUdpStats);

        if (pUdpStats)
            free(pUdpStats);
    }
    else
        printf("IpStat: no available statistics for %s.\n", pszProto);

}


void DumpTcpTable(PMIB_TCPTABLE pTcpTable)
{
    char    strState[MAX_STRLEN];
    struct  in_addr    inadLocal, inadRemote;
    DWORD   dwRemotePort = 0;
    char    szLocalIp[MAX_STRLEN];
    char    szRemIp[MAX_STRLEN];

    if (pTcpTable != NULL)
    {
        printf("TCP TABLE\n");
        printf("%20s %10s %20s %10s %s\n", "Loc Addr", "Loc Port", "Rem Addr",
                "Rem Port", "State");
        for (UINT i = 0; i < pTcpTable->dwNumEntries; ++i)
        {
            switch (pTcpTable->table[i].dwState)
            {
            case MIB_TCP_STATE_CLOSED:
                strcpy(strState, "CLOSED");
                break;
            case MIB_TCP_STATE_TIME_WAIT:
                strcpy(strState, "TIME_WAIT");
                break;
            case MIB_TCP_STATE_LAST_ACK:
                strcpy(strState, "LAST_ACK");
                break;
            case MIB_TCP_STATE_CLOSING:
                strcpy(strState, "CLOSING");
                break;
            case MIB_TCP_STATE_CLOSE_WAIT:
                strcpy(strState, "CLOSE_WAIT");
                break;
            case MIB_TCP_STATE_FIN_WAIT1:
                strcpy(strState, "FIN_WAIT1");
                break;
            case MIB_TCP_STATE_ESTAB:
                strcpy(strState, "ESTAB");
                break;
            case MIB_TCP_STATE_SYN_RCVD:
                strcpy(strState, "SYN_RCVD");
                break;
            case MIB_TCP_STATE_SYN_SENT:
                strcpy(strState, "SYN_SENT");
                break;
            case MIB_TCP_STATE_LISTEN:
                strcpy(strState, "LISTEN");
                break;