/*++
    tcpdel.c

    Author:
        Matt Edman

    Date:
        2/18/2005

    Description:
        This little tool intends to serve the same purpose as the tcpdrop utility on OpenBSD; 
        it just kills an open TCP connection. There weren't many utilities (that I could find)
        to do this on Windows except for one shareware utility that claims this is some magical 
        trick. As you can see, it's really not difficult. I also added support for wildcards
        when specifying which TCP connection to terminate.
--*/

#include <stdio.h>
#include <Winsock2.h>
#include <Ws2tcpip.h>
#include <iphlpapi.h>

/* Needed libraries (VC7) */
#pragma comment(lib, "iphlpapi.lib")
#pragma comment(lib, "Ws2_32.lib")

#define MATCH_ANY   -1

/* Display usage information to a confused user */
void print_usage(const char *exename) {
    printf("\nusage: %s localAddr localPort foreignAddr foreignPort\n\n", exename);
    printf("You can specify * for any or all of the fields. If a hostname or a\n");
    printf("service name is supplied, %s will attempt to resolve them to their\n", exename);
    printf("appropriate values.\n\n");
    printf("Ex:\t%s 10.0.0.10 * 10.0.0.11 1234\n", exename);
    printf("\t%s * * blah.com http\n", exename);
    printf("\t%s * * * *\n", exename);
}

/* Resolve a given hostname and service name */
int gethostinfo(const char *host, 
                const char *service, 
                /*OUT*/ DWORD *ip, 
                /*OUT*/ DWORD *port) {
    struct addrinfo *aiOut;
    struct addrinfo aih;
    struct sockaddr_in *sock;
    *ip = 0;
    *port = 0;

    /* Don't pass the wildcards to getaddrinfo, it doesn't like them */
    if (!strcmp(host, "*")) {
        *ip = MATCH_ANY;
        host = "\0";
    }
    if (!strcmp(service, "*")) {
        *port = MATCH_ANY;
        service = "\0";
    }

    /* Get the address info, if there is anything to get */
    if (*ip != MATCH_ANY || *port != MATCH_ANY) {
        /* Zero the addrinfo hints structure */
        memset(&aih, 0, sizeof(aih));
        aih.ai_family   = AF_INET;
        aih.ai_socktype = SOCK_STREAM;
        aih.ai_protocol = IPPROTO_TCP;
    
        if (getaddrinfo(host, service, &aih, &aiOut) != 0)
            return WSAGetLastError();

        sock  = (struct sockaddr_in *)aiOut->ai_addr;

        /* The address and port values are stored as DWORDs in the MIB_TCPROW *
         * structure, so it is necessary to cast to these types here.         */
        if (*ip != MATCH_ANY) *ip = (DWORD)sock->sin_addr.S_un.S_addr;
        if (*port != MATCH_ANY) *port = (DWORD)sock->sin_port;

        freeaddrinfo(aiOut);
    }
   return 0;
}

/* Delete all matching established TCP connections */
int tcpdelete(MIB_TCPTABLE *pTcpTable, 
              DWORD localIP, 
              DWORD localPort, 
              DWORD remoteIP, 
              DWORD remotePort) {

    MIB_TCPROW *row;
    DWORD i;
    int nDeleted = 0;
   
    /* Iterate through the table rows and delete one if a match is found */
    for (i = 0; i < pTcpTable->dwNumEntries; i++) {
        row = &pTcpTable->table[i];

        if ((row->dwLocalAddr  == localIP    || localIP    == MATCH_ANY) &&
            (row->dwLocalPort  == localPort  || localPort  == MATCH_ANY) &&
            (row->dwRemoteAddr == remoteIP   || remoteIP   == MATCH_ANY) &&
            (row->dwRemotePort == remotePort || remotePort == MATCH_ANY)) {
                row->dwState = MIB_TCP_STATE_DELETE_TCB;
                if (SetTcpEntry(row) == NO_ERROR)
                    nDeleted++;
            }
    }
    return nDeleted;
}

int main(int argc, char *argv[]) {
    WSADATA WSAData;
    MIB_TCPTABLE *pTcpTable = (MIB_TCPTABLE *)malloc(sizeof(MIB_TCPTABLE));
    DWORD dwBuffSize = sizeof(MIB_TCPTABLE);
    DWORD dwRet;
    DWORD dwLocalIP, dwLocalPort, dwRemoteIP, dwRemotePort;
    int   numDeleted;
    
    if (argc != 5) {      
        /* Display usage information */
        print_usage(argv[0]);
        return -1;
    }

    /* This silly exercise is necessary before Winsock will let us resolve any  *
     * host names.                                                              */    
    if (WSAStartup(MAKEWORD(2,0), &WSAData) != 0) {
        printf("Winsock failed to initialize.\n");
        return -1;
  }
    
    /* Retrieve the current TCP table */
    dwRet = GetTcpTable(pTcpTable, &dwBuffSize, TRUE);
    if (dwRet  == ERROR_INSUFFICIENT_BUFFER) {
        /* Need to make the table buffer bigger */
        free(pTcpTable);
        pTcpTable = (MIB_TCPTABLE *)malloc(dwBuffSize);
        dwRet = GetTcpTable(pTcpTable, &dwBuffSize, TRUE);
    }
    
    if (dwRet != NO_ERROR) {
        /* Bad stuff happened */
        free(pTcpTable);
        printf("Error retrieving the TCP table.\n");
        return -1;
    }
        
    /* Resolve the specified local and remote address information */
    if (gethostinfo(argv[1], argv[2], &dwLocalIP, &dwLocalPort) != 0) {
        free(pTcpTable);
        printf("Error resolving local address information.\n");
        return -1;
    }
    if (gethostinfo(argv[3], argv[4], &dwRemoteIP, &dwRemotePort) != 0) {
        free(pTcpTable);
        printf("Error resolving remote address information.\n");
        return -1;
    }

    /* Attempt to delete the TCP connections and display the results */
    numDeleted = tcpdelete(pTcpTable, dwLocalIP, dwLocalPort, dwRemoteIP, dwRemotePort);
    printf("%d TCP connection%sdeleted.\n", numDeleted, (numDeleted == 1 ? " " : "s "));

    free(pTcpTable);
    return numDeleted;
}