// list_socket.cpp: определяет точку входа для консольного приложения.
//

#include "stdafx.h"

#include <windows.h>
#include <winternl.h>
#include <Ntstatus.h>
#include <Winsock.h>


#pragma comment(lib, "ws2_32.lib")

#define SystemHandleInformation 16
#define BUF_SIZE 102400

typedef enum _POOL_TYPE
{
         NonPagedPool = 0,
         PagedPool = 1,
         NonPagedPoolMustSucceed = 2,
         DontUseThisType = 3,
         NonPagedPoolCacheAligned = 4,
         PagedPoolCacheAligned = 5,
         NonPagedPoolCacheAlignedMustS = 6,
         MaxPoolType = 7,
         NonPagedPoolSession = 32,
         PagedPoolSession = 33,
         NonPagedPoolMustSucceedSession = 34,
         DontUseThisTypeSession = 35,
         NonPagedPoolCacheAlignedSession = 36,
         PagedPoolCacheAlignedSession = 37,
         NonPagedPoolCacheAlignedMustSSession = 38
} POOL_TYPE;

typedef struct _OBJECT_NAME_INFORMATION {


  UNICODE_STRING          Name;
  WCHAR                   NameBuffer[0];

} OBJECT_NAME_INFORMATION, *POBJECT_NAME_INFORMATION;

typedef struct _OBJECT_TYPE_INFORMATION {



  UNICODE_STRING          TypeName;
  ULONG                   TotalNumberOfHandles;
  ULONG                   TotalNumberOfObjects;
  WCHAR                   Unused1[8];
  ULONG                   HighWaterNumberOfHandles;
  ULONG                   HighWaterNumberOfObjects;
  WCHAR                   Unused2[8];
  ACCESS_MASK             InvalidAttributes;
  GENERIC_MAPPING         GenericMapping;
  ACCESS_MASK             ValidAttributes;
  BOOLEAN                 SecurityRequired;
  BOOLEAN                 MaintainHandleCount;
  USHORT                  MaintainTypeList;
  POOL_TYPE               PoolType;
  ULONG                   DefaultPagedPoolCharge;
  ULONG                   DefaultNonPagedPoolCharge;



} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

typedef enum _OBJECT_INFORMATION_CLASS {
    ObjectBasicInformation,
    ObjectNameInformation,
    ObjectTypeInformation,
    ObjectAllInformation,
    ObjectDataInformation
} OBJECT_INFORMATION_CLASS, *POBJECT_INFORMATION_CLASS;


typedef NTSTATUS (CALLBACK *LPFNZwQuerySystemInformation)(
  SYSTEM_INFORMATION_CLASS SystemInformationClass,
  PVOID SystemInformation,
  ULONG SystemInformationLength,
  PULONG ReturnLength
);

typedef NTSTATUS (CALLBACK *LPFNNtQueryObject)(
  HANDLE               ObjectHandle,
  OBJECT_INFORMATION_CLASS ObjectInformationClass,
  PVOID               ObjectInformation,
  ULONG                Length,
  PULONG              ResultLength );

typedef NTSTATUS (CALLBACK *LPFNZwClose)(HANDLE  Handle);

LPFNZwQuerySystemInformation ZwQuerySystemInformation;
LPFNNtQueryObject NtQueryObject;
LPFNZwClose ZwClose;




typedef struct _SYSTEM_HANDLE_INFORMATION {
  USHORT ProcessId;
  USHORT CreatorBackTraceIndex;
  UCHAR ObjectTypeNumber;
  UCHAR Flags;
  USHORT Handle;
  PVOID Object;
  ACCESS_MASK GrantedAccess;
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

typedef struct _SYSTEM_HANDLE_INFORMATION_EX {

  ULONG NumberOfHandles;
  SYSTEM_HANDLE_INFORMATION Information[1];
  
} SYSTEM_HANDLE_INFORMATION_EX, *PSYSTEM_HANDLE_INFORMATION_EX;




int _tmain(int argc, _TCHAR* argv[])
{

DWORD  PID,ret;
NTSTATUS  ss;
int i, j, ress, rem_port, loc_port;
PUNICODE_STRING us;
POBJECT_TYPE_INFORMATION ot;
char *remaddr,*locaddr;

sockaddr_in sockname, locname;
WSAData   WSData;
LPVOID   pinfo, codeadr;
DWORD    tid, wsa_error;

HINSTANCE hNTdll = LoadLibrary(L"Ntdll.dll"); 
HANDLE    hProcess, hcp, ObjHandle, hh;
if(!hNTdll) return 1;

ZwQuerySystemInformation = (LPFNZwQuerySystemInformation) 
GetProcAddress(hNTdll, "ZwQuerySystemInformation");

NtQueryObject=(LPFNNtQueryObject) 
GetProcAddress(hNTdll, "NtQueryObject");

ZwClose=(LPFNZwClose) 
GetProcAddress(hNTdll, "ZwClose");

 WSAStartup(0x202, &WSData);
  printf("There are the following sockets opened on system:\n");

 // Get the number of handles on the system
  DWORD buffer_size = 0;
  SYSTEM_HANDLE_INFORMATION_EX temp_info;


  NTSTATUS status = ZwQuerySystemInformation(
      (SYSTEM_INFORMATION_CLASS)SystemHandleInformation, &temp_info, 
      sizeof(temp_info), &buffer_size);


  SYSTEM_HANDLE_INFORMATION_EX *system_handles =
    (SYSTEM_HANDLE_INFORMATION_EX*)(new BYTE[buffer_size]);


  status = ZwQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemHandleInformation,
                                    system_handles,
                                    buffer_size, &buffer_size);

  ot=(POBJECT_TYPE_INFORMATION)malloc(BUF_SIZE);
  us=(PUNICODE_STRING)malloc(BUF_SIZE);


  //printf("nb of handles = %d", system_handles->NumberOfHandles);
for(i=0;i<system_handles->NumberOfHandles;i++)
{
	 hProcess = OpenProcess(PROCESS_DUP_HANDLE, FALSE, system_handles->Information[i].ProcessId);
	 if (hProcess != INVALID_HANDLE_VALUE){
          hh=(HANDLE)system_handles->Information[i].Handle;
		 if(DuplicateHandle(hProcess,hh,
                       INVALID_HANDLE_VALUE, &ObjHandle, DUPLICATE_SAME_ACCESS, 0, 0))
		 {
			ss = NtQueryObject(ObjHandle,
                          ObjectTypeInformation, ot, BUF_SIZE, &ret);

			if (ss == STATUS_SUCCESS)
			{
				if(lstrcmp(ot->TypeName.Buffer,L"File")==0){
					ss=NtQueryObject(ObjHandle,ObjectNameInformation, us, BUF_SIZE, &ret);
					
					

					if(ss==STATUS_SUCCESS){
						if(lstrcmp(us->Buffer,L"\\Device\\Afd")==0)
						{
							
							

							 ress = sizeof(sockaddr_in);//--- commenting this line would cause an 10014 error.
							getpeername((SOCKET)ObjHandle, (sockaddr *)&sockname, &ress);

							//--- determine local ip and port
							ress = sizeof(sockaddr_in);//--- commenting this line would cause an 10014 error.
							getsockname((SOCKET)ObjHandle, (sockaddr *)&locname, &ress);
				


							//--- type conversion
							remaddr = inet_ntoa(sockname.sin_addr);
							rem_port = ntohs(sockname.sin_port);

							locaddr = inet_ntoa(locname.sin_addr);
							loc_port = ntohs(locname.sin_port);

							
							wprintf(L"PID=%d; Local=%S:%d; Peer=%S:%d\n",system_handles->Information[i].ProcessId,locaddr,loc_port,remaddr,rem_port);
					

						}
					}
				}

			
			//ZwClose(ObjHandle); //--- we don't need this handle any more
			}


		 }


		 CloseHandle(hProcess);
	 }

}


free(us);
free(ot);
   WSACleanup();

FreeLibrary(hNTdll);

getchar();

return 0;
}