/*---------------------------------------------------------------------------
    Copyright (c) 1998  Microsoft Corporation
---------------------------------------------------------------------------*/
#include <windows.h>
// Link to ws2_32.lib
#include <winsock.h>
// http://msdn.microsoft.com/en-us/library/aa365872(VS.85).aspx
// Link to Iphlpapi.lib
#include <iphlpapi.h>
#include <assert.h>
#include <stdlib.h>
#include <malloc.h>
#include <stdio.h>
#include <string.h>
#include <limits.h>
#include <iptypes.h>

// Prototypes
void  PrintIpNetTable(PMIB_IPNETTABLE pIpNetTable);
// Custom made - http://msdn.microsoft.com/en-us/library/aa365956.aspx
DWORD MyGetIpNetTable(PMIB_IPNETTABLE pIpArpTab, BOOL fOrder);
// Custome made - http://msdn.microsoft.com/en-us/library/aa365949.aspx
DWORD MyGetIpAddrTable(PMIB_IPADDRTABLE pIpAddrTable, BOOL fOrder);
void  PrintIpAddrTable(PMIB_IPADDRTABLE pIpAddrTable);
void DoGetIpNetTable();
void DoSetIpNetEntry(char* pszDottedInetAddr, char* pszPhysAddr, char* pszInterface);
void DoDeleteIpNetEntry(char* pszInetAddr, char* pszInterface);
BOOL PhysAddrToString(BYTE PhysAddr[], DWORD PhysAddrLen, char str[]);
int  StringToPhysAddr(char* szInEther, char* szOutEther);
BOOL InterfaceIdxToInterfaceIp(PMIB_IPADDRTABLE pIpAddrTable, DWORD dwIndex, char str[]);

void Usage(char * pszProgramName)
{
    printf("%s -s inet_addr eth_addr [if_addr]\n", pszProgramName);
    printf("%s -d inet_addr [if_addr]\n", pszProgramName);
    printf("%s -a\n", pszProgramName);

    printf("-a            Displays current ARP entries by interrogating the current\n");
    printf("              protocol data.\n");
    printf("-d            Deletes the host specified by inet_addr.\n");
    printf("-s            Adds the host and associates the Internet address inet_addr\n");
    printf("              with the Physical address eth_addr.  The Physical address is\n");
    printf("              given as 6 hexadecimal bytes separated by hyphens. The entry\n");
    printf("              is permanent.\n");
    printf("eth_addr      Specifies a physical address.\n");
    printf("if_addr       If present, this specifies the Internet address of the\n");
    printf("              interface whose address translation table should be modified.\n");
    printf("              If not present, the first applicable interface will be used.\n");
    printf("Example:\n");
    printf("   >IpArp -s 157.55.85.212   00-aa-bb-cc-dd-ee 0x2000003 .... Add a static\n");
    printf("                                  arp entry on interface number 0x2000003.\n");
    printf("   >IpArp -a                                    ....Displays the arp table.\n");
    printf("   >IpArp -d 157.55.85.212                      ....Delete an entry.\n");
    WSACleanup();
    exit(1);
}

int main(int argc, char **argv)
{
	WSADATA wsaData;

    if ((argc < 2) || (argv[1][0] != '-'))
        Usage(argv[0]);
    if (strlen(argv[1]) > 2)
        Usage(argv[0]);
	
	
	if(WSAStartup(MAKEWORD(2,2), &wsaData) != 0)
    {
        printf("WSAStartup() is OK!\n");
        return -1;
    }
	else
		printf("WSAStartup() is OK!\n");

    switch(argv[1][1])
    {
	case 'a':
        // Print arp table
        DoGetIpNetTable();
        break;
    case 's':
        //Update or add an ARP Internet/Ethernet Address entry
        if (argc == 4)
            DoSetIpNetEntry(argv[2], argv[3], NULL);
        else if (argc == 5)
            DoSetIpNetEntry(argv[2], argv[3], argv[4]);
        else
            Usage(argv[0]);
        break;
    case 'd':
        //Delete an Internet/Ethernet Address pair from the ARP table
        if (argc == 3)
            DoDeleteIpNetEntry(argv[2], NULL);
        else if (argc == 4)
            DoDeleteIpNetEntry(argv[2], argv[3]);
        else
            Usage(argv[0]);
        break;
    default:
        // help
        Usage(argv[0]);
        break;
    }
    if(WSACleanup() == 0)
		printf("WSACleanup() is pretty fine!\n");
	else
		printf("WSACleanup() failed with error code %d\n", WSAGetLastError());
}

void DoGetIpNetTable()
{
    DWORD dwStatus;
    PMIB_IPNETTABLE pIpArpTab = NULL;


    if ( (dwStatus = MyGetIpNetTable(pIpArpTab, TRUE)) == NO_ERROR)
    {
		printf("Some entries in arp cache...\n");
        PrintIpNetTable(pIpArpTab);
        free(pIpArpTab);
        return;
    }
    else if ( dwStatus == ERROR_NO_DATA)
    {
        printf("No entries in arp cache...\n");
        if (pIpArpTab)
            free (pIpArpTab);
        return;
    }
    else
    {
        if (pIpArpTab)
            free (pIpArpTab);
        printf("IpArp returned 0x%x\n", dwStatus);
        return;
    }
}

//----------------------------------------------------------------------------
// Add an arp entry with ip dotted decimal address of "pszDottedInetAddr" and
// physical address of "pszPhysAddr" in 00-aa-bb-cc-dd-ee format on interface
// index of "pszInterface" in hex form.
//----------------------------------------------------------------------------
void DoSetIpNetEntry(char* pszDottedInetAddr, char* pszPhysAddr, char* pszInterface)
{
    DWORD dwInetAddr = 0; // ip address
    DWORD dwStatus;
    BYTE bPhysAddr[MAXLEN_PHYSADDR];
	MIB_IPNETROW arpEntry; // an arp entry
	PMIB_IPADDRTABLE pIpAddrTable;

	pszInterface = NULL;

    if (pszDottedInetAddr == NULL || pszPhysAddr == NULL)
    {
        printf("Bad Argument!\n");
        return;
    }

    dwInetAddr = inet_addr(pszDottedInetAddr); // convert dotted ip addr. to ip addr.
    if (dwInetAddr == INADDR_NONE)
    {
        printf("Bad Argument %s!\n", pszDottedInetAddr);
        return;
    }

    if (StringToPhysAddr(pszPhysAddr, (char*)bPhysAddr ) != 0)
    {
        printf("Bad Argument %s!\n", pszPhysAddr);
        return;
    }
    if (pszInterface)
    {
        // User provides a interface index number
        sscanf_s(pszInterface, "%X",&(arpEntry.dwIndex));
    }
    else
    {
        // add this to the first available interface
        pIpAddrTable = NULL;

        if ( (dwStatus = MyGetIpAddrTable(pIpAddrTable, FALSE)) != NO_ERROR)
        {
            printf("Couldn't find a interface number to add your arp entry\n");
            return;
        }
		else
			printf("MyGetIpAddrTable() is OK!\n"); 

        arpEntry.dwIndex = pIpAddrTable->table[0].dwIndex;
        free(pIpAddrTable);
    }

    arpEntry.dwPhysAddrLen = 6;
    memcpy(arpEntry.bPhysAddr, bPhysAddr, 6);
    arpEntry.dwAddr = dwInetAddr;
    arpEntry.dwType = MIB_IPNET_TYPE_STATIC; //static arp entry
    dwStatus = SetIpNetEntry(&arpEntry);
    if (dwStatus != NO_ERROR)
    {
        printf("Couldn't add (%s, %s), dwStatus = %lu.\n", pszDottedInetAddr, pszPhysAddr, dwStatus);
    }
}
//----------------------------------------------------------------------------
// Delete an arp entry with ip dotted decimal address of "pszDottedInetAddr"
// and interface index of "pszInterface" in hex form.
//----------------------------------------------------------------------------
void DoDeleteIpNetEntry(char* pszDottedInetAddr, char* pszInterface)
{
    DWORD dwInetAddr = 0; // ip address
    DWORD dwStatus;
	    MIB_IPNETROW arpEntry; // an arp entry


	pszInterface = NULL;

    if (pszDottedInetAddr == NULL)
    {
        printf("Bad Argument!\n");
        return;
    }
    dwInetAddr = inet_addr(pszDottedInetAddr); // convert dotted ip addr. to ip addr.
    if (dwInetAddr == INADDR_NONE)
    {
        printf("Bad Argument %s!\n", pszDottedInetAddr);
        return;
    }

    if (pszInterface)
    {
        // User provides a interface index number
        sscanf_s(pszInterface, "%X",&(arpEntry.dwIndex));
    }
    else
    {
        // try to delete this from first available interface
        PMIB_IPADDRTABLE pIpAddrTable = NULL;
        if ( (dwStatus = MyGetIpAddrTable(pIpAddrTable, FALSE)) != NO_ERROR)
        {
            printf("Couldn't find a interface number to add your arp entry\n");
            return;
        }
		else
			printf("MyGetIpAddrTable() is OK!\n");

        arpEntry.dwIndex = pIpAddrTable->table[0].dwIndex;
        free(pIpAddrTable);
    }

    arpEntry.dwAddr = dwInetAddr;
    dwStatus = DeleteIpNetEntry(&arpEntry);
    if (dwStatus != NO_ERROR)
    {
        printf("Couldn't delete (%s), dwStatus = %lu.\n", pszDottedInetAddr, dwStatus);
    }
}

//----------------------------------------------------------------------------
// Inputs: pIpAddrTable is the IP address table dwIndex is the Interface Number
// Output: If it returns TRUE, str contains the ip address of the interface
//----------------------------------------------------------------------------
BOOL InterfaceIdxToInterfaceIp(PMIB_IPADDRTABLE pIpAddrTable, DWORD dwIndex, char str[])
{
    struct in_addr inadTmp;
    char* szIpAddr;
	DWORD dwIdx;

    if (pIpAddrTable == NULL ||  str == NULL)
        return FALSE;

    str[0] = '\0';

    for (dwIdx = 0; dwIdx < pIpAddrTable->dwNumEntries; dwIdx++)
    {
        if (dwIndex == pIpAddrTable->table[dwIdx].dwIndex)
        {
            inadTmp.s_addr = pIpAddrTable->table[dwIdx].dwAddr;
            szIpAddr = inet_ntoa(inadTmp);
            if (szIpAddr)
            {
                strcpy_s(str, sizeof(str), szIpAddr);
                return TRUE;
            }
            else
                return FALSE;
        }
    }
    return FALSE;

}

//-------------------------------------------------------------------------
// Input: str points to an ethernet address of the form 00-aa-bb-cc-dd-ee
// If it returns 0, ret contains the 6 bytes ethernet address.
//-------------------------------------------------------------------------
int StringToPhysAddr(char* szInEther,  char* szOutEther)
{
    const char DASH = '-';
    register char c;
    register int val;
	int i;

    // check szInEther for the correct format
    if (strlen(szInEther) != 17)
        return (-1);
    if (szInEther[2] != DASH || szInEther[5] != DASH || szInEther[8] != DASH ||
        szInEther[8] != DASH || szInEther[14] != DASH)
        return (-1);
    if (!isxdigit(szInEther[0]) || !isxdigit(szInEther[1]) ||
        !isxdigit(szInEther[3]) || !isxdigit(szInEther[4]) ||
        !isxdigit(szInEther[6]) || !isxdigit(szInEther[7]) ||
        !isxdigit(szInEther[9]) || !isxdigit(szInEther[10]) ||
        !isxdigit(szInEther[12]) || !isxdigit(szInEther[13]) ||
        !isxdigit(szInEther[15]) || !isxdigit(szInEther[16]))
        return (-1);
    // convert the 12 hex decimals back to 6 digit decimals
    for (i = 0; i < 6; i++)
    {
        val = 0;
        c = toupper(szInEther[i*3]);
        c = c - (isdigit(c) ? '0' : ('A' - 10)); //offset adjustment
        val += c;
        val = (val << 4); // val * 16
        c = toupper(szInEther[i*3 + 1]);
        c = c - (isdigit(c) ? '0' : ('A' - 10)); // offset adjustement
        val += c;
        szOutEther[i] = val;
    }
    return (0);
}

//------------------------------------------------------------------------------------
// Inputs: PhysAddr is the hardware address in bytes
//         PhysAddrLen is the length of the PhysAddr
// Outputs: if it returns TRUE, str is the hex formated string of the hardware address.
// NOTE: make sure str is TRIPLE as big as PhysAddrLen
//-------------------------------------------------------------------------------------
BOOL PhysAddrToString(BYTE PhysAddr[], DWORD PhysAddrLen, char str[])
{
	DWORD dwIdx;

    if (PhysAddr == NULL || PhysAddrLen == 0 || str == NULL)
        return FALSE;

    str[0] = '\0';

    for (dwIdx = 0; dwIdx < PhysAddrLen; dwIdx++)
    {
        if (dwIdx == PhysAddrLen-1)
            sprintf_s(str+(dwIdx*3), sizeof(str+(dwIdx*3)), "%02X", ((int)PhysAddr[dwIdx])&0xff);
        else
            sprintf_s(str+(dwIdx*3), sizeof(str+(dwIdx*3)), "%02X-", ((int)PhysAddr[dwIdx])&0xff);

    }
    return TRUE;
}
//----------------------------------------------------------------------------
//   arp table format to be printed:
// Interface: 157.61.239.34 on Interface 2
//   Internet Address      Physical Address      Type
//   159.61.230.39         00-aa-00-61-5d-a4     dynamic
//
// Interface: 157.54.178.219 on Interface 3
//   Internet Address      Physical Address      Type
//   159.54.170.1          00-10-54-42-c0-88     dynamic
//   159.54.170.113        00-aa-00-c0-80-2e     dynamic
//----------------------------------------------------------------------------
void PrintIpNetTable(PMIB_IPNETTABLE pIpNetTable)
{
    DWORD i, dwStatus, dwCurrIndex;
    struct in_addr inadTmp;
    char szPrintablePhysAddr[256];
    char szType[128];
    char szIpAddr[128];
    PMIB_IPADDRTABLE pIpAddrTable = NULL;

    if (pIpNetTable == NULL)
    {
        printf( "pIpNetTable == NULL in line %d\n", __LINE__);
        return;
    }

    // get IP Address Table for mapping interface index number to ip address
    if ( (dwStatus = MyGetIpAddrTable(pIpAddrTable, FALSE)) != NO_ERROR)
    {
        printf("GetIpAddrTable() returned 0x%x\n", dwStatus);
        if (pIpAddrTable)
            free(pIpAddrTable);
        return;
    }
    assert(pIpAddrTable);

    // Note: the ARP table should be sorted in interface index
    dwCurrIndex = pIpNetTable->table[0].dwIndex;
    if (InterfaceIdxToInterfaceIp(pIpAddrTable, dwCurrIndex, szIpAddr))
    {
        printf("\nInterface: %s on Interface 0x%X\n", szIpAddr, dwCurrIndex);
        printf("  Internet Address      Physical Address      Type\n");
    }
    else
    {
        printf("Error: Could not convert Interface number 0x%X to IP address.\n", pIpNetTable->table[0].dwIndex);
        return;
    }

    for (i = 0; i < pIpNetTable->dwNumEntries; ++i)
    {
        if (pIpNetTable->table[i].dwIndex != dwCurrIndex)
        {
            dwCurrIndex = pIpNetTable->table[i].dwIndex;
            if (InterfaceIdxToInterfaceIp(pIpAddrTable, dwCurrIndex, szIpAddr))
            {
                printf("Interface: %s on Interface 0x%X\n", szIpAddr, dwCurrIndex);
                printf("  Internet Address      Physical Address      Type\n");
            }
            else
            {
                printf("Error: Could not convert Interface number 0x%X to IP address.\n",
                    pIpNetTable->table[0].dwIndex);
                return;
            }
        }
        PhysAddrToString(pIpNetTable->table[i].bPhysAddr, pIpNetTable->table[i].dwPhysAddrLen, szPrintablePhysAddr);
        inadTmp.s_addr = pIpNetTable->table[i].dwAddr;
        switch (pIpNetTable->table[i].dwType)
        {
        case 1:
            strcpy_s(szType, sizeof(szType), "Oother");
            break;
        case 2:
            strcpy_s(szType,sizeof(szType), "Invalidated");
            break;
        case 3:
            strcpy_s(szType,sizeof(szType), "Dynamic");
            break;
        case 4:
            strcpy_s(szType,sizeof(szType), "Static");
            break;
        default:
            strcpy_s(szType, sizeof(szType), "InvalidType");
        }
        printf("  %-16s      %-17s     %-11s\n", inet_ntoa(inadTmp), szPrintablePhysAddr, szType);

    }
    if (pIpAddrTable)
        free(pIpAddrTable);
}

//--------------------------------------------------------------------------------------
//  Format of IP Address Table:
//  ipAdEntAddr   ifAdEntIfIndex  ipAdEntNetMask   ipAdEntBcastAddr  ipAdEntReasmMaxSize
//--------------------------------------------------------------------------------------
void PrintIpAddrTable(PMIB_IPADDRTABLE pIpAddrTable)
{
    DWORD i;
    struct in_addr inadTmp1;
    struct in_addr inadTmp2;
    char szAddr[128];
    char szMask[128];

    if (pIpAddrTable == NULL)
    {
        printf("pIpAddrTable == NULL in line %d\n", __LINE__);
        return;
    }
    printf("ipAdEntAddr\t ifAdEntIfIndex\t ipAdEntNetMask\t ipAdEntBcastAddr\t ipAdEntReasmMaxSize\n");
    for (i = 0; i < pIpAddrTable->dwNumEntries; ++i)
    {
        inadTmp1.s_addr = pIpAddrTable->table[i].dwAddr;
        strcpy_s(szAddr, sizeof(szAddr), inet_ntoa(inadTmp1));
        inadTmp2.s_addr = pIpAddrTable->table[i].dwMask;
        strcpy_s(szMask, sizeof(szMask), inet_ntoa(inadTmp2));
        printf("  %s\t 0x%X\t %s\t %s\t %u\n",
                szAddr,
                pIpAddrTable->table[i].dwIndex,
                szMask,
                (pIpAddrTable->table[i].dwBCastAddr ? "255.255.255.255" : "0.0.0.0"),
                pIpAddrTable->table[i].dwReasmSize);

    }
}

//----------------------------------------------------------------------------
// Input : fOrder -- sorts the output IP Addr Table
// Output: If it returns NO_ERROR, pIpAddrTable points to the IP Addr Table
//----------------------------------------------------------------------------
// Redundant - http://msdn.microsoft.com/en-us/library/aa365949.aspx
DWORD MyGetIpAddrTable(PMIB_IPADDRTABLE pIpAddrTable, BOOL fOrder)
{
    DWORD status = NO_ERROR;
    DWORD statusRetry = NO_ERROR;
    DWORD dwActualSize = 0;	
	fOrder = FALSE;

    // query for buffer size needed
    status = GetIpAddrTable(pIpAddrTable, &dwActualSize, fOrder);

    if (status == NO_ERROR)
    {
        printf("No error!\n");
        return status;
    }
    else if (status == ERROR_INSUFFICIENT_BUFFER)
    {
        // need more space
        pIpAddrTable = (PMIB_IPADDRTABLE) malloc(dwActualSize);
        assert(pIpAddrTable);

        statusRetry = GetIpAddrTable(pIpAddrTable, &dwActualSize, fOrder);
        return statusRetry;
    }
    else
    {
        return status;
    }
}

//----------------------------------------------------------------------------
// Input : fOrder -- sorts the output IP Net Table
// Output: If it returns NO_ERROR, pIpNetTable points to the IP Net Table
//----------------------------------------------------------------------------
// Redundant: http://msdn.microsoft.com/en-us/library/aa365956.aspx
DWORD MyGetIpNetTable(PMIB_IPNETTABLE pIpNetTable, BOOL fOrder)
{
    DWORD status = NO_ERROR;
    DWORD statusRetry = NO_ERROR;
    DWORD dwActualSize = 0;
	fOrder = FALSE;

    // query for buffer size needed
    dwActualSize = 0;
    status = GetIpNetTable(pIpNetTable, &dwActualSize, fOrder);

    if (status == NO_ERROR)
    {
        return status;
    }
    else if (status == ERROR_INSUFFICIENT_BUFFER)
    {
        // need more space

        pIpNetTable = (PMIB_IPNETTABLE) malloc(dwActualSize);
        assert(pIpNetTable);

        statusRetry = GetIpNetTable(pIpNetTable, &dwActualSize, fOrder);

        if (statusRetry != NO_ERROR)
        {
#ifdef _DEBUG
            printf("Retry failed...\n");
#endif
            return statusRetry;
        }
        else
        {
            return statusRetry;
        }
    }
    else
    {
#ifdef _DEBUG
        printf("First getipnettable() call failed!\n");
#endif
        return status;
    }
}
