C:/work/RevealedSystems/EyeSpy-RS/eyespy/src/swig/extendedTCPTableLib/ExtendedXportTable.cxx

Go to the documentation of this file.

/*
    Copyright (C) <2008>  <Mark Hanegraaff>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "Winsock2.h"
#include "Iphlpapi.h"
#include <stdio.h>
#include <tchar.h>
#include <psapi.h>


#define MALLOC(x) HeapAlloc(GetProcessHeap(), 0, (x))
#define FREE(x) HeapFree(GetProcessHeap(), 0, (x))

//FWD Declarations
void clearERROR_MESSAGE();
bool tableArrayOutOfRange(unsigned int idx);
bool indexOutOfRange(unsigned int idx);
void populateTCPV4Table();
void populateUDPV4Table();
BOOL GetProcessName( DWORD processID, TCHAR* szProcessName);
void updateV4Tables();



char ERROR_MESSAGE[4096];
const unsigned int maxSize = 10000;
unsigned int tableSize;

typedef struct V4Table
{
        unsigned int    Pid;
        char                    processName[256];
        char                    transport[9]; //TCP or UDP
        unsigned int    srcPort;
        unsigned int    state;  


} V4Table;


V4Table tableArray[maxSize];



//
// Getters
//

unsigned int getV4TableSize() {return tableSize;}


unsigned int getV4Pid(unsigned int idx) throw (const char *)
{
        clearERROR_MESSAGE();

        if (tableArrayOutOfRange(idx) || indexOutOfRange(idx)) throw ERROR_MESSAGE;

        clearERROR_MESSAGE();

        return tableArray[idx].Pid;
                
}

char * getV4ProcessName(unsigned int idx) throw (const char *)
{
        clearERROR_MESSAGE();

        if (tableArrayOutOfRange(idx) || indexOutOfRange(idx)) throw ERROR_MESSAGE;

        clearERROR_MESSAGE();

        return tableArray[idx].processName;
                
}


unsigned int getV4SrcPort(unsigned int idx) throw (const char *)
{
        clearERROR_MESSAGE();

        if (tableArrayOutOfRange(idx) || indexOutOfRange(idx)) throw ERROR_MESSAGE;

        clearERROR_MESSAGE();

        return tableArray[idx].srcPort; 
}


unsigned int getV4State(unsigned int idx) throw (const char *)
{
        clearERROR_MESSAGE();

        if (tableArrayOutOfRange(idx) || indexOutOfRange(idx)) throw ERROR_MESSAGE;

        clearERROR_MESSAGE();

        return tableArray[idx].state;   
}

char * getV4Transport(unsigned int idx) throw (const char *)
{
        clearERROR_MESSAGE();

        if (tableArrayOutOfRange(idx) || indexOutOfRange(idx)) throw ERROR_MESSAGE;

        clearERROR_MESSAGE();

        return tableArray[idx].transport;       
}


//
// Exported work function. This does all the actual work
//

void updateV4Tables()
{
        tableSize = 0;

        // initialize the table
        memset(tableArray, 0, sizeof(tableArray) );

        populateTCPV4Table();
        populateUDPV4Table();   
}


void populateTCPV4Table() throw (const char *)
{       
        clearERROR_MESSAGE();   

        MIB_TCPTABLE_OWNER_PID * pTcpTable;

        DWORD dwSize = 0;
        DWORD dwRetVal = 0;     

        pTcpTable = ( MIB_TCPTABLE_OWNER_PID *) MALLOC(sizeof ( MIB_TCPTABLE_OWNER_PID));

        if (pTcpTable == NULL) 
        {
                strcpy(ERROR_MESSAGE, "Could not allocate memory for TCP Table");
                throw ERROR_MESSAGE;
        }

        dwSize = sizeof (MIB_TCPTABLE_OWNER_PID);

        if ((dwRetVal = GetExtendedTcpTable (pTcpTable, &dwSize, TRUE, AF_INET, TCP_TABLE_OWNER_PID_ALL , 0)) == ERROR_INSUFFICIENT_BUFFER) 
        {
                FREE(pTcpTable);
                pTcpTable = (MIB_TCPTABLE_OWNER_PID *) MALLOC(dwSize);
    
                if (pTcpTable == NULL) 
                {            
                        strcpy(ERROR_MESSAGE, "Could not allocate memory for TCP Table");
                        throw ERROR_MESSAGE;

                        return;
                }
        }

        // Make a second call to GetTcpTable to get
        // the actual data we require

        if ((dwRetVal = GetExtendedTcpTable (pTcpTable, &dwSize, TRUE, AF_INET, TCP_TABLE_OWNER_PID_ALL , 0)) == NO_ERROR)
        {
                //
        }
        else
        {
                strcpy(ERROR_MESSAGE, "Could not allocate memory for TCP Table");
                FREE(pTcpTable);

                throw ERROR_MESSAGE;
        }

        char processName[1024];


        for (int i = 0; i < pTcpTable->dwNumEntries; i++)
        {
                GetProcessName(pTcpTable->table[i].dwOwningPid, processName);                           

                //getTCPProcName(&pTcpTable->table[i], processName, processPath);


                tableArray[tableSize].Pid = pTcpTable->table[i].dwOwningPid;
                strcpy(tableArray[i].processName, processName);
                tableArray[tableSize].srcPort = ntohs(pTcpTable->table[i].dwLocalPort);
                strcpy(tableArray[tableSize].transport, "TCP");
                tableArray[tableSize].state = pTcpTable->table[i].dwState;

                tableSize ++;
        }

        FREE(pTcpTable);
}





void populateUDPV4Table() throw (const char *)
{       
        clearERROR_MESSAGE();

        MIB_UDPTABLE_OWNER_PID* pUDPTable;

        DWORD dwSize = 0;
        DWORD dwRetVal = 0;     

        pUDPTable = ( MIB_UDPTABLE_OWNER_PID  *) MALLOC(sizeof ( MIB_UDPTABLE_OWNER_PID ));

        if (pUDPTable == NULL) 
        {               
                strcpy(ERROR_MESSAGE, "Error allocating UDP table memory");
        throw ERROR_MESSAGE;
        }

        dwSize = sizeof (MIB_UDPTABLE_OWNER_PID);

        if ((dwRetVal = GetExtendedUdpTable (pUDPTable, &dwSize, TRUE, AF_INET,  UDP_TABLE_OWNER_PID, 0)) == ERROR_INSUFFICIENT_BUFFER) 
        {
                FREE(pUDPTable);
                pUDPTable = (MIB_UDPTABLE_OWNER_PID *) MALLOC(dwSize);
    
                if (pUDPTable == NULL) 
                {            
                        strcpy(ERROR_MESSAGE, "Error allocating UDP table memory");
                        throw ERROR_MESSAGE;
                }
        }


        // Make a second call to GetTcpTable to get
        // the actual data we require

        if ((dwRetVal = GetExtendedUdpTable (pUDPTable, &dwSize, TRUE, AF_INET, UDP_TABLE_OWNER_PID, 0)) == NO_ERROR)
        {
                //
        }
        else
        {
                strcpy(ERROR_MESSAGE, "Error allocating UDP table memory");
                FREE(pUDPTable);

        throw ERROR_MESSAGE;
        }

        char procName[1024];

        for (int i = 0; i < pUDPTable->dwNumEntries; i++)
        {
                GetProcessName(pUDPTable->table[i].dwOwningPid, procName);              
                strcpy(tableArray[tableSize].processName, procName);

                tableArray[tableSize].Pid = pUDPTable->table[i].dwOwningPid;            
                tableArray[tableSize].srcPort =  ntohs(pUDPTable->table[i].dwLocalPort);
                tableArray[tableSize].state = -1;
                strcpy(tableArray[tableSize].transport,  "UDP");

                tableSize++;
        }

        FREE(pUDPTable);

}


//
// Helpers
//

void clearERROR_MESSAGE()
{
        //Nulls string may cause problems back in Pythonland
        strcpy(ERROR_MESSAGE, "NOERROR");
}

bool tableArrayOutOfRange(unsigned int idx)
{
        if (idx >= maxSize)
        {
                sprintf(ERROR_MESSAGE, "Index exceeded size of table: Index:%d, Max Table Size:%d", idx, maxSize);
                return true;
        }

        return false;
}

bool indexOutOfRange(unsigned int idx)
{
        if (idx >= tableSize)
        {
                sprintf(ERROR_MESSAGE, "Index out of Range: Index:%d, Current Table Size:%d", idx, tableSize);
                return true;
        }

        return false;
}


//      Returns the process name give its pid

BOOL GetProcessName( DWORD processID, TCHAR* szProcessName)
{
 
    BOOL bRC = FALSE;
 
    // Get a handle to the process.
 
    HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                                   PROCESS_VM_READ,
                                   FALSE, processID );
 
    // Get the process name.
 
    if (NULL != hProcess )
    {
        HMODULE hMod;
        DWORD cbNeeded;
 
        if ( EnumProcessModules( hProcess, &hMod, sizeof(hMod),
             &cbNeeded) )
        {
            GetModuleBaseName( hProcess, hMod, szProcessName,
                               MAX_PATH );
 
            bRC = TRUE;
        }
                else
                {
                        strcpy(szProcessName, "UNKNOWN");
                }
    }
        else
        {
                strcpy(szProcessName, "UNKNOWN");
        }
 
    CloseHandle( hProcess );
 
    return bRC;
}

---